diff --git a/internal/vector/sqlitevec/client.go b/internal/vector/sqlitevec/client.go index a233c38..b88bcb5 100644 --- a/internal/vector/sqlitevec/client.go +++ b/internal/vector/sqlitevec/client.go @@ -139,9 +139,22 @@ func NewClient(cfg Config, embedSvc *embedding.Service) (*Client, error) { return c, nil } -// AddDocuments adds documents with their embeddings to the vector store. +// AddDocuments embeds the documents and inserts them. A document whose ID already exists is an error (the vec0 +// table does not replace on insert); use ReplaceDocuments to write the documents of a note that may have some. func (c *Client) AddDocuments(ctx context.Context, docs []Document) error { - if len(docs) == 0 { + return c.writeDocuments(ctx, docs, "", 0) +} + +// ReplaceDocuments writes the documents of one note (observation, summary, ...) in place of the ones it had: the old +// documents are deleted and the new ones inserted in one transaction, so an edit that removes a fact leaves no stale +// document behind and a failure leaves the old ones untouched. docs may be empty, which only deletes. +func (c *Client) ReplaceDocuments(ctx context.Context, docType string, sqliteID int64, docs []Document) error { + return c.writeDocuments(ctx, docs, docType, sqliteID) +} + +// writeDocuments inserts docs; when docType is not empty it first deletes every document of that note. +func (c *Client) writeDocuments(ctx context.Context, docs []Document, docType string, sqliteID int64) (err error) { + if len(docs) == 0 && docType == "" { return nil } @@ -154,9 +167,11 @@ func (c *Client) AddDocuments(ctx context.Context, docs []Document) error { // Compute embeddings OUTSIDE the write lock for better concurrency. // Embedding is ONNX inference (slow, mutex-protected internally) — holding // writeMu during inference blocks all concurrent writes and reads. - embeddings, err := c.embedSvc.EmbedBatchWithContext(ctx, texts) - if err != nil { - return fmt.Errorf("generate embeddings: %w", err) + var embeddings [][]float32 + if len(texts) > 0 { + if embeddings, err = c.embedSvc.EmbedBatchWithContext(ctx, texts); err != nil { + return fmt.Errorf("generate embeddings: %w", err) + } } // Acquire write lock for DB operations only @@ -165,7 +180,7 @@ func (c *Client) AddDocuments(ctx context.Context, docs []Document) error { // Insert into vectors table with model version tracking const insertQuery = ` - INSERT OR REPLACE INTO vectors (doc_id, embedding, sqlite_id, doc_type, field_type, project, scope, model_version) + INSERT INTO vectors (doc_id, embedding, sqlite_id, doc_type, field_type, project, scope, model_version) VALUES (?, ?, ?, ?, ?, ?, ?, ?) ` @@ -176,6 +191,8 @@ func (c *Client) AddDocuments(ctx context.Context, docs []Document) error { if err != nil { return fmt.Errorf("begin transaction: %w", err) } + // err is the named result, so every failed return below reaches this and the write lock is released. (A shadowed + // err in the loop once left the transaction open, and the database locked, after a failed insert.) defer func() { if err != nil { if rbErr := tx.Rollback(); rbErr != nil { @@ -185,6 +202,12 @@ func (c *Client) AddDocuments(ctx context.Context, docs []Document) error { } }() + if docType != "" { + if _, err = tx.ExecContext(ctx, `DELETE FROM vectors WHERE doc_type = ? AND sqlite_id = ?`, docType, sqliteID); err != nil { + return fmt.Errorf("delete old documents: %w", err) + } + } + stmt, err := tx.PrepareContext(ctx, insertQuery) if err != nil { return fmt.Errorf("prepare statement: %w", err) @@ -193,29 +216,28 @@ func (c *Client) AddDocuments(ctx context.Context, docs []Document) error { for i, doc := range docs { // Serialize embedding to blob format - embBlob, err := sqlite_vec.SerializeFloat32(embeddings[i]) - if err != nil { + var embBlob []byte + if embBlob, err = sqlite_vec.SerializeFloat32(embeddings[i]); err != nil { return fmt.Errorf("serialize embedding for %s: %w", doc.ID, err) } // Extract metadata - sqliteID, _ := doc.Metadata["sqlite_id"].(int64) - docType, _ := doc.Metadata["doc_type"].(string) + docSQLiteID, _ := doc.Metadata["sqlite_id"].(int64) + docDocType, _ := doc.Metadata["doc_type"].(string) fieldType, _ := doc.Metadata["field_type"].(string) project, _ := doc.Metadata["project"].(string) scope, _ := doc.Metadata["scope"].(string) - _, err = stmt.ExecContext(ctx, + if _, err = stmt.ExecContext(ctx, doc.ID, embBlob, - sqliteID, - docType, + docSQLiteID, + docDocType, fieldType, project, scope, modelVersion, - ) - if err != nil { + ); err != nil { return fmt.Errorf("insert document %s: %w", doc.ID, err) } } diff --git a/internal/vector/sqlitevec/client_test.go b/internal/vector/sqlitevec/client_test.go index 7b1ba53..4c3e126 100644 --- a/internal/vector/sqlitevec/client_test.go +++ b/internal/vector/sqlitevec/client_test.go @@ -2152,3 +2152,67 @@ func TestGetOrComputeEmbedding_ContextCancelDuringSingleflight(t *testing.T) { // Wait for first call to finish <-firstDone } + +func TestClient_AddDocuments_FailedBatchIsRolledBackAndDoesNotLockTheDatabase(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) + + doc := func(id string, sqliteID int64) Document { + return Document{ID: id, Content: "text of " + id, Metadata: map[string]any{ + "sqlite_id": sqliteID, "doc_type": "observation", "field_type": "narrative", "project": "alpha", "scope": "project", + }} + } + require.NoError(t, client.AddDocuments(context.Background(), []Document{doc("obs_1_narrative", 1)})) + + // The second document of the batch exists already, so the batch fails after the first was inserted. + err = client.AddDocuments(context.Background(), []Document{doc("obs_2_narrative", 2), doc("obs_1_narrative", 1)}) + require.Error(t, err) + + var n int + require.NoError(t, db.QueryRow("SELECT COUNT(*) FROM vectors WHERE doc_id = 'obs_2_narrative'").Scan(&n)) + assert.Zero(t, n, "nothing of the failed batch is kept") + + // With the transaction left open, this would wait for the lock until the context ends. + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + require.NoError(t, client.AddDocuments(ctx, []Document{doc("obs_3_narrative", 3)}), "a write after a failed one goes through") +} + +func TestClient_ReplaceDocuments_ReplacesOnlyTheDocumentsOfThatNote(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) + ctx := context.Background() + + doc := func(id string, sqliteID int64, docType, scope string) Document { + return Document{ID: id, Content: "text of " + id, Metadata: map[string]any{ + "sqlite_id": sqliteID, "doc_type": docType, "field_type": "narrative", "project": "alpha", "scope": scope, + }} + } + require.NoError(t, client.AddDocuments(ctx, []Document{ + doc("obs_1_narrative", 1, "observation", "global"), doc("obs_1_fact_0", 1, "observation", "global"), + doc("obs_2_narrative", 2, "observation", "global"), + doc("summary_1_request", 1, "session_summary", ""), // the same sqlite id, another kind of document + })) + + require.NoError(t, client.ReplaceDocuments(ctx, "observation", 1, []Document{doc("obs_1_narrative", 1, "observation", "project")})) + + count := func(where string, args ...any) (n int) { + require.NoError(t, db.QueryRow("SELECT COUNT(*) FROM vectors WHERE "+where, args...).Scan(&n)) + return n + } + assert.Equal(t, 1, count("doc_type = 'observation' AND sqlite_id = 1"), "the fact document is gone") + assert.Equal(t, 1, count("doc_id = 'obs_1_narrative' AND scope = 'project'"), "and the narrative carries the new scope") + assert.Equal(t, 1, count("doc_id = 'obs_2_narrative' AND scope = 'global'"), "another note is untouched") + assert.Equal(t, 1, count("doc_id = 'summary_1_request'"), "so is a summary with the same id") + + require.NoError(t, client.ReplaceDocuments(ctx, "observation", 1, nil), "no documents left: only the old ones are deleted") + assert.Zero(t, count("doc_type = 'observation' AND sqlite_id = 1")) +} diff --git a/internal/vector/sqlitevec/sync.go b/internal/vector/sqlitevec/sync.go index fe824de..3aaeee9 100644 --- a/internal/vector/sqlitevec/sync.go +++ b/internal/vector/sqlitevec/sync.go @@ -22,11 +22,8 @@ func NewSync(client *Client) *Sync { // SyncObservation syncs a single observation to the vector store. func (s *Sync) SyncObservation(ctx context.Context, obs *models.Observation) error { docs := s.formatObservationDocs(obs) - if len(docs) == 0 { - return nil - } - - if err := s.client.AddDocuments(ctx, docs); err != nil { + // Replace, not add: a note that is synced again (after an edit, a re-scope or a re-import) already has documents. + if err := s.client.ReplaceDocuments(ctx, "observation", obs.ID, docs); err != nil { return fmt.Errorf("add observation docs: %w", err) } @@ -102,11 +99,7 @@ func (s *Sync) formatObservationDocs(obs *models.Observation) []Document { // SyncSummary syncs a single session summary to the vector store. func (s *Sync) SyncSummary(ctx context.Context, summary *models.SessionSummary) error { docs := s.formatSummaryDocs(summary) - if len(docs) == 0 { - return nil - } - - if err := s.client.AddDocuments(ctx, docs); err != nil { + if err := s.client.ReplaceDocuments(ctx, "session_summary", summary.ID, docs); err != nil { return fmt.Errorf("add summary docs: %w", err) } diff --git a/internal/vector/sqlitevec/sync_test.go b/internal/vector/sqlitevec/sync_test.go index 5a9ae63..2447483 100644 --- a/internal/vector/sqlitevec/sync_test.go +++ b/internal/vector/sqlitevec/sync_test.go @@ -346,3 +346,57 @@ func TestSync_FormatObservationDocs_EmptyScope(t *testing.T) { assert.Len(t, docs, 1) assert.Equal(t, "project", docs[0].Metadata["scope"]) } + +func vectorCount(t *testing.T, c *Client, where string, args ...any) int { + t.Helper() + var n int + require.NoError(t, c.db.QueryRow("SELECT COUNT(*) FROM vectors WHERE "+where, args...).Scan(&n)) + return n +} + +func TestSync_SyncObservation_AgainReplacesTheDocumentsOfTheNote(t *testing.T) { + client, cleanup := testClient(t) + defer cleanup() + sync := NewSync(client) + ctx := context.Background() + + obs := &models.Observation{ + ID: 7, SDKSessionID: "s", Project: "alpha", Type: models.ObsTypeDiscovery, Scope: models.ScopeGlobal, + Narrative: sql.NullString{String: "The queue keeps one consumer per tenant.", Valid: true}, + Facts: []string{"one consumer per tenant", "retries are capped", "dead letters go to a second queue"}, + } + other := &models.Observation{ + ID: 8, SDKSessionID: "s", Project: "alpha", Type: models.ObsTypeDiscovery, Scope: models.ScopeProject, + Narrative: sql.NullString{String: "Releases are cut on Thursdays.", Valid: true}, + } + require.NoError(t, sync.SyncObservation(ctx, obs)) + require.NoError(t, sync.SyncObservation(ctx, other)) + require.Equal(t, 4, vectorCount(t, client, "doc_type = 'observation' AND sqlite_id = ?", 7)) + + // The note was edited: one fact fewer, and it is a project note now. + obs.Facts = obs.Facts[:2] + obs.Scope = models.ScopeProject + require.NoError(t, sync.SyncObservation(ctx, obs), "syncing a note that already has documents is not an error") + + assert.Equal(t, 3, vectorCount(t, client, "doc_type = 'observation' AND sqlite_id = ?", 7), "narrative and two facts; the dropped fact's document is gone") + assert.Equal(t, 0, vectorCount(t, client, "doc_id = ?", "obs_7_fact_2")) + assert.Equal(t, 3, vectorCount(t, client, "sqlite_id = ? AND scope = 'project'", 7), "the new scope is what the vectors carry") + assert.Equal(t, 1, vectorCount(t, client, "sqlite_id = ?", 8), "other notes are not touched") +} + +func TestSync_SyncSummary_AgainReplacesTheDocumentsOfTheSummary(t *testing.T) { + client, cleanup := testClient(t) + defer cleanup() + sync := NewSync(client) + ctx := context.Background() + + summary := &models.SessionSummary{ + ID: 3, SDKSessionID: "s", Project: "alpha", + Request: sql.NullString{String: "Add the queue", Valid: true}, Learned: sql.NullString{String: "One consumer per tenant", Valid: true}, + } + require.NoError(t, sync.SyncSummary(ctx, summary)) + summary.Learned = sql.NullString{} + require.NoError(t, sync.SyncSummary(ctx, summary)) + + assert.Equal(t, 1, vectorCount(t, client, "doc_type = 'session_summary' AND sqlite_id = ?", 3), "the field that was emptied has no document any more") +}