diff --git a/engine/engine.go b/engine/engine.go index 59d70ff..fe009e8 100644 --- a/engine/engine.go +++ b/engine/engine.go @@ -33,18 +33,32 @@ type EngineConfig struct { ReEvalSimilarityMin float64 // min similarity for re-eval candidates (default: 0.5) ReEvalSimilarityMax float64 // max similarity for re-eval candidates (default: 0.85) ReEvalImportanceMin float64 // min importance for re-eval candidates (default: 0.5) + + // Extraction prompt budgets — cap inputs into ExtractMemories so repeated + // same-entity writes don't balloon the prompt and push local LLMs over + // their usable latency window. See issue #41. + ExtractionExistingMaxCount int // max existing-memory neighbors considered (default: 10) + ExtractionExistingMinScore float64 // min similarity for existing-memory neighbors (default: 0.5) + 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) } // DefaultEngineConfig returns a default engine configuration. func DefaultEngineConfig() EngineConfig { return EngineConfig{ - ContextTurns: 5, - DefaultMinScore: 0.3, - FTSBaselineSimilarity: 0.4, - DiversityThreshold: 0.9, - ReEvalSimilarityMin: 0.5, - ReEvalSimilarityMax: 0.85, - ReEvalImportanceMin: 0.5, + ContextTurns: 5, + DefaultMinScore: 0.3, + FTSBaselineSimilarity: 0.4, + DiversityThreshold: 0.9, + ReEvalSimilarityMin: 0.5, + ReEvalSimilarityMax: 0.85, + ReEvalImportanceMin: 0.5, + ExtractionExistingMaxCount: 10, + ExtractionExistingMinScore: 0.5, + ExtractionExistingMaxBytes: 4096, + ExtractionContextMaxBytes: 4096, + ExtractionMemoryContentTrim: 240, } } @@ -100,6 +114,27 @@ func NewEngine( if config.ReEvalImportanceMin <= 0 { config.ReEvalImportanceMin = 0.5 } + if config.ExtractionExistingMaxCount <= 0 { + config.ExtractionExistingMaxCount = 10 + } + if config.ExtractionExistingMinScore <= 0 { + config.ExtractionExistingMinScore = 0.5 + } + if config.ExtractionExistingMaxBytes < 0 { + config.ExtractionExistingMaxBytes = 0 + } else if config.ExtractionExistingMaxBytes == 0 { + config.ExtractionExistingMaxBytes = 4096 + } + if config.ExtractionContextMaxBytes < 0 { + config.ExtractionContextMaxBytes = 0 + } else if config.ExtractionContextMaxBytes == 0 { + config.ExtractionContextMaxBytes = 4096 + } + if config.ExtractionMemoryContentTrim < 0 { + config.ExtractionMemoryContentTrim = 0 + } else if config.ExtractionMemoryContentTrim == 0 { + config.ExtractionMemoryContentTrim = 240 + } graphEngine := NewGraphEngine(store, DefaultGraphConfig()) diff --git a/engine/engine_add.go b/engine/engine_add.go index 997f5eb..a21c763 100644 --- a/engine/engine_add.go +++ b/engine/engine_add.go @@ -6,6 +6,7 @@ package engine import ( "context" "fmt" + "log" "time" "github.com/keyoku-ai/keyoku-engine/llm" @@ -80,21 +81,22 @@ func (e *Engine) Add(ctx context.Context, entityID string, req AddRequest) (*Add } // Get conversation context + t0 := time.Now() contextMsgs, err := e.store.GetRecentSessionMessages(ctx, entityID, e.config.ContextTurns) if err != nil { return nil, fmt.Errorf("failed to get session context: %w", err) } + sessionMs := time.Since(t0).Milliseconds() - var conversationCtx []string - for _, msg := range contextMsgs { - conversationCtx = append(conversationCtx, fmt.Sprintf("[%s]: %s", msg.Role, msg.Content)) - } + conversationCtx, ctxBytes := buildConversationContext(contextMsgs, e.config.ExtractionContextMaxBytes) // Embed content to find similar existing memories + t0 = time.Now() queryEmbedding, err := e.embedder.Embed(ctx, req.Content) if err != nil { return nil, fmt.Errorf("failed to embed query: %w", err) } + embedMs := time.Since(t0).Milliseconds() // Build visibility context for similarity search so dedup/conflict only considers visible memories var visibilityFor *storage.VisibilityContext @@ -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) } diff --git a/engine/engine_add_prompt_budget.go b/engine/engine_add_prompt_budget.go new file mode 100644 index 0000000..2dbcd01 --- /dev/null +++ b/engine/engine_add_prompt_budget.go @@ -0,0 +1,78 @@ +// SPDX-License-Identifier: BSL-1.1 +// Copyright (c) 2026 Keyoku. All rights reserved. + +package engine + +import ( + "fmt" + + "github.com/keyoku-ai/keyoku-engine/storage" +) + +// Prompt-budget helpers for the Add() / ExtractMemories path. +// +// Rationale (issue #41): repeated same-entity writes returned steadily more +// 0.5+ similarity neighbors, which inflated the formatted existingMemories +// block inside the extraction prompt and pushed local LLMs (e.g. qwen2.5:7b) +// past their usable latency window (25–55s, then 60s timeouts). The v2 harness +// flagged k=10/minScore=0.5 as hardcoded and observability.md did not require +// per-stage timings. These helpers cap the prompt inputs by byte budget +// (most-similar first, since similarMemories arrives ordered by score) and are +// pure functions so they're cheap to test without provider stubs. +// +// Byte-budget semantics: +// - budget > 0 → hard cap; stop once the formatted block would exceed it. +// - budget == 0 → treated as "unlimited" (no trim). Defaults are applied in +// NewEngine, so a zero-valued config struct still gets sensible caps. + +// buildConversationContext formats recent session messages for the extraction +// prompt and trims by total byte budget (includes newline separators in the +// running count, approximated by adding 1 per entry). +func buildConversationContext(msgs []*storage.SessionMessage, maxBytes int) ([]string, int) { + if len(msgs) == 0 { + return nil, 0 + } + out := make([]string, 0, len(msgs)) + total := 0 + for _, m := range msgs { + if m == nil { + continue + } + line := fmt.Sprintf("[%s]: %s", m.Role, m.Content) + if maxBytes > 0 && total+len(line)+1 > maxBytes { + break + } + out = append(out, line) + total += len(line) + 1 + } + return out, total +} + +// buildExistingMemoriesBlock formats the similar-memory neighbors for the +// extraction prompt and trims by total byte budget. Each memory's content is +// trimmed per-entry via perMemContentTrim (0 = no trim). The format is kept +// compact: `[Type, State, importance:%.1f] Content`. +func buildExistingMemoriesBlock(sims []*storage.SimilarityResult, maxBytes int, perMemContentTrim int) ([]string, int) { + if len(sims) == 0 { + return nil, 0 + } + out := make([]string, 0, len(sims)) + total := 0 + for _, sm := range sims { + if sm == nil || sm.Memory == nil { + continue + } + content := sm.Memory.Content + if perMemContentTrim > 0 && len(content) > perMemContentTrim { + content = content[:perMemContentTrim] + "..." + } + line := fmt.Sprintf("[%s, %s, importance:%.1f] %s", + sm.Memory.Type, sm.Memory.State, sm.Memory.Importance, content) + if maxBytes > 0 && total+len(line)+1 > maxBytes { + break + } + out = append(out, line) + total += len(line) + 1 + } + return out, total +} diff --git a/engine/engine_add_prompt_budget_test.go b/engine/engine_add_prompt_budget_test.go new file mode 100644 index 0000000..8409dba --- /dev/null +++ b/engine/engine_add_prompt_budget_test.go @@ -0,0 +1,109 @@ +// SPDX-License-Identifier: BSL-1.1 +// Copyright (c) 2026 Keyoku. All rights reserved. + +package engine + +import ( + "strings" + "testing" + + "github.com/keyoku-ai/keyoku-engine/storage" +) + +func TestBuildConversationContext_Unlimited(t *testing.T) { + msgs := []*storage.SessionMessage{ + {Role: "user", Content: "hello"}, + {Role: "assistant", Content: "hi there"}, + } + out, total := buildConversationContext(msgs, 0) + if len(out) != 2 { + t.Fatalf("len(out) = %d, want 2", len(out)) + } + if total == 0 { + t.Fatalf("total bytes = 0, want > 0") + } + if out[0] != "[user]: hello" { + t.Fatalf("out[0] = %q", out[0]) + } +} + +func TestBuildConversationContext_TrimsByBudget(t *testing.T) { + // Each formatted line is "[user]: msgN" ~= 13 bytes + 1 newline accounting + msgs := []*storage.SessionMessage{ + {Role: "user", Content: "msg1"}, + {Role: "user", Content: "msg2"}, + {Role: "user", Content: "msg3"}, + {Role: "user", Content: "msg4"}, + } + out, total := buildConversationContext(msgs, 30) + // 30 bytes budget should admit roughly 2 entries (~14 bytes each with newline) + if len(out) >= len(msgs) { + t.Fatalf("expected trimming, got all %d entries", len(out)) + } + if total > 30 { + t.Fatalf("total %d exceeds budget 30", total) + } +} + +func TestBuildExistingMemoriesBlock_TrimsPerMemoryContent(t *testing.T) { + longContent := strings.Repeat("x", 500) + sims := []*storage.SimilarityResult{ + {Memory: &storage.Memory{Type: storage.TypeEvent, State: "active", Importance: 0.8, Content: longContent}}, + } + out, _ := buildExistingMemoriesBlock(sims, 0, 100) + if len(out) != 1 { + t.Fatalf("len(out) = %d, want 1", len(out)) + } + // 100 trim + "..." = 103-char content. Line prefix adds more, but content portion must be bounded. + if !strings.Contains(out[0], strings.Repeat("x", 100)+"...") { + t.Fatalf("expected per-memory trim to 100 chars with ellipsis, got: %q", out[0]) + } + if strings.Count(out[0], "x") > 101 { + t.Fatalf("trim leaked: too many x in output: %d", strings.Count(out[0], "x")) + } +} + +func TestBuildExistingMemoriesBlock_TrimsByByteBudget(t *testing.T) { + sims := []*storage.SimilarityResult{ + {Memory: &storage.Memory{Type: storage.TypeEvent, State: "active", Importance: 0.8, Content: "one"}}, + {Memory: &storage.Memory{Type: storage.TypeEvent, State: "active", Importance: 0.8, Content: "two"}}, + {Memory: &storage.Memory{Type: storage.TypeEvent, State: "active", Importance: 0.8, Content: "three"}}, + } + // Each formatted line is ~"[event, active, importance:0.8] xxx" ~= 35 bytes + 1. + // Budget 40 should allow one, not more. + out, total := buildExistingMemoriesBlock(sims, 40, 0) + if len(out) != 1 { + t.Fatalf("len(out) = %d, want 1 under tight budget", len(out)) + } + if total > 40 { + t.Fatalf("total %d exceeds budget 40", total) + } +} + +func TestBuildExistingMemoriesBlock_NilSafe(t *testing.T) { + sims := []*storage.SimilarityResult{ + nil, + {Memory: nil}, + {Memory: &storage.Memory{Type: storage.TypeEvent, State: "active", Importance: 0.5, Content: "ok"}}, + } + out, _ := buildExistingMemoriesBlock(sims, 0, 0) + if len(out) != 1 { + t.Fatalf("len(out) = %d, want 1 (nil entries skipped)", len(out)) + } +} + +func TestBuildExistingMemoriesBlock_PreservesMostSimilarFirst(t *testing.T) { + // Input order mirrors store ordering (highest similarity first). We must + // keep that order so tight budgets retain the most relevant neighbors. + sims := []*storage.SimilarityResult{ + {Memory: &storage.Memory{Type: storage.TypeEvent, State: "active", Importance: 0.9, Content: "most_similar"}}, + {Memory: &storage.Memory{Type: storage.TypeEvent, State: "active", Importance: 0.5, Content: "least_similar"}}, + } + out, _ := buildExistingMemoriesBlock(sims, 50, 0) + if len(out) < 1 { + t.Fatalf("expected at least 1 entry") + } + if !strings.Contains(out[0], "most_similar") { + t.Fatalf("first entry should be most_similar, got: %q", out[0]) + } +}