mirror of
https://github.com/lukaszraczylo/claude-mnemonic.git
synced 2026-10-12 04:55:18 +00:00
fix(vector): answer project-filtered queries with two searches (#70)
A project filter was '(project = ? OR scope = global)', which vec0 rejects in a nearest-neighbour query ('A LIMIT or k = ? constraint is required'). The only caller that passes a project is the de-duplication of new observations, which treats an error as 'insert', so semantic de-duplication never took effect.
Run one search for the project and one for global documents and merge them by distance without repeats, cut to the limit. QueryMultiField asks for enough field documents and keeps the best per observation; its two skipped tests are enabled. Because de-duplication now runs, canMergeInto skips a superseded or deleted note as a merge target.
Co-authored-by: Claude Sonnet 5.5 <[email protected]>
This commit is contained in:
Holger Reinhardt
co-authored by
Claude Sonnet 5.5
Claude Sonnet 5.5
parent
055103cb3b
commit
9b6e5f7648
@@ -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.
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user