From 1c587cbb7a978dd4229d60d4d6c062470847c246 Mon Sep 17 00:00:00 2001 From: Audrey Tang Date: Mon, 28 Sep 2026 23:41:00 -0400 Subject: [PATCH 1/5] feat(memory): keep the embedding model loaded and add gc --compact A cold recall paid twice before any ranking started. Ollama unloads a model five minutes after its last request, so an agent that recalls once per turn often reloads a multi-gigabyte embedding model first; the Ollama client now sends keep_alive (MNEMON_EMBED_KEEP_ALIVE, default 30m; not sent to OpenAI-compatible servers). Separately, recall bumps access counters on every hit and nothing ever vacuumed, so a long-lived store ends up with table pages scattered across the file and cold reads run at random-I/O speed. `mnemon gc --compact` runs VACUUM and reports the size before and after. On a large long-lived store on an Apple M4 laptop, VACUUM alone cut a cold-cache recall from about 3.6s to 1.4s. Validated with make test plus new tests for keep_alive (default, override, omitted for OpenAI) and compaction. --- README.md | 1 + cmd/memory/gc.go | 28 +++++++++++++++- docs/USAGE.md | 4 +++ docs/zh/README.md | 1 + docs/zh/USAGE.md | 4 +++ internal/memory/embed/ollama.go | 45 ++++++++++++++++--------- internal/memory/embed/ollama_test.go | 48 +++++++++++++++++++++++++++ internal/memory/store/compact.go | 36 ++++++++++++++++++++ internal/memory/store/compact_test.go | 39 ++++++++++++++++++++++ 9 files changed, 190 insertions(+), 16 deletions(-) create mode 100644 internal/memory/store/compact.go create mode 100644 internal/memory/store/compact_test.go diff --git a/README.md b/README.md index dd292d51c..477f2c327 100644 --- a/README.md +++ b/README.md @@ -485,6 +485,7 @@ reported by ID in the triggering command's `auto_pruned_ids` field. | `MNEMON_EMBED_MODEL` | `nomic-embed-text` | Embedding model name | | `MNEMON_EMBED_PROTOCOL` | *(auto-detect)* | `ollama` or `openai`; auto-detected from an endpoint ending in `/v1` | | `MNEMON_EMBED_API_KEY` | *(none)* | Bearer token for OpenAI-compatible servers (oMLX, vLLM, etc.) | +| `MNEMON_EMBED_KEEP_ALIVE` | `30m` | How long Ollama keeps the embedding model loaded after each request (ignored for OpenAI-compatible servers) | | `MNEMON_EMBED_DIMENSIONS` | *(native)* | Optional Matryoshka dimension truncation | The embedding client speaks the Ollama API by default and the diff --git a/cmd/memory/gc.go b/cmd/memory/gc.go index 4e457788f..9c2400f19 100644 --- a/cmd/memory/gc.go +++ b/cmd/memory/gc.go @@ -13,6 +13,7 @@ var ( gcThreshold float64 gcLimit int gcKeepID string + gcCompact bool ) var gcCmd = &cobra.Command{ @@ -27,7 +28,13 @@ Suggest mode (default): Keep mode: mnemon gc --keep - Boosts an insight's retention (access_count +3, refreshes timestamp).`, + Boosts an insight's retention (access_count +3, refreshes timestamp). + +Compact mode: + mnemon gc --compact + Rewrites the database with VACUUM. Recall bumps access counters on every + hit, which scatters table pages across the file and makes cold recalls + slow; compacting restores sequential layout. Run it weekly or so.`, RunE: func(cmd *cobra.Command, args []string) error { if err := requirePositiveLimit("--limit", gcLimit); err != nil { return err @@ -42,6 +49,24 @@ Keep mode: } defer db.Close() + if gcCompact { + if err := requireWritableDB(db, "gc --compact"); err != nil { + return err + } + before, after, err := db.Compact() + if err != nil { + return err + } + db.LogOp("gc_compact", "", fmt.Sprintf("bytes %d -> %d", before, after)) + enc := json.NewEncoder(os.Stdout) + enc.SetIndent("", " ") + return enc.Encode(map[string]interface{}{ + "status": "compacted", + "bytes_before": before, + "bytes_after": after, + }) + } + // Keep mode: boost retention for a specific insight if gcKeepID != "" { if err := requireWritableDB(db, "gc --keep"); err != nil { @@ -106,5 +131,6 @@ func init() { gcCmd.Flags().Float64Var(&gcThreshold, "threshold", 0.5, "effective_importance threshold (insights below this are candidates)") gcCmd.Flags().IntVar(&gcLimit, "limit", 20, "max candidates to return") gcCmd.Flags().StringVar(&gcKeepID, "keep", "", "boost retention for this insight ID") + gcCmd.Flags().BoolVar(&gcCompact, "compact", false, "rewrite the database with VACUUM to speed up cold recalls") rootCmd.AddCommand(gcCmd) } diff --git a/docs/USAGE.md b/docs/USAGE.md index 765ac460c..12bcb55ed 100644 --- a/docs/USAGE.md +++ b/docs/USAGE.md @@ -266,6 +266,9 @@ mnemon gc --threshold 0.5 --limit 20 # GC keep — boost an insight's retention mnemon gc --keep + +# GC compact — rewrite the store with VACUUM to keep cold recalls fast +mnemon gc --compact ``` ### Store Management @@ -360,6 +363,7 @@ Nodes are colored by category (decision, fact, insight, preference, context); ed | `MNEMON_EMBED_MODEL` | `nomic-embed-text` | Embedding model | | `MNEMON_EMBED_PROTOCOL` | (auto-detect) | `ollama` or `openai`; endpoints ending in `/v1` select `openai` | | `MNEMON_EMBED_API_KEY` | (none) | Bearer token for OpenAI-compatible servers | +| `MNEMON_EMBED_KEEP_ALIVE` | 30m | How long Ollama keeps the embedding model loaded after each request | | `MNEMON_EMBED_DIMENSIONS` | (native) | Embedding dimensions; set to truncate (e.g., `256` for Matryoshka models) | | `MNEMON_MAX_INSIGHTS` | `1000` | Active-insight ceiling before auto-pruning starts; `0` disables auto-pruning | | `MNEMON_AUTO_PRUNE_MIN_AGE` | `24h` | Minimum age before automatic pruning; accepts durations such as `24h`, integer days such as `7d`, or `0` to disable the grace period | diff --git a/docs/zh/README.md b/docs/zh/README.md index ff26ac1e6..cc26cf1bf 100644 --- a/docs/zh/README.md +++ b/docs/zh/README.md @@ -420,6 +420,7 @@ Sub-agent 委派是可选执行策略。当 runtime 支持时，主 agent 可以 | `MNEMON_EMBED_MODEL` | `nomic-embed-text` | 嵌入模型名称 | | `MNEMON_EMBED_PROTOCOL` | *（自动探测）* | `ollama` 或 `openai`；端点以 `/v1` 结尾时自动切换 | | `MNEMON_EMBED_API_KEY` | *（无）* | OpenAI 兼容服务器（oMLX、vLLM 等）的 Bearer 令牌 | +| `MNEMON_EMBED_KEEP_ALIVE` | `30m` | 每次请求后 Ollama 保持嵌入模型加载的时长（OpenAI 兼容服务器忽略此项） | | `MNEMON_EMBED_DIMENSIONS` | *（原生维度）* | 可选的 Matryoshka 维度截断 | 每次自动删除均为软删除，以 `prune` 操作记录到 oplog，并通过触发命令的 diff --git a/docs/zh/USAGE.md b/docs/zh/USAGE.md index 98682e265..e5846e392 100644 --- a/docs/zh/USAGE.md +++ b/docs/zh/USAGE.md @@ -238,6 +238,9 @@ mnemon gc --threshold 0.5 --limit 20 # GC keep — 提升某个洞察的保留度 mnemon gc --keep + +# GC compact — 以 VACUUM 重写存储，保持冷启动召回速度 +mnemon gc --compact ``` ### 记忆体管理 @@ -330,6 +333,7 @@ open graph.html | `MNEMON_EMBED_MODEL` | `nomic-embed-text` | 嵌入模型 | | `MNEMON_EMBED_PROTOCOL` | （自动探测） | `ollama` 或 `openai`；以 `/v1` 结尾的端点自动选择 `openai` | | `MNEMON_EMBED_API_KEY` | （无） | OpenAI 兼容服务器的 Bearer 令牌 | +| `MNEMON_EMBED_KEEP_ALIVE` | 30m | 每次请求后 Ollama 保持嵌入模型加载的时长 | | `MNEMON_EMBED_DIMENSIONS` | (原生维度) | 嵌入向量维度；可设置截断值（例如 Matryoshka 模型使用 `256`） | | `MNEMON_MAX_INSIGHTS` | `1000` | 触发自动清理的活跃洞察数量上限；设为 `0` 可关闭自动清理 | | `MNEMON_AUTO_PRUNE_MIN_AGE` | `24h` | 可被自动清理前的最短存活时间；支持 `24h`、整数天 `7d`，设为 `0` 可关闭保护期 | diff --git a/internal/memory/embed/ollama.go b/internal/memory/embed/ollama.go index fe4378064..f71583225 100644 --- a/internal/memory/embed/ollama.go +++ b/internal/memory/embed/ollama.go @@ -31,21 +31,27 @@ const DefaultModel = "nomic-embed-text" // DefaultEndpoint is the default Ollama API endpoint. const DefaultEndpoint = "http://localhost:11434" +// DefaultKeepAlive is how long Ollama keeps the embedding model loaded after +// a request. Ollama's own five-minute default lets a large embedding model +// unload between agent turns, and reloading it costs seconds on a cold recall. +const DefaultKeepAlive = "30m" + // Client communicates with an embedding server (an Ollama instance or an // OpenAI-compatible server) for embedding generation. type Client struct { - endpoint string - model string - dims int // 0 means use native dimensions - apiKey string - protocol Protocol - http *http.Client + endpoint string + model string + dims int // 0 means use native dimensions + apiKey string + keepAlive string + protocol Protocol + http *http.Client } // NewClient creates an embedding client. // It checks MNEMON_EMBED_ENDPOINT, MNEMON_EMBED_MODEL, -// MNEMON_EMBED_DIMENSIONS, MNEMON_EMBED_API_KEY, and -// MNEMON_EMBED_PROTOCOL env vars. +// MNEMON_EMBED_DIMENSIONS, MNEMON_EMBED_API_KEY, MNEMON_EMBED_KEEP_ALIVE, +// and MNEMON_EMBED_PROTOCOL env vars. func NewClient() *Client { return NewClientWithModel("") } @@ -75,6 +81,10 @@ func NewClientWithModel(model string) *Client { dims = v } } + keepAlive := os.Getenv("MNEMON_EMBED_KEEP_ALIVE") + if keepAlive == "" { + keepAlive = DefaultKeepAlive + } protocol := ProtocolOllama explicit := false if p := os.Getenv("MNEMON_EMBED_PROTOCOL"); p != "" { @@ -97,11 +107,12 @@ func NewClientWithModel(model string) *Client { } } return &Client{ - endpoint: endpoint, - model: model, - dims: dims, - apiKey: os.Getenv("MNEMON_EMBED_API_KEY"), - protocol: protocol, + endpoint: endpoint, + model: model, + dims: dims, + apiKey: os.Getenv("MNEMON_EMBED_API_KEY"), + keepAlive: keepAlive, + protocol: protocol, http: &http.Client{ Timeout: 30 * time.Second, Transport: &http.Transport{ @@ -216,6 +227,9 @@ func (c *Client) embedWithContext(ctx context.Context, text string) ([]float64, if c.dims > 0 { req.Dimensions = c.dims } + if c.protocol == ProtocolOllama { + req.KeepAlive = c.keepAlive + } body, err := json.Marshal(req) if err != nil { return nil, fmt.Errorf("marshal request: %w", err) @@ -267,6 +281,7 @@ type embedRequest struct { Model string `json:"model"` Input string `json:"input"` Dimensions int `json:"dimensions,omitempty"` + KeepAlive string `json:"keep_alive,omitempty"` } type ollamaEmbedResponse struct { @@ -288,8 +303,8 @@ func (c *Client) embedRequestRoute() string { } // Embed generates an embedding vector for the given text. -// The request body is identical for both protocols; only the endpoint -// path and the response shape differ. +// The request body is identical for both protocols apart from Ollama's +// keep_alive; the endpoint path and the response shape differ. func (c *Client) Embed(text string) ([]float64, error) { return c.embedWithContext(context.Background(), text) } diff --git a/internal/memory/embed/ollama_test.go b/internal/memory/embed/ollama_test.go index 006befc1a..2bf308e6f 100644 --- a/internal/memory/embed/ollama_test.go +++ b/internal/memory/embed/ollama_test.go @@ -1,6 +1,7 @@ package embed import ( + "encoding/json" "net/http" "net/http/httptest" "testing" @@ -119,3 +120,50 @@ func TestOllamaEndpointWithTrailingSlash(t *testing.T) { t.Fatalf("expected 3 dims, got %d", len(vec)) } } + +// keepAliveSent embeds once against a stub server under the given protocol +// and returns the keep_alive field of the request body ("" when absent). +func keepAliveSent(t *testing.T, protocol, route string) string { + t.Helper() + var got map[string]any + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if err := json.NewDecoder(r.Body).Decode(&got); err != nil { + t.Errorf("decode request: %v", err) + } + w.Header().Set("Content-Type", "application/json") + if protocol == "openai" { + _, _ = w.Write([]byte(`{"data":[{"embedding":[0.1]}]}`)) + return + } + _, _ = w.Write([]byte(`{"embeddings":[[0.1]]}`)) + })) + defer srv.Close() + t.Setenv("MNEMON_EMBED_PROTOCOL", protocol) + t.Setenv("MNEMON_EMBED_ENDPOINT", srv.URL+route) + if _, err := NewClient().Embed("hello"); err != nil { + t.Fatalf("Embed: %v", err) + } + s, _ := got["keep_alive"].(string) + return s +} + +func TestOllamaEmbed_SendsDefaultKeepAlive(t *testing.T) { + t.Setenv("MNEMON_EMBED_KEEP_ALIVE", "") + if got := keepAliveSent(t, "ollama", ""); got != DefaultKeepAlive { + t.Fatalf("keep_alive = %q, want %q", got, DefaultKeepAlive) + } +} + +func TestOllamaEmbed_KeepAliveEnvOverride(t *testing.T) { + t.Setenv("MNEMON_EMBED_KEEP_ALIVE", "2h") + if got := keepAliveSent(t, "ollama", ""); got != "2h" { + t.Fatalf("keep_alive = %q, want %q", got, "2h") + } +} + +func TestOpenAIEmbed_OmitsKeepAlive(t *testing.T) { + t.Setenv("MNEMON_EMBED_KEEP_ALIVE", "") + if got := keepAliveSent(t, "openai", "/v1"); got != "" { + t.Fatalf("keep_alive = %q, want it omitted", got) + } +} diff --git a/internal/memory/store/compact.go b/internal/memory/store/compact.go new file mode 100644 index 000000000..4d7b674ba --- /dev/null +++ b/internal/memory/store/compact.go @@ -0,0 +1,36 @@ +package store + +import ( + "fmt" + "os" +) + +// Compact rewrites the database with VACUUM and returns its size in bytes +// before and after. Recall updates access counters on every hit, so pages of +// the insights and edges tables end up scattered across the file; a cold recall +// then reads them at random-I/O speed. VACUUM restores table order and drops +// free pages. It needs an exclusive lock and briefly doubles disk usage. +func (db *DB) Compact() (before, after int64, err error) { + if db.readOnly { + return 0, 0, fmt.Errorf("compact: database is read-only") + } + if _, err := db.conn.Exec(`PRAGMA wal_checkpoint(TRUNCATE)`); err != nil { + return 0, 0, fmt.Errorf("checkpoint: %w", err) + } + before = db.fileSize() + if _, err := db.conn.Exec(`VACUUM`); err != nil { + return 0, 0, fmt.Errorf("vacuum: %w", err) + } + if _, err := db.conn.Exec(`PRAGMA wal_checkpoint(TRUNCATE)`); err != nil { + return 0, 0, fmt.Errorf("checkpoint: %w", err) + } + return before, db.fileSize(), nil +} + +func (db *DB) fileSize() int64 { + fi, err := os.Stat(db.path) + if err != nil { + return 0 + } + return fi.Size() +} diff --git a/internal/memory/store/compact_test.go b/internal/memory/store/compact_test.go new file mode 100644 index 000000000..1144baedb --- /dev/null +++ b/internal/memory/store/compact_test.go @@ -0,0 +1,39 @@ +package store + +import ( + "fmt" + "testing" +) + +func TestCompact_ShrinksAfterDeletes(t *testing.T) { + db, err := Open(t.TempDir()) + if err != nil { + t.Fatal(err) + } + defer db.Close() + if _, err := db.conn.Exec(`CREATE TABLE filler (b BLOB)`); err != nil { + t.Fatal(err) + } + for i := 0; i < 200; i++ { + if _, err := db.conn.Exec(`INSERT INTO filler VALUES (zeroblob(8192))`); err != nil { + t.Fatal(fmt.Errorf("insert %d: %w", i, err)) + } + } + if _, err := db.conn.Exec(`DELETE FROM filler`); err != nil { + t.Fatal(err) + } + before, after, err := db.Compact() + if err != nil { + t.Fatalf("Compact: %v", err) + } + if after >= before { + t.Fatalf("size after = %d, want below %d", after, before) + } + var free int + if err := db.conn.QueryRow(`PRAGMA freelist_count`).Scan(&free); err != nil { + t.Fatal(err) + } + if free != 0 { + t.Fatalf("freelist_count = %d, want 0", free) + } +} From 1648e886a60d1b19ba03abdc272087d53320aec5 Mon Sep 17 00:00:00 2001 From: Audrey Tang Date: Mon, 28 Sep 2026 23:41:28 -0400 Subject: [PATCH 2/5] fix(recall): cut per-node SQLite reads that made cold recall slow Beam search fetched each neighbour with GetInsightByID although recall had already loaded every active insight, re-read a node's edges for every anchor whose neighbourhood reached it, and read edges through single-column indexes that sent each row to a random page of the edges table. Recall also scanned the insights table a second time for the known-entity set. Neighbours now resolve from the loaded snapshot, edges are cached per recall, and the entity set comes from the same snapshot (IntentAwareRecallFrom, store.KnownEntities). New covering indexes on (source_id|target_id, other end, edge_type, weight) answer GetNeighborEdges without touching the table and replace the single-column edge indexes. Its fixed order also makes the visit-budgeted beam deterministic. The superseded check reads a small partial index of supersedes edges once instead of probing the target index per candidate, falling back to the scoped query on stores a read-only open has not migrated. Both open modes set mmap_size. On a large long-lived store (Apple M4, model already loaded, medians of 5-7 runs over three queries) recall went from 0.82-0.87s to 0.53-0.56s warm and from 1.92-1.96s to 1.22-1.53s after purging the page cache, with the same top-10 set. Validated with make test plus tests for neighbour order, query plans, the known-entity set and the missing-index fallback. --- cmd/memory/recall.go | 13 +- internal/memory/search/integration_test.go | 21 +++- internal/memory/search/recall.go | 44 +++++-- internal/memory/store/db.go | 36 +++++- internal/memory/store/edge.go | 87 +++++++++++-- internal/memory/store/node.go | 14 +++ internal/memory/store/recall_index_test.go | 134 +++++++++++++++++++++ 7 files changed, 316 insertions(+), 33 deletions(-) create mode 100644 internal/memory/store/recall_index_test.go diff --git a/cmd/memory/recall.go b/cmd/memory/recall.go index 9acc07cf0..7b0b88ab1 100644 --- a/cmd/memory/recall.go +++ b/cmd/memory/recall.go @@ -190,12 +190,15 @@ meta.intent and meta.intent_source (auto or override). --basic bypasses intent.` // Extract query entities at cmd layer (avoid graph->search circular dep). // Load the known-entity set so the indexed extractor's fourth path can // admit user vocabulary (single-segment CamelCase, lowercase project - // names) that techDictionary does not cover. The lookup is read-only; - // on error we fall through to the default regex+dictionary extractor. - knownEntities, _ := db.LoadKnownEntities() - queryEntities := graph.ExtractEntitiesIndexed(keyword, knownEntities) + // names) that techDictionary does not cover. The set is derived from the + // same insight snapshot recall ranks, so the table is scanned once. + all, err := db.GetAllActiveInsights() + if err != nil { + return fmt.Errorf("recall: %w", err) + } + queryEntities := graph.ExtractEntitiesIndexed(keyword, store.KnownEntities(all)) - resp, err := search.IntentAwareRecall(db, keyword, queryVec, queryEntities, recLimit, intentOverride) + resp, err := search.IntentAwareRecallFrom(db, all, keyword, queryVec, queryEntities, recLimit, intentOverride) if err != nil { return fmt.Errorf("recall: %w", err) } diff --git a/internal/memory/search/integration_test.go b/internal/memory/search/integration_test.go index 9453db33d..d24eed624 100644 --- a/internal/memory/search/integration_test.go +++ b/internal/memory/search/integration_test.go @@ -442,7 +442,7 @@ func TestBeamSearchFromAnchor_ScorePropagation(t *testing.T) { params := TraversalParams{BeamWidth: 10, MaxDepth: 3, MaxVisited: 100} scoreMap["bs-1"] = 1.0 - beamSearchFromAnchor(db, "bs-1", 1.0, nil, weights, params, scoreMap, viaMap, insightMap, nil) + beamSearchFromAnchor(db, "bs-1", 1.0, nil, weights, params, scoreMap, viaMap, insightMap, nil, activeMap(t, db), map[string][]*model.Edge{}) // Neighbor should be discovered with score > 0 if _, ok := scoreMap["bs-2"]; !ok { @@ -484,7 +484,7 @@ func TestBeamSearchFromAnchor_BeamWidthPruning(t *testing.T) { params := TraversalParams{BeamWidth: 3, MaxDepth: 3, MaxVisited: 500} scoreMap["bw-center"] = 1.0 - beamSearchFromAnchor(db, "bw-center", 1.0, nil, weights, params, scoreMap, viaMap, insightMap, nil) + beamSearchFromAnchor(db, "bw-center", 1.0, nil, weights, params, scoreMap, viaMap, insightMap, nil, activeMap(t, db), map[string][]*model.Edge{}) // Count deep nodes discovered — should be limited by beam width deepCount := 0 @@ -522,7 +522,7 @@ func TestBeamSearchFromAnchor_MaxVisitedBudget(t *testing.T) { params := TraversalParams{BeamWidth: 10, MaxDepth: 20, MaxVisited: 5} scoreMap["mv-0"] = 1.0 - beamSearchFromAnchor(db, "mv-0", 1.0, nil, weights, params, scoreMap, viaMap, insightMap, nil) + beamSearchFromAnchor(db, "mv-0", 1.0, nil, weights, params, scoreMap, viaMap, insightMap, nil, activeMap(t, db), map[string][]*model.Edge{}) // scoreMap includes the anchor itself, so discovered nodes (excluding anchor) should be <= 4 discovered := len(scoreMap) - 1 // subtract anchor @@ -612,3 +612,18 @@ func scoreOf(results []RecallResult, id string) float64 { } return 0 } + +// activeMap indexes the store's active insights the way IntentAwareRecallFrom +// hands them to beam search. +func activeMap(t *testing.T, db *store.DB) map[string]*model.Insight { + t.Helper() + all, err := db.GetAllActiveInsights() + if err != nil { + t.Fatal(err) + } + m := make(map[string]*model.Insight, len(all)) + for _, ins := range all { + m[ins.ID] = ins + } + return m +} diff --git a/internal/memory/search/recall.go b/internal/memory/search/recall.go index ebe70ed89..c0721f443 100644 --- a/internal/memory/search/recall.go +++ b/internal/memory/search/recall.go @@ -129,6 +129,18 @@ type RecallResult struct { // 6. Sparse hint detection func IntentAwareRecall(db *store.DB, query string, queryVec []float64, queryEntities []string, limit int, intentOverride *Intent) (RecallResponse, error) { + all, err := db.GetAllActiveInsights() + if err != nil { + return RecallResponse{}, err + } + return IntentAwareRecallFrom(db, all, query, queryVec, queryEntities, limit, intentOverride) +} + +// IntentAwareRecallFrom is IntentAwareRecall over a caller-loaded snapshot of +// the active insights, so a caller that already scanned them (for example to +// build the known-entity set) does not scan the table twice. +func IntentAwareRecallFrom(db *store.DB, all []*model.Insight, query string, queryVec []float64, + queryEntities []string, limit int, intentOverride *Intent) (RecallResponse, error) { // Step 1: Intent determination var intent Intent @@ -143,10 +155,9 @@ func IntentAwareRecall(db *store.DB, query string, queryVec []float64, weights := GetWeights(intent) params := getTraversalParams(intent) - // Get all active insights - all, err := db.GetAllActiveInsights() - if err != nil { - return RecallResponse{}, err + activeByID := make(map[string]*model.Insight, len(all)) + for _, ins := range all { + activeByID[ins.ID] = ins } // Pre-load all embeddings once (avoids N+1 queries in beam search and reranking). @@ -257,9 +268,11 @@ func IntentAwareRecall(db *store.DB, query string, queryVec []float64, insightMap[id] = a.insight } - // Step 3: Beam search from each anchor + // Step 3: Beam search from each anchor. Anchors' neighbourhoods overlap + // heavily, so each node's edges are read from the store once per recall. + edgeCache := make(map[string][]*model.Edge) for id, a := range anchorMap { - beamSearchFromAnchor(db, id, a.score, queryVec, weights, params, scoreMap, viaMap, insightMap, embedCache) + beamSearchFromAnchor(db, id, a.score, queryVec, weights, params, scoreMap, viaMap, insightMap, embedCache, activeByID, edgeCache) } traversedCount := len(scoreMap) @@ -508,6 +521,9 @@ func causalTopologicalSort(db *store.DB, results []RecallResult) []RecallResult // beamSearchFromAnchor performs beam search starting from a single anchor node. // It uses a priority queue to keep the top beamWidth candidates at each depth level. // embedCache provides pre-loaded embedding vectors (nil = no embeddings). +// activeByID resolves neighbours without a query; a neighbour missing from it +// is soft-deleted and never enters insightMap. edgeCache memoises edge reads +// across anchors within one recall. func beamSearchFromAnchor( db *store.DB, startID string, @@ -519,6 +535,8 @@ func beamSearchFromAnchor( viaMap map[string]string, insightMap map[string]*model.Insight, embedCache map[string][]float64, + activeByID map[string]*model.Insight, + edgeCache map[string][]*model.Edge, ) { visited := map[string]bool{startID: true} totalVisited := 1 @@ -545,9 +563,14 @@ func beamSearchFromAnchor( break } - edges, err := db.GetEdgesByNode(cur.id) - if err != nil { - continue + edges, cached := edgeCache[cur.id] + if !cached { + var err error + edges, err = db.GetNeighborEdges(cur.id) + if err != nil { + continue + } + edgeCache[cur.id] = edges } for _, e := range edges { @@ -578,8 +601,7 @@ func beamSearchFromAnchor( scoreMap[neighborID] = neighborScore viaMap[neighborID] = string(e.EdgeType) if _, loaded := insightMap[neighborID]; !loaded { - ins, err := db.GetInsightByID(neighborID) - if err == nil && ins != nil { + if ins, ok := activeByID[neighborID]; ok { insightMap[neighborID] = ins } } diff --git a/internal/memory/store/db.go b/internal/memory/store/db.go index 9cb57b991..20ad6106b 100644 --- a/internal/memory/store/db.go +++ b/internal/memory/store/db.go @@ -231,6 +231,7 @@ func OpenReadOnly(dataDir string) (*DB, error) { query.Set("mode", "ro") query.Set("immutable", "1") query.Add("_pragma", "foreign_keys(1)") + query.Add("_pragma", "mmap_size(268435456)") dsn.RawQuery = query.Encode() conn, err := sql.Open("sqlite", dsn.String()) @@ -248,7 +249,9 @@ func Open(dataDir string) (*DB, error) { } dbPath := filepath.Join(dataDir, "mnemon.db") - conn, err := sql.Open("sqlite", dbPath+"?_pragma=journal_mode(WAL)&_pragma=foreign_keys(1)") + // mmap lets cold reads skip a syscall and a copy per page; recall's random + // edge and insight reads are where a cold open spends its time. + conn, err := sql.Open("sqlite", dbPath+"?_pragma=journal_mode(WAL)&_pragma=foreign_keys(1)&_pragma=mmap_size(268435456)") if err != nil { return nil, fmt.Errorf("open database: %w", err) } @@ -312,8 +315,6 @@ CREATE INDEX IF NOT EXISTS idx_insights_importance ON insights(importance); CREATE INDEX IF NOT EXISTS idx_insights_created ON insights(created_at); CREATE INDEX IF NOT EXISTS idx_insights_deleted ON insights(deleted_at); CREATE INDEX IF NOT EXISTS idx_insights_source ON insights(source); -CREATE INDEX IF NOT EXISTS idx_edges_source ON edges(source_id); -CREATE INDEX IF NOT EXISTS idx_edges_target ON edges(target_id); CREATE INDEX IF NOT EXISTS idx_edges_type ON edges(edge_type); CREATE INDEX IF NOT EXISTS idx_edges_source_type ON edges(source_id, edge_type); CREATE INDEX IF NOT EXISTS idx_edges_target_type ON edges(target_id, edge_type); @@ -399,6 +400,10 @@ CREATE INDEX IF NOT EXISTS idx_oplog_created ON oplog(created_at); return fmt.Errorf("add supersedes edge type: %w", err) } + if err := db.migrateEdgeRecallIndexes(); err != nil { + return fmt.Errorf("edge recall indexes: %w", err) + } + // One-time cleanup: soft-delete narrative category insights from legacy databases. // Only runs the UPDATE when narrative insights actually exist (avoids needless writes). var narrativeCount int @@ -412,6 +417,31 @@ CREATE INDEX IF NOT EXISTS idx_oplog_created ON oplog(created_at); return nil } +// migrateEdgeRecallIndexes gives recall's hot edge queries indexes that answer +// them without touching the edges table. Beam search reads every edge of each +// visited node; with single-column indexes each hit was a random read into the +// table, which dominates a cold-cache recall on a large graph. The covering +// indexes supersede the single-column ones (the (id, edge_type) indexes still +// serve plain id lookups), so those are dropped; the table-rebuild migrations +// above recreate them, hence this runs last on every open. +// +// The partial index holds only 'supersedes' edges, so recall's superseded +// check reads a few hundred KB instead of probing the whole target index. +func (db *DB) migrateEdgeRecallIndexes() error { + for _, stmt := range []string{ + `CREATE INDEX IF NOT EXISTS idx_edges_source_cover ON edges(source_id, target_id, edge_type, weight)`, + `CREATE INDEX IF NOT EXISTS idx_edges_target_cover ON edges(target_id, source_id, edge_type, weight)`, + `CREATE INDEX IF NOT EXISTS idx_edges_supersedes ON edges(target_id, source_id) WHERE edge_type = 'supersedes'`, + `DROP INDEX IF EXISTS idx_edges_source`, + `DROP INDEX IF EXISTS idx_edges_target`, + } { + if _, err := db.conn.Exec(stmt); err != nil { + return fmt.Errorf("%s: %w", stmt, err) + } + } + return nil +} + // addColumnIfNotExists runs an ALTER TABLE ADD COLUMN statement, // ignoring "duplicate column" errors (column already exists). func addColumnIfNotExists(conn *sql.DB, stmt string) error { diff --git a/internal/memory/store/edge.go b/internal/memory/store/edge.go index 8dca1aa83..a7055e77a 100644 --- a/internal/memory/store/edge.go +++ b/internal/memory/store/edge.go @@ -39,24 +39,89 @@ func (db *DB) GetEdgesByNode(nodeID string) ([]*model.Edge, error) { return scanEdges(rows) } +// GetNeighborEdges returns every edge touching nodeID with the columns beam +// search scores on (metadata and created_at are left zero). Both halves are +// answered from the covering indexes alone, and the fixed order makes recall +// deterministic: beam search stops at a visit budget, so which neighbours it +// reaches depends on the order edges arrive in. +func (db *DB) GetNeighborEdges(nodeID string) ([]*model.Edge, error) { + rows, err := db.execer().Query( + `SELECT source_id, target_id, edge_type, weight FROM edges WHERE source_id = ? + UNION ALL + SELECT source_id, target_id, edge_type, weight FROM edges WHERE target_id = ? AND source_id != ? + ORDER BY weight DESC, source_id, target_id, edge_type`, + nodeID, nodeID, nodeID) + if err != nil { + return nil, err + } + defer rows.Close() + var results []*model.Edge + for rows.Next() { + var e model.Edge + var edgeType string + if err := rows.Scan(&e.SourceID, &e.TargetID, &edgeType, &e.Weight); err != nil { + return nil, err + } + e.EdgeType = model.EdgeType(edgeType) + results = append(results, &e) + } + return results, rows.Err() +} + +// GetSupersededIDs returns which of the given ids are the target of at least +// one 'supersedes' edge from a different insight, i.e. which of them some +// other insight claims to replace. Recall uses this to demote stale content; +// the rows are kept so the lineage stays inspectable. +// +// It reads every supersedes edge from the partial idx_edges_supersedes (a few +// thousand rows, a few hundred KB) and filters to the caller's ids in memory: +// probing each candidate through the whole-graph target index costs a random +// read per candidate on a cold cache. SQLite's planner does not pick a partial +// index by itself, hence INDEXED BY. A read-only open skips migrations, so a +// store no writable open has touched since the index was introduced lacks it; +// that store falls back to the scoped lookup. +func (db *DB) GetSupersededIDs(ids []string) (map[string]bool, error) { + superseded := make(map[string]bool) + if len(ids) == 0 { + return superseded, nil + } + want := make(map[string]bool, len(ids)) + for _, id := range ids { + want[id] = true + } + ex := db.execer() + rows, err := ex.Query( + `SELECT target_id FROM edges INDEXED BY idx_edges_supersedes + WHERE edge_type = 'supersedes' AND source_id != target_id`) + if err != nil && strings.Contains(err.Error(), "no such index") { + return scopedSupersededIDs(ex, ids) + } + if err != nil { + return nil, err + } + defer rows.Close() + for rows.Next() { + var id string + if err := rows.Scan(&id); err != nil { + return nil, err + } + if want[id] { + superseded[id] = true + } + } + return superseded, rows.Err() +} + // supersededLookupChunk bounds how many ids go into one IN clause. SQLite's // host-parameter ceiling is 32766 on current builds and 999 on older ones; // 500 stays inside both. Recall's candidate set is normally far smaller, so // the loop below runs once. const supersededLookupChunk = 500 -// GetSupersededIDs returns which of the given ids are the target of at least -// one 'supersedes' edge, i.e. which of them some other insight claims to -// replace. Recall uses this to demote stale content; the rows are kept so the -// lineage stays inspectable. -// -// The lookup is scoped to the ids the caller holds. Reading every supersedes -// edge in the store would cost time proportional to its whole supersession -// history on a path that only needs a verdict for the current candidates, -// and idx_edges_target_type answers the scoped form from the index. -func (db *DB) GetSupersededIDs(ids []string) (map[string]bool, error) { +// scopedSupersededIDs answers GetSupersededIDs through idx_edges_target_type +// for stores that predate idx_edges_supersedes. +func scopedSupersededIDs(ex dbExecer, ids []string) (map[string]bool, error) { superseded := make(map[string]bool) - ex := db.execer() for start := 0; start < len(ids); start += supersededLookupChunk { end := min(start+supersededLookupChunk, len(ids)) if err := collectSupersededIDs(ex, ids[start:end], superseded); err != nil { diff --git a/internal/memory/store/node.go b/internal/memory/store/node.go index bc4e1fff5..72c1fb294 100644 --- a/internal/memory/store/node.go +++ b/internal/memory/store/node.go @@ -245,6 +245,20 @@ func (db *DB) LoadKnownEntities() (map[string]bool, error) { return known, nil } +// KnownEntities is LoadKnownEntities over insights the caller already loaded +// (GetAllActiveInsights), sparing recall a second scan of the table. +func KnownEntities(active []*model.Insight) map[string]bool { + known := make(map[string]bool) + for _, ins := range active { + for _, e := range ins.Entities { + if e != "" { + known[e] = true + } + } + } + return known +} + // IncrementAccessCount bumps the access count and refreshes last_accessed_at. // No-op when the database is read-only. func (db *DB) IncrementAccessCount(id string) error { diff --git a/internal/memory/store/recall_index_test.go b/internal/memory/store/recall_index_test.go new file mode 100644 index 000000000..253450271 --- /dev/null +++ b/internal/memory/store/recall_index_test.go @@ -0,0 +1,134 @@ +package store + +import ( + "reflect" + "strings" + "testing" + "time" + + "github.com/mnemon-dev/mnemon/internal/memory/model" +) + +func queryPlan(t *testing.T, db *DB, query string, args ...any) string { + t.Helper() + rows, err := db.conn.Query("EXPLAIN QUERY PLAN "+query, args...) + if err != nil { + t.Fatal(err) + } + defer rows.Close() + var plan []string + for rows.Next() { + var id, parent, unused int + var detail string + if err := rows.Scan(&id, &parent, &unused, &detail); err != nil { + t.Fatal(err) + } + plan = append(plan, detail) + } + return strings.Join(plan, "\n") +} + +func TestGetNeighborEdges_BothDirectionsHeaviestFirst(t *testing.T) { + db := testDB(t) + for _, id := range []string{"n-a", "n-b", "n-c", "n-d"} { + if err := db.InsertInsight(makeInsight(id, id, 3)); err != nil { + t.Fatal(err) + } + } + now := time.Now().UTC() + for _, e := range []*model.Edge{ + {SourceID: "n-a", TargetID: "n-b", EdgeType: model.EdgeSemantic, Weight: 0.2}, + {SourceID: "n-c", TargetID: "n-a", EdgeType: model.EdgeCausal, Weight: 0.9}, + {SourceID: "n-a", TargetID: "n-d", EdgeType: model.EdgeEntity, Weight: 0.5}, + } { + e.Metadata, e.CreatedAt = map[string]string{}, now + if err := db.InsertEdge(e); err != nil { + t.Fatal(err) + } + } + edges, err := db.GetNeighborEdges("n-a") + if err != nil { + t.Fatal(err) + } + var got []string + for _, e := range edges { + got = append(got, e.SourceID+">"+e.TargetID) + } + want := []string{"n-c>n-a", "n-a>n-d", "n-a>n-b"} + if !reflect.DeepEqual(got, want) { + t.Fatalf("edges = %v, want %v", got, want) + } +} + +func TestRecallEdgeQueries_UseCoveringIndexes(t *testing.T) { + db := testDB(t) + plan := queryPlan(t, db, `SELECT source_id, target_id, edge_type, weight FROM edges WHERE source_id = ? + UNION ALL + SELECT source_id, target_id, edge_type, weight FROM edges WHERE target_id = ? AND source_id != ?`, "x", "x", "x") + if strings.Count(plan, "COVERING INDEX") != 2 { + t.Fatalf("neighbour query not answered from covering indexes:\n%s", plan) + } + plan = queryPlan(t, db, `SELECT target_id FROM edges INDEXED BY idx_edges_supersedes + WHERE edge_type = 'supersedes' AND source_id != target_id`) + if !strings.Contains(plan, "COVERING INDEX idx_edges_supersedes") { + t.Fatalf("superseded lookup is not answered from the partial index alone:\n%s", plan) + } +} + +func TestKnownEntities_MatchesLoadKnownEntities(t *testing.T) { + db := testDB(t) + a := makeInsight("ke-a", "first", 3) + a.Entities = []string{"Athena", "Hestia"} + b := makeInsight("ke-b", "second", 3) + b.Entities = []string{"Hestia", ""} + for _, ins := range []*model.Insight{a, b} { + if err := db.InsertInsight(ins); err != nil { + t.Fatal(err) + } + } + want, err := db.LoadKnownEntities() + if err != nil { + t.Fatal(err) + } + all, err := db.GetAllActiveInsights() + if err != nil { + t.Fatal(err) + } + if got := KnownEntities(all); !reflect.DeepEqual(got, want) { + t.Fatalf("KnownEntities = %v, want %v", got, want) + } +} + +// A read-only open skips migrations, so a store last written by an older +// build has no idx_edges_supersedes; the lookup must still answer. +func TestGetSupersededIDs_FallsBackWithoutPartialIndex(t *testing.T) { + db := testDB(t) + for _, id := range []string{"fb-new", "fb-old", "fb-self"} { + if err := db.InsertInsight(makeInsight(id, id, 3)); err != nil { + t.Fatal(err) + } + } + now := time.Now().UTC() + for _, e := range []*model.Edge{ + {SourceID: "fb-new", TargetID: "fb-old", EdgeType: model.EdgeSupersedes, Weight: 1}, + {SourceID: "fb-self", TargetID: "fb-self", EdgeType: model.EdgeSupersedes, Weight: 1}, + } { + if _, err := db.conn.Exec(`INSERT INTO edges (source_id, target_id, edge_type, weight, metadata, created_at) VALUES (?, ?, ?, ?, '{}', ?)`, + e.SourceID, e.TargetID, string(e.EdgeType), e.Weight, now.Format(time.RFC3339)); err != nil { + t.Fatal(err) + } + } + want := map[string]bool{"fb-old": true} + ids := []string{"fb-new", "fb-old", "fb-self"} + got, err := db.GetSupersededIDs(ids) + if err != nil || !reflect.DeepEqual(got, want) { + t.Fatalf("with index: %v, %v; want %v", got, err, want) + } + if _, err := db.conn.Exec(`DROP INDEX idx_edges_supersedes`); err != nil { + t.Fatal(err) + } + got, err = db.GetSupersededIDs(ids) + if err != nil || !reflect.DeepEqual(got, want) { + t.Fatalf("without index: %v, %v; want %v", got, err, want) + } +} From 06a1059bd7b1ebd16cc2e93126f9c723edc8d8ca Mon Sep 17 00:00:00 2001 From: GRIVN Date: Sat, 3 Oct 2026 19:08:12 +0800 Subject: [PATCH 3/5] fix(recall): rank transitions before applying the visit budget Order cached neighborhoods by the full intent and similarity contribution before consuming the visit budget. Relevant causal, temporal, and vector matches must not disappear behind higher raw-weight edges; preserve the existing traversal bounds. Validated focused search/store tests and real CLI comparisons: the causal result returns at rank two in all seven runs, while dense-graph recall retains the performance improvement. --- internal/memory/search/integration_test.go | 6 +- internal/memory/search/recall.go | 40 ++++--------- internal/memory/search/recall_budget_test.go | 62 ++++++++++++++++++++ internal/memory/search/recall_transition.go | 49 ++++++++++++++++ 4 files changed, 126 insertions(+), 31 deletions(-) create mode 100644 internal/memory/search/recall_budget_test.go create mode 100644 internal/memory/search/recall_transition.go diff --git a/internal/memory/search/integration_test.go b/internal/memory/search/integration_test.go index d24eed624..1930e702c 100644 --- a/internal/memory/search/integration_test.go +++ b/internal/memory/search/integration_test.go @@ -442,7 +442,7 @@ func TestBeamSearchFromAnchor_ScorePropagation(t *testing.T) { params := TraversalParams{BeamWidth: 10, MaxDepth: 3, MaxVisited: 100} scoreMap["bs-1"] = 1.0 - beamSearchFromAnchor(db, "bs-1", 1.0, nil, weights, params, scoreMap, viaMap, insightMap, nil, activeMap(t, db), map[string][]*model.Edge{}) + beamSearchFromAnchor(db, "bs-1", 1.0, nil, weights, params, scoreMap, viaMap, insightMap, nil, activeMap(t, db), map[string][]recallTransition{}) // Neighbor should be discovered with score > 0 if _, ok := scoreMap["bs-2"]; !ok { @@ -484,7 +484,7 @@ func TestBeamSearchFromAnchor_BeamWidthPruning(t *testing.T) { params := TraversalParams{BeamWidth: 3, MaxDepth: 3, MaxVisited: 500} scoreMap["bw-center"] = 1.0 - beamSearchFromAnchor(db, "bw-center", 1.0, nil, weights, params, scoreMap, viaMap, insightMap, nil, activeMap(t, db), map[string][]*model.Edge{}) + beamSearchFromAnchor(db, "bw-center", 1.0, nil, weights, params, scoreMap, viaMap, insightMap, nil, activeMap(t, db), map[string][]recallTransition{}) // Count deep nodes discovered — should be limited by beam width deepCount := 0 @@ -522,7 +522,7 @@ func TestBeamSearchFromAnchor_MaxVisitedBudget(t *testing.T) { params := TraversalParams{BeamWidth: 10, MaxDepth: 20, MaxVisited: 5} scoreMap["mv-0"] = 1.0 - beamSearchFromAnchor(db, "mv-0", 1.0, nil, weights, params, scoreMap, viaMap, insightMap, nil, activeMap(t, db), map[string][]*model.Edge{}) + beamSearchFromAnchor(db, "mv-0", 1.0, nil, weights, params, scoreMap, viaMap, insightMap, nil, activeMap(t, db), map[string][]recallTransition{}) // scoreMap includes the anchor itself, so discovered nodes (excluding anchor) should be <= 4 discovered := len(scoreMap) - 1 // subtract anchor diff --git a/internal/memory/search/recall.go b/internal/memory/search/recall.go index c0721f443..c75134e9f 100644 --- a/internal/memory/search/recall.go +++ b/internal/memory/search/recall.go @@ -270,7 +270,7 @@ func IntentAwareRecallFrom(db *store.DB, all []*model.Insight, query string, que // Step 3: Beam search from each anchor. Anchors' neighbourhoods overlap // heavily, so each node's edges are read from the store once per recall. - edgeCache := make(map[string][]*model.Edge) + edgeCache := make(map[string][]recallTransition) for id, a := range anchorMap { beamSearchFromAnchor(db, id, a.score, queryVec, weights, params, scoreMap, viaMap, insightMap, embedCache, activeByID, edgeCache) } @@ -522,8 +522,8 @@ func causalTopologicalSort(db *store.DB, results []RecallResult) []RecallResult // It uses a priority queue to keep the top beamWidth candidates at each depth level. // embedCache provides pre-loaded embedding vectors (nil = no embeddings). // activeByID resolves neighbours without a query; a neighbour missing from it -// is soft-deleted and never enters insightMap. edgeCache memoises edge reads -// across anchors within one recall. +// is soft-deleted and never enters insightMap. edgeCache memoises scored edge +// transitions across anchors within one recall. func beamSearchFromAnchor( db *store.DB, startID string, @@ -536,7 +536,7 @@ func beamSearchFromAnchor( insightMap map[string]*model.Insight, embedCache map[string][]float64, activeByID map[string]*model.Insight, - edgeCache map[string][]*model.Edge, + edgeCache map[string][]recallTransition, ) { visited := map[string]bool{startID: true} totalVisited := 1 @@ -563,43 +563,27 @@ func beamSearchFromAnchor( break } - edges, cached := edgeCache[cur.id] + transitions, cached := edgeCache[cur.id] if !cached { - var err error - edges, err = db.GetNeighborEdges(cur.id) + edges, err := db.GetNeighborEdges(cur.id) if err != nil { continue } - edgeCache[cur.id] = edges + transitions = rankRecallTransitions(cur.id, edges, queryVec, weights, embedCache) + edgeCache[cur.id] = transitions } - for _, e := range edges { + for _, transition := range transitions { if totalVisited >= params.MaxVisited { break } - neighborID := e.TargetID - if neighborID == cur.id { - neighborID = e.SourceID - } - - // MAGMA transition score (P6): additive accumulation - // score_v = score_u + λ₁·φ(edgeType, intent) + λ₂·sim(v_neighbor, v_query) - structural := weights[e.EdgeType] * e.Weight // φ(edgeType, intent) * edge_weight - semantic := 0.0 - if queryVec != nil && embedCache != nil { - if nVec, ok := embedCache[neighborID]; ok { - cosSim := embed.CosineSimilarity(queryVec, nVec) - if cosSim > 0 { - semantic = cosSim - } - } - } - neighborScore := cur.score + lambda1*structural + lambda2*semantic + neighborID := transition.id + neighborScore := cur.score + transition.delta // Update global score map if this path is better if existing, ok := scoreMap[neighborID]; !ok || neighborScore > existing { scoreMap[neighborID] = neighborScore - viaMap[neighborID] = string(e.EdgeType) + viaMap[neighborID] = string(transition.edgeType) if _, loaded := insightMap[neighborID]; !loaded { if ins, ok := activeByID[neighborID]; ok { insightMap[neighborID] = ins diff --git a/internal/memory/search/recall_budget_test.go b/internal/memory/search/recall_budget_test.go new file mode 100644 index 000000000..3d04f0c08 --- /dev/null +++ b/internal/memory/search/recall_budget_test.go @@ -0,0 +1,62 @@ +package search + +import ( + "fmt" + "testing" + "time" + + "github.com/mnemon-dev/mnemon/internal/memory/model" +) + +func TestRecallBudgetKeepsBestTransition(t *testing.T) { + for _, tc := range []struct { + name string + intent Intent + kind model.EdgeType + weight float64 + incoming bool + vector bool + }{ + {"causal outgoing", IntentWhy, model.EdgeCausal, 0.9, false, false}, + {"causal incoming", IntentWhy, model.EdgeCausal, 0.9, true, false}, + {"temporal", IntentWhen, model.EdgeTemporal, 0.9, false, false}, + {"semantic similarity", IntentGeneral, model.EdgeEntity, 0.1, false, true}, + } { + t.Run(tc.name, func(t *testing.T) { + db := testDB(t) + now := time.Now().UTC() + active := make(map[string]*model.Insight) + for _, id := range []string{"root", "gold", "noise1", "noise2", "noise3"} { + active[id] = insertInsight(t, db, id, id, "test", 3, nil, now) + } + for i := 1; i <= 3; i++ { + if err := db.InsertEdge(&model.Edge{SourceID: "root", TargetID: fmt.Sprintf("noise%d", i), EdgeType: model.EdgeEntity, Weight: 1, CreatedAt: now}); err != nil { + t.Fatal(err) + } + } + source, target := "root", "gold" + if tc.incoming { + source, target = target, source + } + if err := db.InsertEdge(&model.Edge{SourceID: source, TargetID: target, EdgeType: tc.kind, Weight: tc.weight, CreatedAt: now}); err != nil { + t.Fatal(err) + } + var query []float64 + var embeddings map[string][]float64 + if tc.vector { + query = []float64{1, 0} + embeddings = map[string][]float64{"gold": {1, 0}} + } + scores := map[string]float64{"root": 1} + via := map[string]string{} + insights := map[string]*model.Insight{"root": active["root"]} + beamSearchFromAnchor(db, "root", 1, query, GetWeights(tc.intent), TraversalParams{BeamWidth: 2, MaxDepth: 1, MaxVisited: 3}, scores, via, insights, embeddings, active, make(map[string][]recallTransition)) + if _, ok := scores["gold"]; !ok { + t.Fatalf("best transition excluded by visit budget: scores=%v", scores) + } + if len(scores) > 3 { + t.Fatalf("visit bound exceeded: %d", len(scores)) + } + }) + } +} diff --git a/internal/memory/search/recall_transition.go b/internal/memory/search/recall_transition.go new file mode 100644 index 000000000..53346ae28 --- /dev/null +++ b/internal/memory/search/recall_transition.go @@ -0,0 +1,49 @@ +package search + +import ( + "sort" + + "github.com/mnemon-dev/mnemon/internal/memory/embed" + "github.com/mnemon-dev/mnemon/internal/memory/model" +) + +type recallTransition struct { + id string + edgeType model.EdgeType + delta float64 +} + +// rankRecallTransitions applies the same intent and similarity terms as beam +// scoring before the visit budget can discard a neighbour. The ordering and +// deltas are reusable across anchors within this one recall, whose query and +// intent do not change. The current path score is added by the caller. +func rankRecallTransitions(nodeID string, edges []*model.Edge, query []float64, + weights IntentWeights, embeddings map[string][]float64) []recallTransition { + transitions := make([]recallTransition, 0, len(edges)) + for _, e := range edges { + id := e.TargetID + if id == nodeID { + id = e.SourceID + } + delta := lambda1 * weights[e.EdgeType] * e.Weight + if query != nil { + if vector, ok := embeddings[id]; ok { + if similarity := embed.CosineSimilarity(query, vector); similarity > 0 { + delta += lambda2 * similarity + } + } + } + transitions = append(transitions, recallTransition{id: id, edgeType: e.EdgeType, delta: delta}) + } + sort.Slice(transitions, func(i, j int) bool { + a, b := transitions[i], transitions[j] + if a.delta != b.delta { + return a.delta > b.delta + } + if a.id != b.id { + return a.id < b.id + } + return a.edgeType < b.edgeType + }) + return transitions +} From 6e6308bb142c468409a505bab67ca6b4f06e40bd Mon Sep 17 00:00:00 2001 From: GRIVN Date: Sat, 3 Oct 2026 19:08:12 +0800 Subject: [PATCH 4/5] fix(store): scope supersedes lookups to recall candidates Seek requested target IDs through the covering partial index in bounded batches instead of scanning unrelated supersession history. Retain the read-only legacy-index fallback and exclude self-supersession. Validated query-plan, chunking, and fallback checks plus real CLI comparisons with 100,000 historical edges: median query time returns from 29.09 ms to 12.69 ms with unchanged results. --- internal/memory/store/edge.go | 69 ++++++---------------- internal/memory/store/recall_index_test.go | 8 +-- 2 files changed, 21 insertions(+), 56 deletions(-) diff --git a/internal/memory/store/edge.go b/internal/memory/store/edge.go index a7055e77a..c03bcd44c 100644 --- a/internal/memory/store/edge.go +++ b/internal/memory/store/edge.go @@ -41,9 +41,8 @@ func (db *DB) GetEdgesByNode(nodeID string) ([]*model.Edge, error) { // GetNeighborEdges returns every edge touching nodeID with the columns beam // search scores on (metadata and created_at are left zero). Both halves are -// answered from the covering indexes alone, and the fixed order makes recall -// deterministic: beam search stops at a visit budget, so which neighbours it -// reaches depends on the order edges arrive in. +// answered from covering indexes in a fixed store order. Recall ranks the full +// intent and similarity transition score before applying its visit budget. func (db *DB) GetNeighborEdges(nodeID string) ([]*model.Edge, error) { rows, err := db.execer().Query( `SELECT source_id, target_id, edge_type, weight FROM edges WHERE source_id = ? @@ -68,48 +67,12 @@ func (db *DB) GetNeighborEdges(nodeID string) ([]*model.Edge, error) { return results, rows.Err() } -// GetSupersededIDs returns which of the given ids are the target of at least -// one 'supersedes' edge from a different insight, i.e. which of them some -// other insight claims to replace. Recall uses this to demote stale content; -// the rows are kept so the lineage stays inspectable. -// -// It reads every supersedes edge from the partial idx_edges_supersedes (a few -// thousand rows, a few hundred KB) and filters to the caller's ids in memory: -// probing each candidate through the whole-graph target index costs a random -// read per candidate on a cold cache. SQLite's planner does not pick a partial -// index by itself, hence INDEXED BY. A read-only open skips migrations, so a -// store no writable open has touched since the index was introduced lacks it; -// that store falls back to the scoped lookup. +// GetSupersededIDs checks only the requested targets, using the covering partial +// index when available. A store can accumulate an unbounded supersession history; +// unrelated history must not be scanned for a small recall candidate set. +// Read-only legacy stores without the partial index retain the scoped fallback. func (db *DB) GetSupersededIDs(ids []string) (map[string]bool, error) { - superseded := make(map[string]bool) - if len(ids) == 0 { - return superseded, nil - } - want := make(map[string]bool, len(ids)) - for _, id := range ids { - want[id] = true - } - ex := db.execer() - rows, err := ex.Query( - `SELECT target_id FROM edges INDEXED BY idx_edges_supersedes - WHERE edge_type = 'supersedes' AND source_id != target_id`) - if err != nil && strings.Contains(err.Error(), "no such index") { - return scopedSupersededIDs(ex, ids) - } - if err != nil { - return nil, err - } - defer rows.Close() - for rows.Next() { - var id string - if err := rows.Scan(&id); err != nil { - return nil, err - } - if want[id] { - superseded[id] = true - } - } - return superseded, rows.Err() + return scopedSupersededIDs(db.execer(), ids) } // supersededLookupChunk bounds how many ids go into one IN clause. SQLite's @@ -118,8 +81,10 @@ func (db *DB) GetSupersededIDs(ids []string) (map[string]bool, error) { // the loop below runs once. const supersededLookupChunk = 500 -// scopedSupersededIDs answers GetSupersededIDs through idx_edges_target_type -// for stores that predate idx_edges_supersedes. +const supersededLookupSQL = `SELECT DISTINCT target_id FROM edges INDEXED BY idx_edges_supersedes + WHERE edge_type = 'supersedes' AND source_id != target_id AND target_id IN (%s)` + +// scopedSupersededIDs bounds each query to the SQLite parameter limit. func scopedSupersededIDs(ex dbExecer, ids []string) (map[string]bool, error) { superseded := make(map[string]bool) for start := 0; start < len(ids); start += supersededLookupChunk { @@ -135,18 +100,18 @@ func scopedSupersededIDs(ex dbExecer, ids []string) (map[string]bool, error) { // are closed before returning: the pool holds a single connection, so an open // cursor would block the next batch. func collectSupersededIDs(ex dbExecer, chunk []string, into map[string]bool) error { - args := make([]any, 0, len(chunk)+1) - args = append(args, string(model.EdgeSupersedes)) + args := make([]any, 0, len(chunk)) placeholders := make([]string, len(chunk)) for i, id := range chunk { placeholders[i] = "?" args = append(args, id) } - rows, err := ex.Query(fmt.Sprintf( - `SELECT DISTINCT target_id FROM edges - WHERE edge_type = ? AND source_id != target_id AND target_id IN (%s)`, - strings.Join(placeholders, ",")), args...) + query := fmt.Sprintf(supersededLookupSQL, strings.Join(placeholders, ",")) + rows, err := ex.Query(query, args...) + if err != nil && strings.Contains(err.Error(), "no such index: idx_edges_supersedes") { + rows, err = ex.Query(strings.Replace(query, " INDEXED BY idx_edges_supersedes", "", 1), args...) + } if err != nil { return err } diff --git a/internal/memory/store/recall_index_test.go b/internal/memory/store/recall_index_test.go index 253450271..ac06d9430 100644 --- a/internal/memory/store/recall_index_test.go +++ b/internal/memory/store/recall_index_test.go @@ -1,6 +1,7 @@ package store import ( + "fmt" "reflect" "strings" "testing" @@ -68,10 +69,9 @@ func TestRecallEdgeQueries_UseCoveringIndexes(t *testing.T) { if strings.Count(plan, "COVERING INDEX") != 2 { t.Fatalf("neighbour query not answered from covering indexes:\n%s", plan) } - plan = queryPlan(t, db, `SELECT target_id FROM edges INDEXED BY idx_edges_supersedes - WHERE edge_type = 'supersedes' AND source_id != target_id`) - if !strings.Contains(plan, "COVERING INDEX idx_edges_supersedes") { - t.Fatalf("superseded lookup is not answered from the partial index alone:\n%s", plan) + plan = queryPlan(t, db, fmt.Sprintf(supersededLookupSQL, "?,?"), "first", "second") + if !strings.Contains(plan, "SEARCH edges USING COVERING INDEX idx_edges_supersedes (target_id=?)") || strings.Contains(plan, "SCAN edges") { + t.Fatalf("superseded lookup must seek requested targets in the partial index:\n%s", plan) } } From 431b81b19b7f73fc240bf6b8387baedf198ae2ec Mon Sep 17 00:00:00 2001 From: GRIVN Date: Sat, 3 Oct 2026 19:08:12 +0800 Subject: [PATCH 5/5] fix(store): reject busy checkpoints during compaction Inspect the checkpoint result rows before and after VACUUM, report blocked or partially completed compaction accurately, and propagate database size errors instead of returning a false success. Validated real concurrent-reader CLI reproduction and SQLite tests for busy rejection, successful retry with committed data preserved, and read-only rejection. Normal compaction preserves content and passes integrity_check. --- internal/memory/store/compact.go | 31 ++++++++--- internal/memory/store/compact_busy_test.go | 62 ++++++++++++++++++++++ 2 files changed, 85 insertions(+), 8 deletions(-) create mode 100644 internal/memory/store/compact_busy_test.go diff --git a/internal/memory/store/compact.go b/internal/memory/store/compact.go index 4d7b674ba..07e30fa63 100644 --- a/internal/memory/store/compact.go +++ b/internal/memory/store/compact.go @@ -14,23 +14,38 @@ func (db *DB) Compact() (before, after int64, err error) { if db.readOnly { return 0, 0, fmt.Errorf("compact: database is read-only") } - if _, err := db.conn.Exec(`PRAGMA wal_checkpoint(TRUNCATE)`); err != nil { + if err := db.compactCheckpoint(); err != nil { return 0, 0, fmt.Errorf("checkpoint: %w", err) } - before = db.fileSize() + before, err = db.fileSize() + if err != nil { + return 0, 0, err + } if _, err := db.conn.Exec(`VACUUM`); err != nil { return 0, 0, fmt.Errorf("vacuum: %w", err) } - if _, err := db.conn.Exec(`PRAGMA wal_checkpoint(TRUNCATE)`); err != nil { - return 0, 0, fmt.Errorf("checkpoint: %w", err) + if err := db.compactCheckpoint(); err != nil { + return before, 0, fmt.Errorf("vacuum completed but checkpoint failed: %w", err) + } + after, err = db.fileSize() + return before, after, err +} + +func (db *DB) compactCheckpoint() error { + var busy, log, checkpointed int + if err := db.conn.QueryRow(`PRAGMA wal_checkpoint(TRUNCATE)`).Scan(&busy, &log, &checkpointed); err != nil { + return err + } + if busy != 0 { + return fmt.Errorf("database busy: checkpointed %d of %d WAL frames; retry after active transactions finish", checkpointed, log) } - return before, db.fileSize(), nil + return nil } -func (db *DB) fileSize() int64 { +func (db *DB) fileSize() (int64, error) { fi, err := os.Stat(db.path) if err != nil { - return 0 + return 0, fmt.Errorf("stat database: %w", err) } - return fi.Size() + return fi.Size(), nil } diff --git a/internal/memory/store/compact_busy_test.go b/internal/memory/store/compact_busy_test.go new file mode 100644 index 000000000..1ddd20442 --- /dev/null +++ b/internal/memory/store/compact_busy_test.go @@ -0,0 +1,62 @@ +package store + +import ( + "database/sql" + "strings" + "testing" +) + +func TestCompactRejectsBusyCheckpoint(t *testing.T) { + db := testDB(t) + if _, err := db.conn.Exec(`CREATE TABLE compact_probe (value INTEGER); INSERT INTO compact_probe VALUES (1); PRAGMA wal_checkpoint(TRUNCATE)`); err != nil { + t.Fatal(err) + } + reader, err := sql.Open("sqlite", db.path) + if err != nil { + t.Fatal(err) + } + defer reader.Close() + tx, err := reader.Begin() + if err != nil { + t.Fatal(err) + } + defer tx.Rollback() + var value int + if err := tx.QueryRow(`SELECT value FROM compact_probe`).Scan(&value); err != nil { + t.Fatal(err) + } + if _, err := db.conn.Exec(`UPDATE compact_probe SET value=2`); err != nil { + t.Fatal(err) + } + if _, _, err := db.Compact(); err == nil || !strings.Contains(err.Error(), "busy") { + t.Fatalf("Compact error=%v, want busy checkpoint", err) + } + if err := tx.Rollback(); err != nil { + t.Fatal(err) + } + if _, _, err := db.Compact(); err != nil { + t.Fatalf("retry after reader closes: %v", err) + } + if err := db.conn.QueryRow(`SELECT value FROM compact_probe`).Scan(&value); err != nil || value != 2 { + t.Fatalf("committed value changed: value=%d err=%v", value, err) + } +} + +func TestCompactRejectsReadOnly(t *testing.T) { + dir := t.TempDir() + db, err := Open(dir) + if err != nil { + t.Fatal(err) + } + if err := db.Close(); err != nil { + t.Fatal(err) + } + ro, err := OpenReadOnly(dir) + if err != nil { + t.Fatal(err) + } + defer ro.Close() + if _, _, err := ro.Compact(); err == nil { + t.Fatal("readonly compaction succeeded") + } +}