From 67b707ef3ff8bd4689922176efd55f6aea9a048a Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 6 Apr 2026 10:02:30 +0000 Subject: [PATCH 1/8] Initial plan From 04014269ad545a53285e9255b6d2271ed63743a5 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 6 Apr 2026 10:10:23 +0000 Subject: [PATCH 2/8] Add batch entity and relation hydration lookups Agent-Logs-Url: https://github.com/vndee/memex/sessions/1ec419ab-55f2-4063-8471-c0e8891d649c Co-authored-by: vndee <28271488+vndee@users.noreply.github.com> --- internal/search/graph.go | 18 +++--- internal/search/graph_test.go | 76 ++++++++++++++++++++++ internal/server/helpers.go | 27 ++++++-- internal/server/helpers_test.go | 108 +++++++++++++++++++++++++++++++ internal/storage/entities.go | 45 +++++++++++++ internal/storage/relations.go | 44 +++++++++++++ internal/storage/sqlite.go | 2 + internal/storage/storage_test.go | 102 +++++++++++++++++++++++++++++ 8 files changed, 409 insertions(+), 13 deletions(-) create mode 100644 internal/search/graph_test.go create mode 100644 internal/server/helpers_test.go diff --git a/internal/search/graph.go b/internal/search/graph.go index 9549061..81e60ba 100644 --- a/internal/search/graph.go +++ b/internal/search/graph.go @@ -2,8 +2,6 @@ package search import ( "context" - "database/sql" - "errors" "fmt" "log/slog" "sort" @@ -184,18 +182,20 @@ func hydrateEntityResults( scores map[string]float64, limit int, ) ([]*domain.SearchResult, error) { + entitiesByID, err := store.GetEntitiesByIDs(ctx, kbID, ids) + if err != nil { + slog.Warn("graph entity batch lookup failed", "count", len(ids), "error", err) + return nil, fmt.Errorf("get entities by ids: %w", err) + } + results := make([]*domain.SearchResult, 0, min(len(ids), limit)) for _, id := range ids { if len(results) >= limit { break } - ent, err := store.GetEntity(ctx, kbID, id) - if err != nil { - if errors.Is(err, sql.ErrNoRows) { - continue - } - slog.Warn("graph entity lookup failed", "id", id, "error", err) - return nil, fmt.Errorf("get entity %s: %w", id, err) + ent, ok := entitiesByID[id] + if !ok { + continue } results = append(results, &domain.SearchResult{ ID: ent.ID, diff --git a/internal/search/graph_test.go b/internal/search/graph_test.go new file mode 100644 index 0000000..b6d4672 --- /dev/null +++ b/internal/search/graph_test.go @@ -0,0 +1,76 @@ +package search + +import ( + "context" + "testing" + "time" + + memex "github.com/vndee/memex" + "github.com/vndee/memex/internal/domain" + "github.com/vndee/memex/internal/storage" +) + +func init() { + storage.MigrationSQL = memex.MigrationSQL() +} + +type countingSearchStore struct { + storage.Store + getEntityCalls int + getEntitiesByIDsCalls int +} + +func (s *countingSearchStore) GetEntity(ctx context.Context, kbID, id string) (*domain.Entity, error) { + s.getEntityCalls++ + return s.Store.GetEntity(ctx, kbID, id) +} + +func (s *countingSearchStore) GetEntitiesByIDs(ctx context.Context, kbID string, ids []string) (map[string]*domain.Entity, error) { + s.getEntitiesByIDsCalls++ + return s.Store.GetEntitiesByIDs(ctx, kbID, ids) +} + +func TestHydrateEntityResults_UsesBatchLookup(t *testing.T) { + sqliteStore, err := storage.NewSQLiteStore(":memory:") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = sqliteStore.Close() }) + + now := time.Now().UTC() + if err := sqliteStore.CreateKB(context.Background(), &domain.KnowledgeBase{ + ID: "kb1", Name: "KB 1", + EmbedConfig: domain.EmbedConfig{Provider: "ollama", Model: "nomic-embed-text"}, + LLMConfig: domain.LLMConfig{Provider: "ollama", Model: "llama3.2"}, + CreatedAt: now, + }); err != nil { + t.Fatal(err) + } + for _, e := range []*domain.Entity{ + {ID: "e1", KBID: "kb1", Name: "Alice", Type: "person", Summary: "Engineer", CreatedAt: now, UpdatedAt: now}, + {ID: "e2", KBID: "kb1", Name: "Bob", Type: "person", Summary: "Manager", CreatedAt: now, UpdatedAt: now}, + } { + if err := sqliteStore.CreateEntity(context.Background(), e); err != nil { + t.Fatal(err) + } + } + + store := &countingSearchStore{Store: sqliteStore} + ids := []string{"e1", "missing", "e2"} + scores := map[string]float64{"e1": 1, "e2": 0.5, "missing": 0.1} + + results, err := hydrateEntityResults(context.Background(), store, "kb1", ids, scores, 10) + if err != nil { + t.Fatal(err) + } + if len(results) != 2 { + t.Fatalf("got %d results, want 2", len(results)) + } + if store.getEntitiesByIDsCalls != 1 { + t.Fatalf("GetEntitiesByIDs calls = %d, want 1", store.getEntitiesByIDsCalls) + } + if store.getEntityCalls != 0 { + t.Fatalf("GetEntity calls = %d, want 0", store.getEntityCalls) + } +} + diff --git a/internal/server/helpers.go b/internal/server/helpers.go index da8796f..47b2b27 100644 --- a/internal/server/helpers.go +++ b/internal/server/helpers.go @@ -2,6 +2,7 @@ package server import ( "context" + "fmt" "time" "github.com/vndee/memex/internal/domain" @@ -47,10 +48,19 @@ func buildKB(id, name, desc, embedProvider, embedModel, llmProvider, llmModel st // HydrateSubgraph enriches a raw SubgraphResult with entity and relation metadata // from storage. Used by MCP, HTTP, and CLI graph traversal handlers. func HydrateSubgraph(ctx context.Context, store storage.Store, kbID string, sg graph.SubgraphResult) (*domain.Subgraph, error) { + nodeIDs := make([]string, 0, len(sg.Nodes)) + for id := range sg.Nodes { + nodeIDs = append(nodeIDs, id) + } + entitiesByID, err := store.GetEntitiesByIDs(ctx, kbID, nodeIDs) + if err != nil { + return nil, fmt.Errorf("get subgraph entities by ids: %w", err) + } + nodes := make([]domain.SubgraphNode, 0, len(sg.Nodes)) for id, dist := range sg.Nodes { - ent, err := store.GetEntity(ctx, kbID, id) - if err != nil { + ent, ok := entitiesByID[id] + if !ok { nodes = append(nodes, domain.SubgraphNode{ ID: id, Distance: dist, }) @@ -65,10 +75,19 @@ func HydrateSubgraph(ctx context.Context, store storage.Store, kbID string, sg g }) } + edgeIDs := make([]string, 0, len(sg.Edges)) + for _, e := range sg.Edges { + edgeIDs = append(edgeIDs, e.RelID) + } + relationsByID, err := store.GetRelationsByIDs(ctx, kbID, edgeIDs) + if err != nil { + return nil, fmt.Errorf("get subgraph relations by ids: %w", err) + } + edges := make([]domain.SubgraphEdge, 0, len(sg.Edges)) for _, e := range sg.Edges { - rel, err := store.GetRelation(ctx, kbID, e.RelID) - if err != nil { + rel, ok := relationsByID[e.RelID] + if !ok { edges = append(edges, domain.SubgraphEdge{ ID: e.RelID, SourceID: e.SourceID, diff --git a/internal/server/helpers_test.go b/internal/server/helpers_test.go new file mode 100644 index 0000000..a191c7c --- /dev/null +++ b/internal/server/helpers_test.go @@ -0,0 +1,108 @@ +package server + +import ( + "context" + "testing" + "time" + + memex "github.com/vndee/memex" + "github.com/vndee/memex/internal/domain" + "github.com/vndee/memex/internal/graph" + "github.com/vndee/memex/internal/storage" +) + +func init() { + storage.MigrationSQL = memex.MigrationSQL() +} + +type countingServerStore struct { + storage.Store + getEntityCalls int + getEntitiesByIDsCalls int + getRelationCalls int + getRelationsByIDsCalls int +} + +func (s *countingServerStore) GetEntity(ctx context.Context, kbID, id string) (*domain.Entity, error) { + s.getEntityCalls++ + return s.Store.GetEntity(ctx, kbID, id) +} + +func (s *countingServerStore) GetEntitiesByIDs(ctx context.Context, kbID string, ids []string) (map[string]*domain.Entity, error) { + s.getEntitiesByIDsCalls++ + return s.Store.GetEntitiesByIDs(ctx, kbID, ids) +} + +func (s *countingServerStore) GetRelation(ctx context.Context, kbID, id string) (*domain.Relation, error) { + s.getRelationCalls++ + return s.Store.GetRelation(ctx, kbID, id) +} + +func (s *countingServerStore) GetRelationsByIDs(ctx context.Context, kbID string, ids []string) (map[string]*domain.Relation, error) { + s.getRelationsByIDsCalls++ + return s.Store.GetRelationsByIDs(ctx, kbID, ids) +} + +func TestHydrateSubgraph_UsesBatchLookups(t *testing.T) { + sqliteStore, err := storage.NewSQLiteStore(":memory:") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = sqliteStore.Close() }) + + now := time.Now().UTC() + if err := sqliteStore.CreateKB(context.Background(), &domain.KnowledgeBase{ + ID: "kb1", Name: "KB 1", + EmbedConfig: domain.EmbedConfig{Provider: "ollama", Model: "nomic-embed-text"}, + LLMConfig: domain.LLMConfig{Provider: "ollama", Model: "llama3.2"}, + CreatedAt: now, + }); err != nil { + t.Fatal(err) + } + for _, e := range []*domain.Entity{ + {ID: "e1", KBID: "kb1", Name: "Alice", Type: "person", Summary: "Engineer", CreatedAt: now, UpdatedAt: now}, + {ID: "e2", KBID: "kb1", Name: "Bob", Type: "person", Summary: "Manager", CreatedAt: now, UpdatedAt: now}, + } { + if err := sqliteStore.CreateEntity(context.Background(), e); err != nil { + t.Fatal(err) + } + } + if err := sqliteStore.CreateRelation(context.Background(), &domain.Relation{ + ID: "r1", KBID: "kb1", SourceID: "e1", TargetID: "e2", Type: "knows", Weight: 1.0, ValidAt: now, CreatedAt: now, + }); err != nil { + t.Fatal(err) + } + + store := &countingServerStore{Store: sqliteStore} + subgraph := graph.SubgraphResult{ + Nodes: map[string]int{"e1": 0, "e2": 1, "missing": 2}, + Edges: []graph.SubgraphEdge{ + {RelID: "r1", SourceID: "e1", TargetID: "e2", Type: "knows", Weight: 1}, + {RelID: "missing-rel", SourceID: "e2", TargetID: "missing", Type: "mentions", Weight: 0.5}, + }, + } + + hydrated, err := HydrateSubgraph(context.Background(), store, "kb1", subgraph) + if err != nil { + t.Fatal(err) + } + if len(hydrated.Nodes) != 3 { + t.Fatalf("got %d nodes, want 3", len(hydrated.Nodes)) + } + if len(hydrated.Edges) != 2 { + t.Fatalf("got %d edges, want 2", len(hydrated.Edges)) + } + if store.getEntitiesByIDsCalls != 1 { + t.Fatalf("GetEntitiesByIDs calls = %d, want 1", store.getEntitiesByIDsCalls) + } + if store.getRelationsByIDsCalls != 1 { + t.Fatalf("GetRelationsByIDs calls = %d, want 1", store.getRelationsByIDsCalls) + } + if store.getEntityCalls != 0 { + t.Fatalf("GetEntity calls = %d, want 0", store.getEntityCalls) + } + if store.getRelationCalls != 0 { + t.Fatalf("GetRelation calls = %d, want 0", store.getRelationCalls) + } +} + diff --git a/internal/storage/entities.go b/internal/storage/entities.go index 8bea152..b150f99 100644 --- a/internal/storage/entities.go +++ b/internal/storage/entities.go @@ -6,6 +6,7 @@ import ( "encoding/binary" "fmt" "math" + "strings" "time" "github.com/vndee/memex/internal/domain" @@ -33,6 +34,50 @@ func (s *SQLiteStore) GetEntity(ctx context.Context, kbID, id string) (*domain.E return scanEntity(row) } +func (s *SQLiteStore) GetEntitiesByIDs(ctx context.Context, kbID string, ids []string) (map[string]*domain.Entity, error) { + entities := make(map[string]*domain.Entity) + if len(ids) == 0 { + return entities, nil + } + + const maxBatchIDs = 900 + for start := 0; start < len(ids); start += maxBatchIDs { + end := min(start+maxBatchIDs, len(ids)) + batch := ids[start:end] + + placeholders := strings.TrimSuffix(strings.Repeat("?,", len(batch)), ",") + query := `SELECT id, kb_id, name, type, summary, embedding, created_at, updated_at + FROM entities WHERE kb_id = ? AND id IN (` + placeholders + `)` + + args := make([]any, 0, len(batch)+1) + args = append(args, kbID) + for _, id := range batch { + args = append(args, id) + } + + rows, err := s.db.QueryContext(ctx, query, args...) + if err != nil { + return nil, fmt.Errorf("query entities by ids: %w", err) + } + for rows.Next() { + e, err := scanEntity(rows) + if err != nil { + rows.Close() + return nil, err + } + entities[e.ID] = e + } + if err := rows.Err(); err != nil { + rows.Close() + return nil, fmt.Errorf("iterate entities by ids: %w", err) + } + if err := rows.Close(); err != nil { + return nil, fmt.Errorf("close entities by ids rows: %w", err) + } + } + return entities, nil +} + func (s *SQLiteStore) UpdateEntity(ctx context.Context, e *domain.Entity) error { res, err := s.db.ExecContext(ctx, `UPDATE entities SET name = ?, type = ?, summary = ?, embedding = ?, updated_at = ? diff --git a/internal/storage/relations.go b/internal/storage/relations.go index edd0715..22db139 100644 --- a/internal/storage/relations.go +++ b/internal/storage/relations.go @@ -74,6 +74,50 @@ func (s *SQLiteStore) GetRelation(ctx context.Context, kbID, id string) (*domain return scanRelation(row) } +func (s *SQLiteStore) GetRelationsByIDs(ctx context.Context, kbID string, ids []string) (map[string]*domain.Relation, error) { + rels := make(map[string]*domain.Relation) + if len(ids) == 0 { + return rels, nil + } + + const maxBatchIDs = 900 + for start := 0; start < len(ids); start += maxBatchIDs { + end := min(start+maxBatchIDs, len(ids)) + batch := ids[start:end] + + placeholders := strings.TrimSuffix(strings.Repeat("?,", len(batch)), ",") + query := `SELECT id, kb_id, source_id, target_id, type, summary, weight, embedding, episode_id, valid_at, invalid_at, created_at + FROM relations WHERE kb_id = ? AND id IN (` + placeholders + `)` + + args := make([]any, 0, len(batch)+1) + args = append(args, kbID) + for _, id := range batch { + args = append(args, id) + } + + rows, err := s.db.QueryContext(ctx, query, args...) + if err != nil { + return nil, fmt.Errorf("query relations by ids: %w", err) + } + for rows.Next() { + r, err := scanRelation(rows) + if err != nil { + rows.Close() + return nil, err + } + rels[r.ID] = r + } + if err := rows.Err(); err != nil { + rows.Close() + return nil, fmt.Errorf("iterate relations by ids: %w", err) + } + if err := rows.Close(); err != nil { + return nil, fmt.Errorf("close relations by ids rows: %w", err) + } + } + return rels, nil +} + func (s *SQLiteStore) InvalidateRelation(ctx context.Context, kbID, id string, invalidAt time.Time) error { res, err := s.db.ExecContext(ctx, `UPDATE relations SET invalid_at = ? WHERE kb_id = ? AND id = ? AND invalid_at IS NULL`, diff --git a/internal/storage/sqlite.go b/internal/storage/sqlite.go index 7ceaf1a..febc11b 100644 --- a/internal/storage/sqlite.go +++ b/internal/storage/sqlite.go @@ -34,6 +34,7 @@ type Store interface { // Entities CreateEntity(ctx context.Context, e *domain.Entity) error GetEntity(ctx context.Context, kbID, id string) (*domain.Entity, error) + GetEntitiesByIDs(ctx context.Context, kbID string, ids []string) (map[string]*domain.Entity, error) UpdateEntity(ctx context.Context, e *domain.Entity) error DeleteEntity(ctx context.Context, kbID, id string) error FindEntitiesByName(ctx context.Context, kbID, name string) ([]*domain.Entity, error) @@ -43,6 +44,7 @@ type Store interface { // Relations CreateRelation(ctx context.Context, r *domain.Relation) error GetRelation(ctx context.Context, kbID, id string) (*domain.Relation, error) + GetRelationsByIDs(ctx context.Context, kbID string, ids []string) (map[string]*domain.Relation, error) InvalidateRelation(ctx context.Context, kbID, id string, invalidAt time.Time) error GetRelationsForEntity(ctx context.Context, kbID, entityID string) ([]*domain.Relation, error) GetValidRelations(ctx context.Context, kbID string, at time.Time) ([]*domain.Relation, error) diff --git a/internal/storage/storage_test.go b/internal/storage/storage_test.go index 8ae1bc4..1daaff4 100644 --- a/internal/storage/storage_test.go +++ b/internal/storage/storage_test.go @@ -3,6 +3,7 @@ package storage_test import ( "context" "database/sql" + "fmt" "testing" "time" @@ -292,6 +293,66 @@ func TestEntityScopeEnforced(t *testing.T) { } } +func TestGetEntitiesByIDs(t *testing.T) { + store := newTestStore(t) + createTestKB(t, store, "kb1") + createTestKB(t, store, "kb2") + + now := time.Now().UTC() + for _, e := range []*domain.Entity{ + {ID: "ent-001", KBID: "kb1", Name: "Alice", Type: "person", CreatedAt: now, UpdatedAt: now}, + {ID: "ent-002", KBID: "kb1", Name: "Bob", Type: "person", CreatedAt: now, UpdatedAt: now}, + {ID: "ent-003", KBID: "kb2", Name: "Charlie", Type: "person", CreatedAt: now, UpdatedAt: now}, + } { + if err := store.CreateEntity(context.Background(), e); err != nil { + t.Fatal("create entity:", err) + } + } + + entities, err := store.GetEntitiesByIDs(context.Background(), "kb1", []string{"ent-001", "ent-002", "ent-003", "ent-missing"}) + if err != nil { + t.Fatal("get entities by ids:", err) + } + if len(entities) != 2 { + t.Fatalf("got %d entities, want 2", len(entities)) + } + if entities["ent-001"] == nil || entities["ent-001"].Name != "Alice" { + t.Fatal("expected ent-001 (Alice) in result") + } + if entities["ent-002"] == nil || entities["ent-002"].Name != "Bob" { + t.Fatal("expected ent-002 (Bob) in result") + } + if _, ok := entities["ent-003"]; ok { + t.Fatal("did not expect kb2 entity in kb1 result") + } +} + +func TestGetEntitiesByIDs_LargeBatch(t *testing.T) { + store := newTestStore(t) + createTestKB(t, store, "kb1") + now := time.Now().UTC() + + const n = 905 + ids := make([]string, 0, n) + for i := 0; i < n; i++ { + id := fmt.Sprintf("ent-%04d", i) + ids = append(ids, id) + if err := store.CreateEntity(context.Background(), &domain.Entity{ + ID: id, KBID: "kb1", Name: id, Type: "entity", CreatedAt: now, UpdatedAt: now, + }); err != nil { + t.Fatal("create entity:", err) + } + } + + entities, err := store.GetEntitiesByIDs(context.Background(), "kb1", ids) + if err != nil { + t.Fatal("get entities by ids:", err) + } + if len(entities) != n { + t.Fatalf("got %d entities, want %d", len(entities), n) + } +} + // --- Relation tests --- func TestRelationCRUD(t *testing.T) { @@ -429,6 +490,47 @@ func TestRelationScopeAndKBBoundaryEnforced(t *testing.T) { } } +func TestGetRelationsByIDs(t *testing.T) { + store := newTestStore(t) + createTestKB(t, store, "kb1") + createTestKB(t, store, "kb2") + + now := time.Now().UTC() + for _, e := range []*domain.Entity{ + {ID: "ent-a", KBID: "kb1", Name: "Alice", Type: "person", CreatedAt: now, UpdatedAt: now}, + {ID: "ent-b", KBID: "kb1", Name: "Bob", Type: "person", CreatedAt: now, UpdatedAt: now}, + {ID: "ent-c", KBID: "kb2", Name: "Carol", Type: "person", CreatedAt: now, UpdatedAt: now}, + {ID: "ent-d", KBID: "kb2", Name: "Dan", Type: "person", CreatedAt: now, UpdatedAt: now}, + } { + if err := store.CreateEntity(context.Background(), e); err != nil { + t.Fatal("create entity:", err) + } + } + + for _, r := range []*domain.Relation{ + {ID: "rel-001", KBID: "kb1", SourceID: "ent-a", TargetID: "ent-b", Type: "knows", Weight: 1, ValidAt: now, CreatedAt: now}, + {ID: "rel-002", KBID: "kb2", SourceID: "ent-c", TargetID: "ent-d", Type: "knows", Weight: 1, ValidAt: now, CreatedAt: now}, + } { + if err := store.CreateRelation(context.Background(), r); err != nil { + t.Fatal("create relation:", err) + } + } + + rels, err := store.GetRelationsByIDs(context.Background(), "kb1", []string{"rel-001", "rel-002", "rel-missing"}) + if err != nil { + t.Fatal("get relations by ids:", err) + } + if len(rels) != 1 { + t.Fatalf("got %d relations, want 1", len(rels)) + } + if rels["rel-001"] == nil || rels["rel-001"].SourceID != "ent-a" || rels["rel-001"].TargetID != "ent-b" { + t.Fatal("expected rel-001 with correct endpoints") + } + if _, ok := rels["rel-002"]; ok { + t.Fatal("did not expect kb2 relation in kb1 result") + } +} + // --- FTS5 Search tests --- func TestSearchFTS(t *testing.T) { From c4fc348daac05d783cd4b5e09753bb75607662af Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 6 Apr 2026 10:11:37 +0000 Subject: [PATCH 3/8] Address validation nits in new tests Agent-Logs-Url: https://github.com/vndee/memex/sessions/1ec419ab-55f2-4063-8471-c0e8891d649c Co-authored-by: vndee <28271488+vndee@users.noreply.github.com> --- internal/search/graph_test.go | 1 - internal/server/helpers_test.go | 1 - 2 files changed, 2 deletions(-) diff --git a/internal/search/graph_test.go b/internal/search/graph_test.go index b6d4672..270c4ac 100644 --- a/internal/search/graph_test.go +++ b/internal/search/graph_test.go @@ -73,4 +73,3 @@ func TestHydrateEntityResults_UsesBatchLookup(t *testing.T) { t.Fatalf("GetEntity calls = %d, want 0", store.getEntityCalls) } } - diff --git a/internal/server/helpers_test.go b/internal/server/helpers_test.go index a191c7c..cc80d80 100644 --- a/internal/server/helpers_test.go +++ b/internal/server/helpers_test.go @@ -105,4 +105,3 @@ func TestHydrateSubgraph_UsesBatchLookups(t *testing.T) { t.Fatalf("GetRelation calls = %d, want 0", store.getRelationCalls) } } - From edb25c9afea436e13092ef98dd8e6914024502d8 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 6 Apr 2026 10:12:40 +0000 Subject: [PATCH 4/8] Strengthen subgraph hydration assertions Agent-Logs-Url: https://github.com/vndee/memex/sessions/1ec419ab-55f2-4063-8471-c0e8891d649c Co-authored-by: vndee <28271488+vndee@users.noreply.github.com> --- internal/server/helpers_test.go | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/internal/server/helpers_test.go b/internal/server/helpers_test.go index cc80d80..184ed01 100644 --- a/internal/server/helpers_test.go +++ b/internal/server/helpers_test.go @@ -92,6 +92,30 @@ func TestHydrateSubgraph_UsesBatchLookups(t *testing.T) { if len(hydrated.Edges) != 2 { t.Fatalf("got %d edges, want 2", len(hydrated.Edges)) } + nodesByID := make(map[string]domain.SubgraphNode, len(hydrated.Nodes)) + for _, n := range hydrated.Nodes { + nodesByID[n.ID] = n + } + if nodesByID["e1"].Name != "Alice" || nodesByID["e1"].Summary != "Engineer" { + t.Fatalf("expected hydrated node for e1, got %+v", nodesByID["e1"]) + } + if nodesByID["e2"].Name != "Bob" || nodesByID["e2"].Summary != "Manager" { + t.Fatalf("expected hydrated node for e2, got %+v", nodesByID["e2"]) + } + if nodesByID["missing"].Name != "" || nodesByID["missing"].Summary != "" { + t.Fatalf("expected fallback node for missing entity, got %+v", nodesByID["missing"]) + } + + edgesByID := make(map[string]domain.SubgraphEdge, len(hydrated.Edges)) + for _, e := range hydrated.Edges { + edgesByID[e.ID] = e + } + if edgesByID["r1"].SourceID != "e1" || edgesByID["r1"].TargetID != "e2" || edgesByID["r1"].Type != "knows" { + t.Fatalf("expected hydrated edge for r1, got %+v", edgesByID["r1"]) + } + if edgesByID["missing-rel"].SourceID != "e2" || edgesByID["missing-rel"].TargetID != "missing" || edgesByID["missing-rel"].Type != "mentions" { + t.Fatalf("expected fallback edge for missing-rel, got %+v", edgesByID["missing-rel"]) + } if store.getEntitiesByIDsCalls != 1 { t.Fatalf("GetEntitiesByIDs calls = %d, want 1", store.getEntitiesByIDsCalls) } From eb4750dbc13bebb2dde89c5ae8d3e113619e7456 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 6 Apr 2026 10:13:55 +0000 Subject: [PATCH 5/8] Refine batch row error handling Agent-Logs-Url: https://github.com/vndee/memex/sessions/1ec419ab-55f2-4063-8471-c0e8891d649c Co-authored-by: vndee <28271488+vndee@users.noreply.github.com> --- internal/storage/entities.go | 1 - internal/storage/relations.go | 1 - 2 files changed, 2 deletions(-) diff --git a/internal/storage/entities.go b/internal/storage/entities.go index b150f99..a42f437 100644 --- a/internal/storage/entities.go +++ b/internal/storage/entities.go @@ -68,7 +68,6 @@ func (s *SQLiteStore) GetEntitiesByIDs(ctx context.Context, kbID string, ids []s entities[e.ID] = e } if err := rows.Err(); err != nil { - rows.Close() return nil, fmt.Errorf("iterate entities by ids: %w", err) } if err := rows.Close(); err != nil { diff --git a/internal/storage/relations.go b/internal/storage/relations.go index 22db139..d95ae46 100644 --- a/internal/storage/relations.go +++ b/internal/storage/relations.go @@ -108,7 +108,6 @@ func (s *SQLiteStore) GetRelationsByIDs(ctx context.Context, kbID string, ids [] rels[r.ID] = r } if err := rows.Err(); err != nil { - rows.Close() return nil, fmt.Errorf("iterate relations by ids: %w", err) } if err := rows.Close(); err != nil { From 0810ac8a5a1c02a2219e80be558ef6c7ee2b9b63 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 6 Apr 2026 10:57:04 +0000 Subject: [PATCH 6/8] Address PR review feedback on batching and row lifecycle Agent-Logs-Url: https://github.com/vndee/memex/sessions/c7178a26-748e-4cfb-a92a-ba8c423be949 Co-authored-by: vndee <28271488+vndee@users.noreply.github.com> --- internal/search/graph.go | 43 +++++++++++++++---------- internal/search/graph_test.go | 59 +++++++++++++++++++++++++++++++++-- internal/storage/entities.go | 33 +++++++++++--------- internal/storage/relations.go | 27 +++++++++------- 4 files changed, 116 insertions(+), 46 deletions(-) diff --git a/internal/search/graph.go b/internal/search/graph.go index 81e60ba..05e50cb 100644 --- a/internal/search/graph.go +++ b/internal/search/graph.go @@ -182,28 +182,37 @@ func hydrateEntityResults( scores map[string]float64, limit int, ) ([]*domain.SearchResult, error) { - entitiesByID, err := store.GetEntitiesByIDs(ctx, kbID, ids) - if err != nil { - slog.Warn("graph entity batch lookup failed", "count", len(ids), "error", err) - return nil, fmt.Errorf("get entities by ids: %w", err) + if limit <= 0 || len(ids) == 0 { + return nil, nil } results := make([]*domain.SearchResult, 0, min(len(ids), limit)) - for _, id := range ids { - if len(results) >= limit { - break + batchSize := min(len(ids), max(limit*2, 64)) + for start := 0; start < len(ids) && len(results) < limit; start += batchSize { + end := min(start+batchSize, len(ids)) + batchIDs := ids[start:end] + + entitiesByID, err := store.GetEntitiesByIDs(ctx, kbID, batchIDs) + if err != nil { + slog.Warn("graph entity batch lookup failed", "count", len(batchIDs), "error", err) + return nil, fmt.Errorf("get entities by ids: %w", err) } - ent, ok := entitiesByID[id] - if !ok { - continue + for _, id := range batchIDs { + if len(results) >= limit { + break + } + ent, ok := entitiesByID[id] + if !ok { + continue + } + results = append(results, &domain.SearchResult{ + ID: ent.ID, + KBID: ent.KBID, + Type: domain.ItemEntity, + Content: ent.Name + ": " + ent.Summary, + Score: scores[id], + }) } - results = append(results, &domain.SearchResult{ - ID: ent.ID, - KBID: ent.KBID, - Type: domain.ItemEntity, - Content: ent.Name + ": " + ent.Summary, - Score: scores[id], - }) } return results, nil } diff --git a/internal/search/graph_test.go b/internal/search/graph_test.go index 270c4ac..22bc13d 100644 --- a/internal/search/graph_test.go +++ b/internal/search/graph_test.go @@ -2,6 +2,7 @@ package search import ( "context" + "fmt" "testing" "time" @@ -16,8 +17,9 @@ func init() { type countingSearchStore struct { storage.Store - getEntityCalls int - getEntitiesByIDsCalls int + getEntityCalls int + getEntitiesByIDsCalls int + getEntitiesByIDsBatchSizes []int } func (s *countingSearchStore) GetEntity(ctx context.Context, kbID, id string) (*domain.Entity, error) { @@ -27,6 +29,7 @@ func (s *countingSearchStore) GetEntity(ctx context.Context, kbID, id string) (* func (s *countingSearchStore) GetEntitiesByIDs(ctx context.Context, kbID string, ids []string) (map[string]*domain.Entity, error) { s.getEntitiesByIDsCalls++ + s.getEntitiesByIDsBatchSizes = append(s.getEntitiesByIDsBatchSizes, len(ids)) return s.Store.GetEntitiesByIDs(ctx, kbID, ids) } @@ -73,3 +76,55 @@ func TestHydrateEntityResults_UsesBatchLookup(t *testing.T) { t.Fatalf("GetEntity calls = %d, want 0", store.getEntityCalls) } } + +func TestHydrateEntityResults_DoesNotFetchAllIDsWhenLimitSmall(t *testing.T) { + sqliteStore, err := storage.NewSQLiteStore(":memory:") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = sqliteStore.Close() }) + + now := time.Now().UTC() + if err := sqliteStore.CreateKB(context.Background(), &domain.KnowledgeBase{ + ID: "kb1", Name: "KB 1", + EmbedConfig: domain.EmbedConfig{Provider: "ollama", Model: "nomic-embed-text"}, + LLMConfig: domain.LLMConfig{Provider: "ollama", Model: "llama3.2"}, + CreatedAt: now, + }); err != nil { + t.Fatal(err) + } + if err := sqliteStore.CreateEntity(context.Background(), &domain.Entity{ + ID: "e001", KBID: "kb1", Name: "Alice", Type: "person", Summary: "Engineer", CreatedAt: now, UpdatedAt: now, + }); err != nil { + t.Fatal(err) + } + + store := &countingSearchStore{Store: sqliteStore} + ids := make([]string, 0, 200) + scores := make(map[string]float64, 200) + for i := 1; i <= 200; i++ { + id := fmt.Sprintf("missing-%03d", i) + if i == 1 { + id = "e001" + } + ids = append(ids, id) + scores[id] = 1 + } + + results, err := hydrateEntityResults(context.Background(), store, "kb1", ids, scores, 1) + if err != nil { + t.Fatal(err) + } + if len(results) != 1 { + t.Fatalf("got %d results, want 1", len(results)) + } + if store.getEntitiesByIDsCalls != 1 { + t.Fatalf("GetEntitiesByIDs calls = %d, want 1", store.getEntitiesByIDsCalls) + } + if len(store.getEntitiesByIDsBatchSizes) != 1 { + t.Fatalf("batch size captures = %d, want 1", len(store.getEntitiesByIDsBatchSizes)) + } + if store.getEntitiesByIDsBatchSizes[0] >= len(ids) { + t.Fatalf("batch fetched %d IDs, expected less than total %d", store.getEntitiesByIDsBatchSizes[0], len(ids)) + } +} diff --git a/internal/storage/entities.go b/internal/storage/entities.go index a42f437..12e21b9 100644 --- a/internal/storage/entities.go +++ b/internal/storage/entities.go @@ -55,23 +55,26 @@ func (s *SQLiteStore) GetEntitiesByIDs(ctx context.Context, kbID string, ids []s args = append(args, id) } - rows, err := s.db.QueryContext(ctx, query, args...) - if err != nil { - return nil, fmt.Errorf("query entities by ids: %w", err) - } - for rows.Next() { - e, err := scanEntity(rows) + if err := func() error { + rows, err := s.db.QueryContext(ctx, query, args...) if err != nil { - rows.Close() - return nil, err + return fmt.Errorf("query entities by ids: %w", err) } - entities[e.ID] = e - } - if err := rows.Err(); err != nil { - return nil, fmt.Errorf("iterate entities by ids: %w", err) - } - if err := rows.Close(); err != nil { - return nil, fmt.Errorf("close entities by ids rows: %w", err) + defer rows.Close() + + for rows.Next() { + e, err := scanEntity(rows) + if err != nil { + return err + } + entities[e.ID] = e + } + if err := rows.Err(); err != nil { + return fmt.Errorf("iterate entities by ids: %w", err) + } + return nil + }(); err != nil { + return nil, err } } return entities, nil diff --git a/internal/storage/relations.go b/internal/storage/relations.go index d95ae46..76904fd 100644 --- a/internal/storage/relations.go +++ b/internal/storage/relations.go @@ -99,19 +99,22 @@ func (s *SQLiteStore) GetRelationsByIDs(ctx context.Context, kbID string, ids [] if err != nil { return nil, fmt.Errorf("query relations by ids: %w", err) } - for rows.Next() { - r, err := scanRelation(rows) - if err != nil { - rows.Close() - return nil, err + if err := func() error { + defer rows.Close() + + for rows.Next() { + r, err := scanRelation(rows) + if err != nil { + return err + } + rels[r.ID] = r } - rels[r.ID] = r - } - if err := rows.Err(); err != nil { - return nil, fmt.Errorf("iterate relations by ids: %w", err) - } - if err := rows.Close(); err != nil { - return nil, fmt.Errorf("close relations by ids rows: %w", err) + if err := rows.Err(); err != nil { + return fmt.Errorf("iterate relations by ids: %w", err) + } + return nil + }(); err != nil { + return nil, err } } return rels, nil From 4fbefaf5bf37180680b3dc7adf8e5d366ae8b771 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 6 Apr 2026 10:58:04 +0000 Subject: [PATCH 7/8] Name hydration batch size constant Agent-Logs-Url: https://github.com/vndee/memex/sessions/c7178a26-748e-4cfb-a92a-ba8c423be949 Co-authored-by: vndee <28271488+vndee@users.noreply.github.com> --- internal/search/graph.go | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/internal/search/graph.go b/internal/search/graph.go index 05e50cb..b6c971d 100644 --- a/internal/search/graph.go +++ b/internal/search/graph.go @@ -16,6 +16,9 @@ const ( pprAlpha = 0.15 // teleport probability (jump back to seed) pprMaxIter = 20 // max power iterations pprEpsilon = 1e-6 // convergence threshold + + // Minimum hydration batch size avoids many tiny IN queries when limit is small. + minHydrationBatchSize = 64 ) // searchGraph expands seed entity IDs via BFS on the knowledge graph, @@ -187,7 +190,7 @@ func hydrateEntityResults( } results := make([]*domain.SearchResult, 0, min(len(ids), limit)) - batchSize := min(len(ids), max(limit*2, 64)) + batchSize := min(len(ids), max(limit*2, minHydrationBatchSize)) for start := 0; start < len(ids) && len(results) < limit; start += batchSize { end := min(start+batchSize, len(ids)) batchIDs := ids[start:end] From c509585f15d88bd0d51eec40cd74fc3d9d299ac2 Mon Sep 17 00:00:00 2001 From: Duy Huynh Date: Mon, 6 Apr 2026 21:34:12 +0700 Subject: [PATCH 8/8] Fix subgraph hydration fallback and memory usage --- internal/server/helpers.go | 124 ++++++++++++------- internal/server/helpers_test.go | 165 +++++++++++++++++++++++-- internal/storage/subgraph_hydration.go | 139 +++++++++++++++++++++ 3 files changed, 375 insertions(+), 53 deletions(-) create mode 100644 internal/storage/subgraph_hydration.go diff --git a/internal/server/helpers.go b/internal/server/helpers.go index 47b2b27..4210ea4 100644 --- a/internal/server/helpers.go +++ b/internal/server/helpers.go @@ -2,7 +2,7 @@ package server import ( "context" - "fmt" + "log/slog" "time" "github.com/vndee/memex/internal/domain" @@ -10,6 +10,11 @@ import ( "github.com/vndee/memex/internal/storage" ) +type subgraphMetadataLoader interface { + GetSubgraphEntitiesByIDs(ctx context.Context, kbID string, ids []string) (map[string]storage.SubgraphEntityMetadata, error) + GetSubgraphRelationsByIDs(ctx context.Context, kbID string, ids []string) (map[string]storage.SubgraphRelationMetadata, error) +} + // buildKB constructs a KnowledgeBase with defaults applied for empty fields. // Shared by both HTTP and MCP handlers. func buildKB(id, name, desc, embedProvider, embedModel, llmProvider, llmModel string) *domain.KnowledgeBase { @@ -48,64 +53,95 @@ func buildKB(id, name, desc, embedProvider, embedModel, llmProvider, llmModel st // HydrateSubgraph enriches a raw SubgraphResult with entity and relation metadata // from storage. Used by MCP, HTTP, and CLI graph traversal handlers. func HydrateSubgraph(ctx context.Context, store storage.Store, kbID string, sg graph.SubgraphResult) (*domain.Subgraph, error) { + nodeIndex := make(map[string]int, len(sg.Nodes)) nodeIDs := make([]string, 0, len(sg.Nodes)) + nodes := make([]domain.SubgraphNode, 0, len(sg.Nodes)) for id := range sg.Nodes { + nodeIndex[id] = len(nodes) nodeIDs = append(nodeIDs, id) - } - entitiesByID, err := store.GetEntitiesByIDs(ctx, kbID, nodeIDs) - if err != nil { - return nil, fmt.Errorf("get subgraph entities by ids: %w", err) - } - - nodes := make([]domain.SubgraphNode, 0, len(sg.Nodes)) - for id, dist := range sg.Nodes { - ent, ok := entitiesByID[id] - if !ok { - nodes = append(nodes, domain.SubgraphNode{ - ID: id, Distance: dist, - }) - continue - } nodes = append(nodes, domain.SubgraphNode{ - ID: ent.ID, - Name: ent.Name, - Type: ent.Type, - Summary: ent.Summary, - Distance: dist, + ID: id, + Distance: sg.Nodes[id], }) } + edgeIndex := make(map[string]int, len(sg.Edges)) edgeIDs := make([]string, 0, len(sg.Edges)) + edges := make([]domain.SubgraphEdge, 0, len(sg.Edges)) for _, e := range sg.Edges { + edgeIndex[e.RelID] = len(edges) edgeIDs = append(edgeIDs, e.RelID) + edges = append(edges, domain.SubgraphEdge{ + ID: e.RelID, + SourceID: e.SourceID, + TargetID: e.TargetID, + Type: e.Type, + Weight: e.Weight, + }) } - relationsByID, err := store.GetRelationsByIDs(ctx, kbID, edgeIDs) - if err != nil { - return nil, fmt.Errorf("get subgraph relations by ids: %w", err) + + if loader, ok := store.(subgraphMetadataLoader); ok { + entitiesByID, err := loader.GetSubgraphEntitiesByIDs(ctx, kbID, nodeIDs) + if err != nil { + slog.Warn("subgraph entity hydration failed", "count", len(nodeIDs), "error", err) + } else { + for id, ent := range entitiesByID { + idx, ok := nodeIndex[id] + if !ok { + continue + } + nodes[idx].ID = ent.ID + nodes[idx].Name = ent.Name + nodes[idx].Type = ent.Type + nodes[idx].Summary = ent.Summary + } + } + + relationsByID, err := loader.GetSubgraphRelationsByIDs(ctx, kbID, edgeIDs) + if err != nil { + slog.Warn("subgraph relation hydration failed", "count", len(edgeIDs), "error", err) + } else { + for id, rel := range relationsByID { + idx, ok := edgeIndex[id] + if !ok { + continue + } + edges[idx].ID = rel.ID + edges[idx].SourceID = rel.SourceID + edges[idx].TargetID = rel.TargetID + edges[idx].Type = rel.Type + edges[idx].Weight = rel.Weight + edges[idx].ValidAt = rel.ValidAt + edges[idx].InvalidAt = rel.InvalidAt + } + } + + return &domain.Subgraph{Nodes: nodes, Edges: edges}, nil } - edges := make([]domain.SubgraphEdge, 0, len(sg.Edges)) - for _, e := range sg.Edges { - rel, ok := relationsByID[e.RelID] - if !ok { - edges = append(edges, domain.SubgraphEdge{ - ID: e.RelID, - SourceID: e.SourceID, - TargetID: e.TargetID, - Type: e.Type, - Weight: e.Weight, - }) + for i := range nodes { + ent, err := store.GetEntity(ctx, kbID, nodes[i].ID) + if err != nil { continue } - edges = append(edges, domain.SubgraphEdge{ - ID: rel.ID, - SourceID: rel.SourceID, - TargetID: rel.TargetID, - Type: rel.Type, - Weight: rel.Weight, - ValidAt: rel.ValidAt, - InvalidAt: rel.InvalidAt, - }) + nodes[i].ID = ent.ID + nodes[i].Name = ent.Name + nodes[i].Type = ent.Type + nodes[i].Summary = ent.Summary + } + + for i := range edges { + rel, err := store.GetRelation(ctx, kbID, edges[i].ID) + if err != nil { + continue + } + edges[i].ID = rel.ID + edges[i].SourceID = rel.SourceID + edges[i].TargetID = rel.TargetID + edges[i].Type = rel.Type + edges[i].Weight = rel.Weight + edges[i].ValidAt = rel.ValidAt + edges[i].InvalidAt = rel.InvalidAt } return &domain.Subgraph{Nodes: nodes, Edges: edges}, nil diff --git a/internal/server/helpers_test.go b/internal/server/helpers_test.go index 184ed01..407a6ff 100644 --- a/internal/server/helpers_test.go +++ b/internal/server/helpers_test.go @@ -2,6 +2,7 @@ package server import ( "context" + "errors" "testing" "time" @@ -17,10 +18,14 @@ func init() { type countingServerStore struct { storage.Store - getEntityCalls int - getEntitiesByIDsCalls int - getRelationCalls int - getRelationsByIDsCalls int + getEntityCalls int + getEntitiesByIDsCalls int + getRelationCalls int + getRelationsByIDsCalls int + getSubgraphEntitiesByIDsCalls int + getSubgraphRelationsByIDsCalls int + failSubgraphEntitiesByIDs bool + failSubgraphRelationsByIDs bool } func (s *countingServerStore) GetEntity(ctx context.Context, kbID, id string) (*domain.Entity, error) { @@ -43,7 +48,35 @@ func (s *countingServerStore) GetRelationsByIDs(ctx context.Context, kbID string return s.Store.GetRelationsByIDs(ctx, kbID, ids) } -func TestHydrateSubgraph_UsesBatchLookups(t *testing.T) { +func (s *countingServerStore) GetSubgraphEntitiesByIDs(ctx context.Context, kbID string, ids []string) (map[string]storage.SubgraphEntityMetadata, error) { + s.getSubgraphEntitiesByIDsCalls++ + if s.failSubgraphEntitiesByIDs { + return nil, errors.New("entity hydration failed") + } + loader, ok := s.Store.(interface { + GetSubgraphEntitiesByIDs(context.Context, string, []string) (map[string]storage.SubgraphEntityMetadata, error) + }) + if !ok { + return nil, errors.New("subgraph entity loader not supported") + } + return loader.GetSubgraphEntitiesByIDs(ctx, kbID, ids) +} + +func (s *countingServerStore) GetSubgraphRelationsByIDs(ctx context.Context, kbID string, ids []string) (map[string]storage.SubgraphRelationMetadata, error) { + s.getSubgraphRelationsByIDsCalls++ + if s.failSubgraphRelationsByIDs { + return nil, errors.New("relation hydration failed") + } + loader, ok := s.Store.(interface { + GetSubgraphRelationsByIDs(context.Context, string, []string) (map[string]storage.SubgraphRelationMetadata, error) + }) + if !ok { + return nil, errors.New("subgraph relation loader not supported") + } + return loader.GetSubgraphRelationsByIDs(ctx, kbID, ids) +} + +func TestHydrateSubgraph_UsesLightweightBatchLookups(t *testing.T) { sqliteStore, err := storage.NewSQLiteStore(":memory:") if err != nil { t.Fatal(err) @@ -116,11 +149,17 @@ func TestHydrateSubgraph_UsesBatchLookups(t *testing.T) { if edgesByID["missing-rel"].SourceID != "e2" || edgesByID["missing-rel"].TargetID != "missing" || edgesByID["missing-rel"].Type != "mentions" { t.Fatalf("expected fallback edge for missing-rel, got %+v", edgesByID["missing-rel"]) } - if store.getEntitiesByIDsCalls != 1 { - t.Fatalf("GetEntitiesByIDs calls = %d, want 1", store.getEntitiesByIDsCalls) + if store.getEntitiesByIDsCalls != 0 { + t.Fatalf("GetEntitiesByIDs calls = %d, want 0", store.getEntitiesByIDsCalls) } - if store.getRelationsByIDsCalls != 1 { - t.Fatalf("GetRelationsByIDs calls = %d, want 1", store.getRelationsByIDsCalls) + if store.getRelationsByIDsCalls != 0 { + t.Fatalf("GetRelationsByIDs calls = %d, want 0", store.getRelationsByIDsCalls) + } + if store.getSubgraphEntitiesByIDsCalls != 1 { + t.Fatalf("GetSubgraphEntitiesByIDs calls = %d, want 1", store.getSubgraphEntitiesByIDsCalls) + } + if store.getSubgraphRelationsByIDsCalls != 1 { + t.Fatalf("GetSubgraphRelationsByIDs calls = %d, want 1", store.getSubgraphRelationsByIDsCalls) } if store.getEntityCalls != 0 { t.Fatalf("GetEntity calls = %d, want 0", store.getEntityCalls) @@ -129,3 +168,111 @@ func TestHydrateSubgraph_UsesBatchLookups(t *testing.T) { t.Fatalf("GetRelation calls = %d, want 0", store.getRelationCalls) } } + +func TestHydrateSubgraph_FallsBackWhenBatchHydrationFails(t *testing.T) { + sqliteStore, err := storage.NewSQLiteStore(":memory:") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = sqliteStore.Close() }) + + now := time.Now().UTC() + if err := sqliteStore.CreateKB(context.Background(), &domain.KnowledgeBase{ + ID: "kb1", Name: "KB 1", + EmbedConfig: domain.EmbedConfig{Provider: "ollama", Model: "nomic-embed-text"}, + LLMConfig: domain.LLMConfig{Provider: "ollama", Model: "llama3.2"}, + CreatedAt: now, + }); err != nil { + t.Fatal(err) + } + if err := sqliteStore.CreateEntity(context.Background(), &domain.Entity{ + ID: "e1", + KBID: "kb1", + Name: "Alice", + Type: "person", + Summary: "Engineer", + Embedding: []float32{1, 2, 3}, + CreatedAt: now, + UpdatedAt: now, + }); err != nil { + t.Fatal(err) + } + if err := sqliteStore.CreateEntity(context.Background(), &domain.Entity{ + ID: "missing", + KBID: "kb1", + Name: "Unhydrated", + Type: "person", + Summary: "Fallback target", + Embedding: []float32{4, 5, 6}, + CreatedAt: now, + UpdatedAt: now, + }); err != nil { + t.Fatal(err) + } + if err := sqliteStore.CreateRelation(context.Background(), &domain.Relation{ + ID: "r1", + KBID: "kb1", + SourceID: "e1", + TargetID: "missing", + Type: "knows", + Summary: "Alice knows someone", + Weight: 1, + Embedding: []float32{1, 2, 3}, + ValidAt: now, + CreatedAt: now, + }); err != nil { + t.Fatal(err) + } + + store := &countingServerStore{ + Store: sqliteStore, + failSubgraphEntitiesByIDs: true, + failSubgraphRelationsByIDs: true, + } + subgraph := graph.SubgraphResult{ + Nodes: map[string]int{"e1": 0, "missing": 1}, + Edges: []graph.SubgraphEdge{ + {RelID: "r1", SourceID: "e1", TargetID: "missing", Type: "knows", Weight: 1}, + }, + } + + hydrated, err := HydrateSubgraph(context.Background(), store, "kb1", subgraph) + if err != nil { + t.Fatal(err) + } + if len(hydrated.Nodes) != 2 { + t.Fatalf("got %d nodes, want 2", len(hydrated.Nodes)) + } + if len(hydrated.Edges) != 1 { + t.Fatalf("got %d edges, want 1", len(hydrated.Edges)) + } + + nodesByID := make(map[string]domain.SubgraphNode, len(hydrated.Nodes)) + for _, n := range hydrated.Nodes { + nodesByID[n.ID] = n + } + if nodesByID["e1"].Name != "" || nodesByID["e1"].Summary != "" { + t.Fatalf("expected raw fallback node for e1, got %+v", nodesByID["e1"]) + } + if nodesByID["e1"].Distance != 0 || nodesByID["missing"].Distance != 1 { + t.Fatalf("expected fallback distances to be preserved, got %+v", hydrated.Nodes) + } + + edge := hydrated.Edges[0] + if edge.ID != "r1" || edge.SourceID != "e1" || edge.TargetID != "missing" || edge.Type != "knows" || edge.Weight != 1 { + t.Fatalf("expected raw fallback edge, got %+v", edge) + } + if !edge.ValidAt.IsZero() || edge.InvalidAt != nil { + t.Fatalf("expected temporal metadata to remain empty on fallback, got %+v", edge) + } + + if store.getSubgraphEntitiesByIDsCalls != 1 { + t.Fatalf("GetSubgraphEntitiesByIDs calls = %d, want 1", store.getSubgraphEntitiesByIDsCalls) + } + if store.getSubgraphRelationsByIDsCalls != 1 { + t.Fatalf("GetSubgraphRelationsByIDs calls = %d, want 1", store.getSubgraphRelationsByIDsCalls) + } + if store.getEntityCalls != 0 || store.getRelationCalls != 0 { + t.Fatalf("expected no per-item fallback queries, got GetEntity=%d GetRelation=%d", store.getEntityCalls, store.getRelationCalls) + } +} diff --git a/internal/storage/subgraph_hydration.go b/internal/storage/subgraph_hydration.go new file mode 100644 index 0000000..cd686e3 --- /dev/null +++ b/internal/storage/subgraph_hydration.go @@ -0,0 +1,139 @@ +package storage + +import ( + "context" + "database/sql" + "fmt" + "strings" + "time" +) + +const maxHydrationLookupIDs = 900 + +// SubgraphEntityMetadata is the lightweight entity shape needed to hydrate a subgraph. +type SubgraphEntityMetadata struct { + ID string + Name string + Type string + Summary string +} + +// SubgraphRelationMetadata is the lightweight relation shape needed to hydrate a subgraph. +type SubgraphRelationMetadata struct { + ID string + SourceID string + TargetID string + Type string + Weight float64 + ValidAt time.Time + InvalidAt *time.Time +} + +// GetSubgraphEntitiesByIDs returns lightweight entity metadata for subgraph hydration. +func (s *SQLiteStore) GetSubgraphEntitiesByIDs(ctx context.Context, kbID string, ids []string) (map[string]SubgraphEntityMetadata, error) { + entities := make(map[string]SubgraphEntityMetadata) + if len(ids) == 0 { + return entities, nil + } + + for start := 0; start < len(ids); start += maxHydrationLookupIDs { + end := min(start+maxHydrationLookupIDs, len(ids)) + batch := ids[start:end] + + placeholders := strings.TrimSuffix(strings.Repeat("?,", len(batch)), ",") + query := `SELECT id, name, type, summary + FROM entities WHERE kb_id = ? AND id IN (` + placeholders + `)` + + args := make([]any, 0, len(batch)+1) + args = append(args, kbID) + for _, id := range batch { + args = append(args, id) + } + + if err := func() error { + rows, err := s.db.QueryContext(ctx, query, args...) + if err != nil { + return fmt.Errorf("query subgraph entities by ids: %w", err) + } + defer rows.Close() + + for rows.Next() { + var meta SubgraphEntityMetadata + if err := rows.Scan(&meta.ID, &meta.Name, &meta.Type, &meta.Summary); err != nil { + return fmt.Errorf("scan subgraph entity: %w", err) + } + entities[meta.ID] = meta + } + if err := rows.Err(); err != nil { + return fmt.Errorf("iterate subgraph entities by ids: %w", err) + } + return nil + }(); err != nil { + return nil, err + } + } + + return entities, nil +} + +// GetSubgraphRelationsByIDs returns lightweight relation metadata for subgraph hydration. +func (s *SQLiteStore) GetSubgraphRelationsByIDs(ctx context.Context, kbID string, ids []string) (map[string]SubgraphRelationMetadata, error) { + rels := make(map[string]SubgraphRelationMetadata) + if len(ids) == 0 { + return rels, nil + } + + for start := 0; start < len(ids); start += maxHydrationLookupIDs { + end := min(start+maxHydrationLookupIDs, len(ids)) + batch := ids[start:end] + + placeholders := strings.TrimSuffix(strings.Repeat("?,", len(batch)), ",") + query := `SELECT id, source_id, target_id, type, weight, valid_at, invalid_at + FROM relations WHERE kb_id = ? AND id IN (` + placeholders + `)` + + args := make([]any, 0, len(batch)+1) + args = append(args, kbID) + for _, id := range batch { + args = append(args, id) + } + + if err := func() error { + rows, err := s.db.QueryContext(ctx, query, args...) + if err != nil { + return fmt.Errorf("query subgraph relations by ids: %w", err) + } + defer rows.Close() + + for rows.Next() { + var meta SubgraphRelationMetadata + var validAt string + var invalidAt sql.NullString + if err := rows.Scan( + &meta.ID, + &meta.SourceID, + &meta.TargetID, + &meta.Type, + &meta.Weight, + &validAt, + &invalidAt, + ); err != nil { + return fmt.Errorf("scan subgraph relation: %w", err) + } + meta.ValidAt, _ = time.Parse(time.RFC3339Nano, validAt) + if invalidAt.Valid { + t, _ := time.Parse(time.RFC3339Nano, invalidAt.String) + meta.InvalidAt = &t + } + rels[meta.ID] = meta + } + if err := rows.Err(); err != nil { + return fmt.Errorf("iterate subgraph relations by ids: %w", err) + } + return nil + }(); err != nil { + return nil, err + } + } + + return rels, nil +}