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/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/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/search/integration_test.go b/internal/memory/search/integration_test.go index 9453db33d..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) + 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) + 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) + 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 @@ -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..c75134e9f 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][]recallTransition) 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 scored edge +// transitions 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][]recallTransition, ) { visited := map[string]bool{startID: true} totalVisited := 1 @@ -545,41 +563,29 @@ func beamSearchFromAnchor( break } - edges, err := db.GetEdgesByNode(cur.id) - if err != nil { - continue + transitions, cached := edgeCache[cur.id] + if !cached { + edges, err := db.GetNeighborEdges(cur.id) + if err != nil { + continue + } + 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 { - ins, err := db.GetInsightByID(neighborID) - if err == nil && ins != nil { + 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 +} diff --git a/internal/memory/store/compact.go b/internal/memory/store/compact.go new file mode 100644 index 000000000..07e30fa63 --- /dev/null +++ b/internal/memory/store/compact.go @@ -0,0 +1,51 @@ +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.compactCheckpoint(); err != nil { + return 0, 0, fmt.Errorf("checkpoint: %w", err) + } + 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.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 nil +} + +func (db *DB) fileSize() (int64, error) { + fi, err := os.Stat(db.path) + if err != nil { + return 0, fmt.Errorf("stat database: %w", err) + } + 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") + } +} 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) + } +} 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..c03bcd44c 100644 --- a/internal/memory/store/edge.go +++ b/internal/memory/store/edge.go @@ -39,24 +39,54 @@ 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 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 = ? + 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 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) { + return scopedSupersededIDs(db.execer(), ids) +} + // 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) { +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) - 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 { @@ -70,18 +100,18 @@ func (db *DB) GetSupersededIDs(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/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..ac06d9430 --- /dev/null +++ b/internal/memory/store/recall_index_test.go @@ -0,0 +1,134 @@ +package store + +import ( + "fmt" + "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, 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) + } +} + +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) + } +}