diff --git a/internal/vector/sqlitevec/client.go b/internal/vector/sqlitevec/client.go index b88bcb5..5d63e59 100644 --- a/internal/vector/sqlitevec/client.go +++ b/internal/vector/sqlitevec/client.go @@ -5,6 +5,7 @@ import ( "context" "database/sql" "fmt" + "sort" "strconv" "strings" "sync" @@ -314,67 +315,13 @@ func (c *Client) Query(ctx context.Context, query string, limit int, where map[s } defer c.readMu.RUnlock() - // Build query with filters - // vec0 supports WHERE clauses on metadata columns - args := []any{queryBlob} - - sqlQuery := ` - SELECT - doc_id, - distance, - sqlite_id, - doc_type, - field_type, - project, - scope - FROM vectors - WHERE embedding MATCH ? - ` - - // Add filters - these work with vec0 metadata columns - if docType, ok := where["doc_type"].(string); ok && docType != "" { - sqlQuery += " AND doc_type = ?" - args = append(args, docType) - } - if project, ok := where["project"].(string); ok && project != "" { - // Include project-specific OR global scope - sqlQuery += " AND (project = ? OR scope = 'global')" - args = append(args, project) - } - - sqlQuery += " ORDER BY distance LIMIT ?" - args = append(args, limit) - - rows, err := c.db.QueryContext(ctx, sqlQuery, args...) + // vec0 answers equality filters on metadata columns, but not an OR between them, so a project filter + // (the project's own documents, plus global ones) is asked as two searches and merged. + docType, _ := where["doc_type"].(string) + project, _ := where["project"].(string) + results, err := c.nearest(ctx, queryBlob, docType, project, limit) if err != nil { - return nil, fmt.Errorf("query vectors: %w", err) - } - defer rows.Close() - - var results []QueryResult - for rows.Next() { - var r QueryResult - var sqliteID int64 - var docType, fieldType, project, scope sql.NullString - - if err := rows.Scan(&r.ID, &r.Distance, &sqliteID, &docType, &fieldType, &project, &scope); err != nil { - return nil, fmt.Errorf("scan row: %w", err) - } - - r.Similarity = DistanceToSimilarity(r.Distance) - r.Metadata = map[string]any{ - "sqlite_id": float64(sqliteID), // Keep as float64 for compatibility - "doc_type": docType.String, - "field_type": fieldType.String, - "project": project.String, - "scope": scope.String, - } - - results = append(results, r) - } - - if err = rows.Err(); err != nil { - return nil, fmt.Errorf("iterate rows: %w", err) + return nil, err } // Cache the results @@ -529,60 +476,108 @@ func (c *Client) QueryMultiField(ctx context.Context, query string, limit int, d } defer c.readMu.RUnlock() - // Query with field type aggregation - get best match per document - sqlQuery := ` - WITH ranked_results AS ( - SELECT - doc_id, - distance, - sqlite_id, - doc_type, - field_type, - project, - scope, - ROW_NUMBER() OVER (PARTITION BY sqlite_id ORDER BY distance ASC) as rn - FROM vectors - WHERE embedding MATCH ? - AND doc_type = ? - AND (project = ? OR scope = 'global') - ) - SELECT doc_id, distance, sqlite_id, doc_type, field_type, project, scope - FROM ranked_results - WHERE rn = 1 - ORDER BY distance - LIMIT ? - ` + // Each observation has several vector documents (one per field). Ask for enough of them that `limit` + // distinct observations are likely to be among them, then keep the best match per observation. + found, err := c.nearest(ctx, queryBlob, docType, project, limit*multiFieldFetchFactor) + if err != nil { + return nil, err + } + results := make([]QueryResult, 0, limit) + seen := make(map[float64]bool, limit) + for _, r := range found { // best first + id, _ := r.Metadata["sqlite_id"].(float64) + if seen[id] { + continue + } + seen[id] = true + results = append(results, r) + if len(results) == limit { + break + } + } + return results, nil +} - rows, err := c.db.QueryContext(ctx, sqlQuery, queryBlob, docType, project, limit) +// multiFieldFetchFactor is how many vector documents QueryMultiField asks for per wanted observation. +const multiFieldFetchFactor = 6 + +// nearest runs the nearest-neighbour search for a document type and, when project is not empty, for that +// project's documents and the global ones. The caller holds the read lock. Results are best first, at most limit. +func (c *Client) nearest(ctx context.Context, queryBlob []byte, docType, project string, limit int) ([]QueryResult, error) { + base := " WHERE embedding MATCH ?" + baseArgs := []any{queryBlob} + if docType != "" { + base += " AND doc_type = ?" + baseArgs = append(baseArgs, docType) + } + if project == "" { + return c.knn(ctx, base, baseArgs, limit) + } + own, err := c.knn(ctx, base+" AND project = ?", append(append([]any{}, baseArgs...), project), limit) + if err != nil { + return nil, err + } + global, err := c.knn(ctx, base+" AND scope = 'global'", baseArgs, limit) + if err != nil { + return nil, err + } + return mergeNearest(limit, own, global), nil +} + +// knn runs one vec0 nearest-neighbour query. where holds the MATCH and filters vec0 can answer on their own. +func (c *Client) knn(ctx context.Context, where string, args []any, limit int) ([]QueryResult, error) { + rows, err := c.db.QueryContext(ctx, + "SELECT doc_id, distance, sqlite_id, doc_type, field_type, project, scope FROM vectors"+where+" ORDER BY distance LIMIT ?", + append(append([]any{}, args...), limit)...) if err != nil { return nil, fmt.Errorf("query vectors: %w", err) } defer rows.Close() - // Pre-allocate with limit to avoid repeated slice growth - results := make([]QueryResult, 0, limit) + var results []QueryResult for rows.Next() { var r QueryResult var sqliteID int64 - var docTypeVal, fieldType, projectVal, scope sql.NullString + var docType, fieldType, project, scope sql.NullString - if err := rows.Scan(&r.ID, &r.Distance, &sqliteID, &docTypeVal, &fieldType, &projectVal, &scope); err != nil { + if err := rows.Scan(&r.ID, &r.Distance, &sqliteID, &docType, &fieldType, &project, &scope); err != nil { return nil, fmt.Errorf("scan row: %w", err) } r.Similarity = DistanceToSimilarity(r.Distance) r.Metadata = map[string]any{ - "sqlite_id": float64(sqliteID), - "doc_type": docTypeVal.String, + "sqlite_id": float64(sqliteID), // Keep as float64 for compatibility + "doc_type": docType.String, "field_type": fieldType.String, - "project": projectVal.String, + "project": project.String, "scope": scope.String, } results = append(results, r) } + if err = rows.Err(); err != nil { + return nil, fmt.Errorf("iterate rows: %w", err) + } + return results, nil +} - return results, rows.Err() +// mergeNearest combines result lists into one, best first, without repeating a document, at most limit long. +func mergeNearest(limit int, lists ...[]QueryResult) []QueryResult { + var all []QueryResult + seen := map[string]bool{} + for _, list := range lists { + for _, r := range list { + if !seen[r.ID] { + seen[r.ID] = true + all = append(all, r) + } + } + } + sort.SliceStable(all, func(i, j int) bool { return all[i].Distance < all[j].Distance }) + if len(all) > limit { + all = all[:limit] + } + return all } // acquireRLockWithContext acquires a read lock on mu, respecting ctx cancellation. diff --git a/internal/vector/sqlitevec/client_test.go b/internal/vector/sqlitevec/client_test.go index 4c3e126..7b8cbe2 100644 --- a/internal/vector/sqlitevec/client_test.go +++ b/internal/vector/sqlitevec/client_test.go @@ -405,6 +405,50 @@ func TestClient_Query_WithProjectFilter(t *testing.T) { results, err := client.Query(context.Background(), "authentication security", 10, nil) require.NoError(t, err) assert.NotEmpty(t, results, "Should find some results") + assert.Len(t, results, 3) + + ids := func(rs []QueryResult) []string { + var out []string + for _, r := range rs { + out = append(out, r.ID) + } + return out + } + + // With a project filter: the project's own documents and the global ones, not another project's. + results, err = client.Query(context.Background(), "authentication security", 10, BuildWhereFilter(DocTypeObservation, "project-a")) + require.NoError(t, err, "a project filter must be answerable (vec0 rejects an OR between metadata filters)") + assert.ElementsMatch(t, []string{"obs-1", "obs-3"}, ids(results), "project-a's note and the global one, not project-b's own") + + results, err = client.Query(context.Background(), "authentication security", 10, BuildWhereFilter(DocTypeObservation, "project-b")) + require.NoError(t, err) + assert.ElementsMatch(t, []string{"obs-2", "obs-3"}, ids(results)) + + // Best first, and the limit applies to the merged list. + results, err = client.Query(context.Background(), "authentication security", 1, BuildWhereFilter(DocTypeObservation, "project-a")) + require.NoError(t, err) + require.Len(t, results, 1) + all, err := client.Query(context.Background(), "authentication security", 10, BuildWhereFilter(DocTypeObservation, "project-a")) + require.NoError(t, err) + assert.Equal(t, all[0].ID, results[0].ID, "the one result is the best of the merged list") + assert.LessOrEqual(t, all[0].Distance, all[1].Distance) + + // A project nobody wrote to still sees the global documents. + results, err = client.Query(context.Background(), "authentication security", 10, BuildWhereFilter(DocTypeObservation, "project-z")) + require.NoError(t, err) + assert.Equal(t, []string{"obs-3"}, ids(results)) +} + +func TestMergeNearest(t *testing.T) { + r := func(id string, d float64) QueryResult { return QueryResult{ID: id, Distance: d} } + got := mergeNearest(3, []QueryResult{r("a", 0.5), r("b", 0.9)}, []QueryResult{r("b", 0.9), r("c", 0.1), r("d", 0.7)}) + var ids []string + for _, x := range got { + ids = append(ids, x.ID) + } + assert.Equal(t, []string{"c", "a", "d"}, ids, "best first, a document found twice counts once, cut to the limit") + assert.Empty(t, mergeNearest(5)) + assert.Len(t, mergeNearest(5, []QueryResult{r("a", 1)}, nil), 1) } func TestClient_IsConnected(t *testing.T) { @@ -1105,8 +1149,6 @@ func TestClient_QueryBatch_WithContextCancellation(t *testing.T) { // ============================================================================= func TestClient_QueryMultiField_Basic(t *testing.T) { - t.Skip("QueryMultiField SQL query needs 'k' parameter fix for sqlite-vec") - db, dbCleanup := testDB(t) defer dbCleanup() @@ -1173,8 +1215,6 @@ func TestClient_QueryMultiField_Basic(t *testing.T) { } func TestClient_QueryMultiField_WithGlobalScope(t *testing.T) { - t.Skip("QueryMultiField SQL query needs 'k' parameter fix for sqlite-vec") - db, dbCleanup := testDB(t) defer dbCleanup() @@ -1217,7 +1257,37 @@ func TestClient_QueryMultiField_WithGlobalScope(t *testing.T) { require.NoError(t, err) // Should include both project-scoped (matching project) and global - assert.NotEmpty(t, results) + require.Len(t, results, 2) + assert.ElementsMatch(t, []float64{1, 2}, []float64{ + results[0].Metadata["sqlite_id"].(float64), results[1].Metadata["sqlite_id"].(float64), + }) +} + +// The query the de-duplication of new observations makes: the same project, a few results, and a note that +// says nearly the same thing must come back with a high similarity. +func TestClient_Query_ProjectFilterFindsANearDuplicate(t *testing.T) { + db, dbCleanup := testDB(t) + defer dbCleanup() + embedSvc, embedCleanup := testEmbeddingService(t) + defer embedCleanup() + client, err := NewClient(Config{DB: db}, embedSvc) + require.NoError(t, err) + + text := "The shipping rate cache lives for sixty minutes and is stored in Redis." + require.NoError(t, client.AddDocuments(context.Background(), []Document{ + {ID: "obs-1", Content: text, Metadata: map[string]any{"sqlite_id": int64(1), "doc_type": "observation", "project": "shop", "scope": "project"}}, + {ID: "obs-2", Content: text, Metadata: map[string]any{"sqlite_id": int64(2), "doc_type": "observation", "project": "elsewhere", "scope": "project"}}, + {ID: "obs-3", Content: "Parcel labels are printed as PDF.", Metadata: map[string]any{"sqlite_id": int64(3), "doc_type": "observation", "project": "shop", "scope": "project"}}, + })) + + results, err := client.Query(context.Background(), text, 3, BuildWhereFilter(DocTypeObservation, "shop")) + require.NoError(t, err) + require.NotEmpty(t, results) + assert.Equal(t, "obs-1", results[0].ID, "the same text in the same project is the best match") + assert.GreaterOrEqual(t, results[0].Similarity, 0.99) + for _, r := range results { + assert.NotEqual(t, "obs-2", r.ID, "another project's note is not a duplicate") + } } // ============================================================================= diff --git a/internal/worker/sdk/dedup.go b/internal/worker/sdk/dedup.go index 3fbd3f3..2323782 100644 --- a/internal/worker/sdk/dedup.go +++ b/internal/worker/sdk/dedup.go @@ -51,7 +51,7 @@ func (p *Processor) checkVectorDeduplication(ctx context.Context, obs *models.Pa for _, r := range results { if r.Similarity >= cfg.DeduplicationThreshold { obsID := extractObservationIDFromVectorDoc(r) - if obsID > 0 { + if obsID > 0 && p.canMergeInto(ctx, obsID) { return &DeduplicationResult{ ExistingID: obsID, Similarity: r.Similarity, @@ -64,6 +64,20 @@ func (p *Processor) checkVectorDeduplication(ctx context.Context, obs *models.Pa return &DeduplicationResult{Action: "insert"} } +// canMergeInto says whether a new observation may be merged into the existing one. A note a person superseded +// is hidden from sessions, so what was merged into it would be hidden as well: such a note is not a target, +// and neither is one that is gone. +func (p *Processor) canMergeInto(ctx context.Context, id int64) bool { + if p.observationStore == nil { + return true + } + existing, err := p.observationStore.GetObservationByID(ctx, id) + if err != nil { + return false + } + return existing != nil && !existing.IsSuperseded +} + // buildObservationSearchText creates searchable text from a parsed observation. func buildObservationSearchText(obs *models.ParsedObservation) string { var parts []string diff --git a/internal/worker/sdk/dedup_test.go b/internal/worker/sdk/dedup_test.go index 59edb06..7cd06d9 100644 --- a/internal/worker/sdk/dedup_test.go +++ b/internal/worker/sdk/dedup_test.go @@ -4,6 +4,7 @@ import ( "context" "testing" + "github.com/lukaszraczylo/claude-mnemonic/internal/db/gorm" "github.com/lukaszraczylo/claude-mnemonic/internal/vector/sqlitevec" "github.com/lukaszraczylo/claude-mnemonic/pkg/models" ) @@ -142,3 +143,35 @@ func TestCheckVectorDeduplication_EmptySearchText(t *testing.T) { t.Errorf("expected Action='insert' for empty observation, got %q", result.Action) } } + +func TestCanMergeInto(t *testing.T) { + ctx := context.Background() + dir := t.TempDir() + store, err := gorm.NewStore(gorm.Config{Path: dir + "/test.db", MaxConns: 1}) + if err != nil { + t.Fatal(err) + } + defer store.Close() + obsStore := gorm.NewObservationStore(store, nil, nil, nil) + id, _, err := obsStore.StoreObservation(ctx, "s1", "proj", &models.ParsedObservation{Type: models.ObsTypeDiscovery, Title: "A note"}, 1, 1) + if err != nil { + t.Fatal(err) + } + p := &Processor{observationStore: obsStore} + + if !p.canMergeInto(ctx, id) { + t.Error("a live observation is a merge target") + } + if p.canMergeInto(ctx, id+1000) { + t.Error("an observation that does not exist is not") + } + if err := store.DB.Exec(`UPDATE observations SET is_superseded = 1 WHERE id = ?`, id).Error; err != nil { + t.Fatal(err) + } + if p.canMergeInto(ctx, id) { + t.Error("a superseded observation is hidden from sessions, so nothing is merged into it") + } + if !(&Processor{}).canMergeInto(ctx, id) { + t.Error("without an observation store there is nothing to check against") + } +}