fix(engine): bound Remember/Add prompt budget + per-stage timings (#41) by tyejcoleman · Pull Request #44 · Keyoku-ai/keyoku-engine

tyejcoleman · 2026-04-24T04:39:14Z

Summary

Closes #41.

Cap existingMemories and conversationCtx by byte budget (4KB each default) before the extraction LLM call; trim per-memory content to 240 chars with ellipsis.
Replace hardcoded k=10, minScore=0.5 in engine_add.go:109 with EngineConfig.ExtractionExistingMaxCount / ExtractionExistingMinScore.
Emit one structured log line per Add() call with session_ms, embed_ms, similar_ms, extract_ms, existing_count/bytes, ctx_msgs/bytes, err — required by the new invariant in harness/40-invariants/observability.md.

Why

moltar-bot's profile (issue #41) showed /remember degrading to 25–55s p95 on repeated same-entity writes under Ollama + qwen2.5:7b-instruct, with a 60s client timeout hit. /search (~40ms) and raw Ollama generate (~6–8s) stayed fast, localizing the bottleneck to extraction prompt assembly. The root is in engine_add.go:109-128: every 0.5+ similarity neighbor was appended verbosely to the prompt, with no byte cap. Harness 45-tunables.md already flagged k=10, minScore=0.5 as hardcoded — fixing it here closes that gap.

Harness changes (companion)

45-tunables.md — new "Extraction prompt budget (Add / Remember)" table.
40-invariants/observability.md — new "Required: Add()/Remember stage timings" subsection defining the mandatory log line.

Validation

6 new unit tests in engine/engine_add_prompt_budget_test.go — unlimited, byte-budget trim (context + existing), per-memory content trim, nil safety, most-similar-first preservation. go test ./engine/ -run TestBuildConversation\|TestBuildExistingMemoriesBlock — all pass.
go test ./engine/... -count=1 — full suite green.
go vet ./... && go build ./... — clean.

Deferred (separate tickets)

Async / lite-extraction degradation path (moltar-bot's item 3). Worth measuring the impact of the budget caps first; if repeat-entity p95 still exceeds an acceptable bound, we'll add either a state=pending_extraction async path or a reduced-prompt "lite" fallback.
Regression benchmark for repeated same-entity writes (moltar-bot's item 4). Blocked by wanting moltar-bot's harness script attached to remember latency bottleneck: ExtractMemories path degrades to 25–55s and times out on repeated same-entity writes #41 for fidelity; easy to add once we have the exact scenario.

Test plan

Unit: byte budget trims existingMemories under pressure.
Unit: per-memory content trim + ellipsis.
Unit: most-similar-first preserved on tight budget.
Unit: conversation ctx trimmed.
Manual: run moltar-bot's reproduction (12× same-entity /remember on Ollama qwen2.5:7b) and compare p95 before/after. Expect extract_ms drop proportional to the existing-memories block shrink.
Grep stage=extract log lines in a live run to confirm the structured timing invariant is intact.

🤖 Generated with Claude Code

github-actions · 2026-04-24T04:39:22Z

All contributors have signed the CLA ✍️ ✅
_{Posted by the CLA Assistant Lite bot.}

Copilot

Pull request overview

This PR addresses /remember latency regressions in the extraction prompt assembly by bounding the prompt inputs (existing similar memories + conversation context) and adding per-stage timing observability in Engine.Add().

Changes:

Add prompt-budget helper functions to cap conversationCtx and existingMemories by byte budget, plus per-memory content trimming with ellipsis.
Replace hardcoded similarity search parameters (k=10, minScore=0.5) with new EngineConfig tunables for extraction-neighbor selection.
Emit a structured, per-stage timing log line around the extraction path in Engine.Add().

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.

File	Description
engine/engine_add_prompt_budget.go	Adds pure helper functions to format and byte-cap extraction prompt components.
engine/engine_add_prompt_budget_test.go	Adds unit tests covering budget trimming, nil-safety, ordering, and per-memory trimming behavior.
engine/engine_add.go	Integrates prompt budgeting, config-driven similarity search knobs, and stage timing logging into `Add()`.
engine/engine.go	Extends `EngineConfig` with extraction prompt budget/tunable fields and applies defaults in `NewEngine()`.

+	ExtractionExistingMaxBytes   int     // byte cap on formatted existingMemories block (default: 4096; 0 = unlimited)
+	ExtractionContextMaxBytes    int     // byte cap on formatted conversationCtx block (default: 4096; 0 = unlimited)
+	ExtractionMemoryContentTrim  int     // per-memory content trim for the existing-memories prompt (default: 240; 0 = no trim)


+		if perMemContentTrim > 0 && len(content) > perMemContentTrim {
+			content = content[:perMemContentTrim] + "..."


@@ -106,26 +108,46 @@ func (e *Engine) Add(ctx context.Context, entityID string, req AddRequest) (*Add
 		visibilityFor = &storage.VisibilityContext{AgentID: agentID, TeamID: req.TeamID}
 	}

-	similarMemories, err := e.store.FindSimilarWithOptions(ctx, queryEmbedding, entityID, 10, 0.5, storage.SimilarityOptions{
-		AgentID:       req.AgentID,
-		VisibilityFor: visibilityFor,
-	})
+	t0 = time.Now()
+	similarMemories, err := e.store.FindSimilarWithOptions(
+		ctx, queryEmbedding, entityID,
+		e.config.ExtractionExistingMaxCount,
+		e.config.ExtractionExistingMinScore,
+		storage.SimilarityOptions{
+			AgentID:       req.AgentID,
+			VisibilityFor: visibilityFor,
+		},
+	)
 	if err != nil {
 		return nil, fmt.Errorf("failed to find similar memories: %w", err)
 	}
+	similarMs := time.Since(t0).Milliseconds()

-	var existingMemories []string
-	for _, sm := range similarMemories {
-		existingMemories = append(existingMemories, fmt.Sprintf("[%s, %s, importance:%.1f] %s",
-			sm.Memory.Type, sm.Memory.State, sm.Memory.Importance, sm.Memory.Content))
-	}
+	existingMemories, existingBytes := buildExistingMemoriesBlock(
+		similarMemories,
+		e.config.ExtractionExistingMaxBytes,
+		e.config.ExtractionMemoryContentTrim,
+	)

 	// Extract memories using LLM
+	t0 = time.Now()
 	extractResp, err := e.provider.ExtractMemories(ctx, llm.ExtractionRequest{
 		Content:          req.Content,
 		ConversationCtx:  conversationCtx,
 		ExistingMemories: existingMemories,
 	})
+	extractMs := time.Since(t0).Milliseconds()
+
+	// Per-stage timing + payload-size trace. See harness/40-invariants/observability.md
+	// (Add() stage timings) and issue #41. Structured key=value so ops/tooling can grep.
+	log.Printf(
+		"INFO [remember/add] stage=extract entity_id=%q session_ms=%d embed_ms=%d similar_ms=%d extract_ms=%d existing_count=%d existing_bytes=%d ctx_msgs=%d ctx_bytes=%d err=%v",
+		entityID, sessionMs, embedMs, similarMs, extractMs,
+		len(existingMemories), existingBytes,
+		len(conversationCtx), ctxBytes,
+		err,
+	)
+
 	if err != nil {
 		return nil, fmt.Errorf("extraction failed: %w", err)
 	}


Closes #41. Before: Engine.Add() called FindSimilarWithOptions with hardcoded k=10 and minScore=0.5, then verbosely formatted every neighbor into the extraction prompt with no byte/token cap. Repeated same-entity writes accumulated more 0.5+ similarity neighbors, inflating the prompt until local LLMs (qwen2.5:7b-instruct) degraded to 25-55s per call and hit the client 60s timeout. /search stayed fast (~40ms) and raw Ollama generate stayed fast (~6-8s), confirming the bottleneck lives in the Add() context assembly, not the provider baseline. Changes: - EngineConfig gains ExtractionExistingMaxCount (k), ExtractionExistingMinScore, ExtractionExistingMaxBytes (4KB default), ExtractionContextMaxBytes (4KB default), ExtractionMemoryContentTrim (240 chars default). NewEngine applies defaults on zero values; negative values opt out (unlimited). - New helpers buildConversationContext and buildExistingMemoriesBlock trim by total byte budget while preserving the most-similar-first ordering that FindSimilarWithOptions returns. Per-memory content is trimmed with an ellipsis so a single verbose memory can't blow the budget on its own. - Engine.Add() now uses the configured k/minScore instead of literals, and emits one structured log line per call with session_ms, embed_ms, similar_ms, extract_ms, existing_count/bytes, ctx_msgs/bytes, err. This matches the new Add() stage-timing invariant in harness/40-invariants/observability.md and gives us a grep-able regression signal without adding Prometheus/OTel. Validation: - Unit: 6 tests in engine_add_prompt_budget_test.go cover unlimited, byte-budget trim, per-memory content trim, nil safety, and most-similar-first preservation. - Full ./engine/... suite still green.

tyejcoleman · 2026-04-24T05:07:09Z

I have read the CLA Document and I hereby sign the CLA

tyejcoleman · 2026-04-24T05:10:14Z

recheck

Copilot AI review requested due to automatic review settings April 24, 2026 04:39

Copilot started reviewing on behalf of tyejcoleman April 24, 2026 04:39 View session

Copilot AI reviewed Apr 24, 2026

View reviewed changes

tyejcoleman force-pushed the fix/remember-latency-prompt-budget branch from bf66aed to 0134c59 Compare April 24, 2026 04:58

tyejcoleman force-pushed the fix/remember-latency-prompt-budget branch from 0134c59 to f75a369 Compare April 24, 2026 05:03

chore: trigger CI refresh

b059afd

tyejcoleman merged commit 9bc5031 into main Apr 24, 2026
1 of 3 checks passed

Navigation Menu

Uh oh!

fix(engine): bound Remember/Add prompt budget + per-stage timings (#41) - #44

fix(engine): bound Remember/Add prompt budget + per-stage timings (#41)#44
tyejcoleman merged 2 commits into
mainfrom
fix/remember-latency-prompt-budget

tyejcoleman commented Apr 24, 2026

Uh oh!

github-actions Bot commented Apr 24, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

tyejcoleman commented Apr 24, 2026

Uh oh!

tyejcoleman commented Apr 24, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		if perMemContentTrim > 0 && len(content) > perMemContentTrim {
		content = content[:perMemContentTrim] + "..."

Uh oh!

Conversation

tyejcoleman commented Apr 24, 2026

Summary

Why

Harness changes (companion)

Validation

Deferred (separate tickets)

Test plan

Uh oh!

github-actions Bot commented Apr 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

tyejcoleman commented Apr 24, 2026

Uh oh!

tyejcoleman commented Apr 24, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

github-actions Bot commented Apr 24, 2026 •

edited

Loading