mirror of
https://github.com/lukaszraczylo/claude-mnemonic.git
synced 2026-10-12 04:55:18 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
8c2c69582a | ||
|
|
fbe938fef2 | ||
|
|
64c2ace22e | ||
|
|
f2886bbedb | ||
|
|
9b6e5f7648 | ||
|
|
055103cb3b | ||
|
|
67de0f0c84 | ||
|
|
eba3136fff |
@@ -252,9 +252,6 @@ func (s *ObservationStore) UpdateObservation(ctx context.Context, id int64, upda
|
|||||||
updates["scope"] = sql.NullString{String: *update.Scope, Valid: true}
|
updates["scope"] = sql.NullString{String: *update.Scope, Valid: true}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Add updated_at timestamp
|
|
||||||
updates["updated_at_epoch"] = sql.NullInt64{Int64: time.Now().Unix(), Valid: true}
|
|
||||||
|
|
||||||
if len(updates) == 0 {
|
if len(updates) == 0 {
|
||||||
// Nothing to update, just return existing observation
|
// Nothing to update, just return existing observation
|
||||||
return toModelObservation(&dbObs), nil
|
return toModelObservation(&dbObs), nil
|
||||||
|
|||||||
@@ -0,0 +1,96 @@
|
|||||||
|
//go:build fts5
|
||||||
|
|
||||||
|
package gorm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/stretchr/testify/assert"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
|
||||||
|
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
||||||
|
)
|
||||||
|
|
||||||
|
func storeForUpdate(t *testing.T) (*ObservationStore, int64, func()) {
|
||||||
|
t.Helper()
|
||||||
|
obsStore, store, cleanup := testObservationStore(t)
|
||||||
|
ctx := context.Background()
|
||||||
|
sessionID, err := NewSessionStore(store).CreateSDKSession(ctx, "claude-1", "test-project", "")
|
||||||
|
require.NoError(t, err)
|
||||||
|
id, _, err := obsStore.StoreObservation(ctx, "claude-1", "test-project", &models.ParsedObservation{
|
||||||
|
Type: models.ObsTypeDiscovery, Title: "Original title", Subtitle: "Original subtitle", Narrative: "Original narrative",
|
||||||
|
Facts: []string{"fact one"}, Concepts: []string{"cache"}, FilesRead: []string{"a.go"},
|
||||||
|
}, int(sessionID), 1)
|
||||||
|
require.NoError(t, err)
|
||||||
|
return obsStore, id, cleanup
|
||||||
|
}
|
||||||
|
|
||||||
|
func ptr[T any](v T) *T { return &v }
|
||||||
|
|
||||||
|
func TestUpdateObservation_ChangesOnlyTheGivenFields(t *testing.T) {
|
||||||
|
obsStore, id, cleanup := storeForUpdate(t)
|
||||||
|
defer cleanup()
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
got, err := obsStore.UpdateObservation(ctx, id, &ObservationUpdate{Title: ptr("New title"), Narrative: ptr("New narrative")})
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, "New title", got.Title.String)
|
||||||
|
assert.Equal(t, "New narrative", got.Narrative.String)
|
||||||
|
assert.Equal(t, "Original subtitle", got.Subtitle.String, "a field that was not given stays")
|
||||||
|
assert.Equal(t, []string{"fact one"}, []string(got.Facts))
|
||||||
|
|
||||||
|
stored, err := obsStore.GetObservationByID(ctx, id)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, "New title", stored.Title.String, "and it is stored, not only returned")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestUpdateObservation_ReplacesListsAndScope(t *testing.T) {
|
||||||
|
obsStore, id, cleanup := storeForUpdate(t)
|
||||||
|
defer cleanup()
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
got, err := obsStore.UpdateObservation(ctx, id, &ObservationUpdate{
|
||||||
|
Facts: &[]string{"a", "b"}, Concepts: &[]string{}, FilesRead: &[]string{"x.go"}, FilesModified: &[]string{"y.go"}, Scope: ptr("global"),
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, []string{"a", "b"}, []string(got.Facts))
|
||||||
|
assert.Empty(t, got.Concepts, "an empty list clears the field")
|
||||||
|
assert.Equal(t, []string{"x.go"}, []string(got.FilesRead))
|
||||||
|
assert.Equal(t, []string{"y.go"}, []string(got.FilesModified))
|
||||||
|
assert.Equal(t, models.ScopeGlobal, got.Scope)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestUpdateObservation_NothingToChangeReturnsTheObservation(t *testing.T) {
|
||||||
|
obsStore, id, cleanup := storeForUpdate(t)
|
||||||
|
defer cleanup()
|
||||||
|
|
||||||
|
got, err := obsStore.UpdateObservation(context.Background(), id, &ObservationUpdate{})
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, "Original title", got.Title.String)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestUpdateObservation_Refusals(t *testing.T) {
|
||||||
|
obsStore, id, cleanup := storeForUpdate(t)
|
||||||
|
defer cleanup()
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
_, err := obsStore.UpdateObservation(ctx, id, nil)
|
||||||
|
assert.Error(t, err, "a nil update")
|
||||||
|
_, err = obsStore.UpdateObservation(ctx, 999999, &ObservationUpdate{Title: ptr("x")})
|
||||||
|
require.Error(t, err)
|
||||||
|
assert.Contains(t, err.Error(), "not found")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestUpdateObservation_KeepsTheSearchIndexInStep(t *testing.T) {
|
||||||
|
obsStore, id, cleanup := storeForUpdate(t)
|
||||||
|
defer cleanup()
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
_, err := obsStore.UpdateObservation(ctx, id, &ObservationUpdate{Title: ptr("Zebrafish enclosure")})
|
||||||
|
require.NoError(t, err)
|
||||||
|
var n int64
|
||||||
|
require.NoError(t, obsStore.db.WithContext(ctx).Raw(
|
||||||
|
`SELECT COUNT(*) FROM observations_fts WHERE observations_fts MATCH 'zebrafish'`).Scan(&n).Error)
|
||||||
|
assert.EqualValues(t, 1, n, "the full-text index sees the new title")
|
||||||
|
}
|
||||||
@@ -0,0 +1,38 @@
|
|||||||
|
package mcp
|
||||||
|
|
||||||
|
import (
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/stretchr/testify/assert"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestIdleTimeoutFromEnv(t *testing.T) {
|
||||||
|
// The idle shutdown is off unless asked for; anything that is not a positive duration leaves it off.
|
||||||
|
for value, want := range map[string]time.Duration{
|
||||||
|
"": 0,
|
||||||
|
"0": 0,
|
||||||
|
"abc": 0,
|
||||||
|
"-5m": 0,
|
||||||
|
"45": 0, // a bare number is not a duration
|
||||||
|
"90m": 90 * time.Minute,
|
||||||
|
"2h": 2 * time.Hour,
|
||||||
|
"30s": 30 * time.Second,
|
||||||
|
" 2h ": 0, // not trimmed: the operator wrote it, so it is wrong rather than guessed
|
||||||
|
} {
|
||||||
|
t.Run(value, func(t *testing.T) {
|
||||||
|
t.Setenv(IdleTimeoutEnv, value)
|
||||||
|
assert.Equal(t, want, idleTimeoutFromEnv())
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestNewServer_HasNoIdleTimeoutByDefault(t *testing.T) {
|
||||||
|
t.Setenv(IdleTimeoutEnv, "")
|
||||||
|
assert.Zero(t, NewServer(nil, "http://localhost:1", "p", "v").idleTimeout)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestNewServer_TakesTheIdleTimeoutFromTheEnvironment(t *testing.T) {
|
||||||
|
t.Setenv(IdleTimeoutEnv, "3h")
|
||||||
|
assert.Equal(t, 3*time.Hour, NewServer(nil, "http://localhost:1", "p", "v").idleTimeout)
|
||||||
|
}
|
||||||
@@ -0,0 +1,97 @@
|
|||||||
|
//go:build unix
|
||||||
|
|
||||||
|
package mcp
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"os"
|
||||||
|
"sync"
|
||||||
|
"syscall"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/stretchr/testify/assert"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
)
|
||||||
|
|
||||||
|
// blockingPipe returns the two ends of a pipe in blocking mode, the way an inherited stdin is: a Read on it cannot be
|
||||||
|
// woken by closing the file (the standard pipe from os.Pipe is pollable and can, which hid this).
|
||||||
|
func blockingPipe(t *testing.T) (r, w *os.File) {
|
||||||
|
t.Helper()
|
||||||
|
var fds [2]int
|
||||||
|
require.NoError(t, syscall.Pipe(fds[:]))
|
||||||
|
r = os.NewFile(uintptr(fds[0]), "stdin-like")
|
||||||
|
w = os.NewFile(uintptr(fds[1]), "stdin-writer")
|
||||||
|
t.Cleanup(func() { _ = w.Close(); _ = r.Close() })
|
||||||
|
return r, w
|
||||||
|
}
|
||||||
|
|
||||||
|
type lockedBuffer struct {
|
||||||
|
buf bytes.Buffer
|
||||||
|
mu sync.Mutex
|
||||||
|
}
|
||||||
|
|
||||||
|
func (b *lockedBuffer) Write(p []byte) (int, error) {
|
||||||
|
b.mu.Lock()
|
||||||
|
defer b.mu.Unlock()
|
||||||
|
return b.buf.Write(p)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (b *lockedBuffer) String() string {
|
||||||
|
b.mu.Lock()
|
||||||
|
defer b.mu.Unlock()
|
||||||
|
return b.buf.String()
|
||||||
|
}
|
||||||
|
|
||||||
|
// The idle timeout used to close stdin and hope the read loop ended. On a real (blocking) stdin the read stays blocked,
|
||||||
|
// so the server logged "shutting down" and then lived on with a closed stdin; the next message from the client woke the
|
||||||
|
// read, which failed with "file already closed" and killed the connection mid-call (seen in Claude Desktop after the
|
||||||
|
// connector sat unused for more than 30 minutes). An idle timeout that is on must end Run for real.
|
||||||
|
func TestRun_IdleTimeoutEndsTheServerCleanly(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
stdin, _ := blockingPipe(t)
|
||||||
|
server := &Server{stdin: stdin, stdout: &lockedBuffer{}, idleTimeout: 50 * time.Millisecond, monitorInterval: 10 * time.Millisecond}
|
||||||
|
|
||||||
|
done := make(chan error, 1)
|
||||||
|
go func() { done <- server.Run(context.Background()) }()
|
||||||
|
|
||||||
|
select {
|
||||||
|
case err := <-done:
|
||||||
|
require.NoError(t, err, "an idle shutdown is a clean exit, not an error")
|
||||||
|
case <-time.After(3 * time.Second):
|
||||||
|
t.Fatal("Run did not return after the idle timeout: stdin was closed but the blocked read never ended")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// By default there is no idle timeout: a stdio server lives as long as its client, and a client that talks to it after a
|
||||||
|
// long pause must get an answer.
|
||||||
|
func TestRun_WithoutAnIdleTimeoutTheServerKeepsAnsweringAfterALongPause(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
stdin, stdinWriter := blockingPipe(t)
|
||||||
|
stdout := &lockedBuffer{}
|
||||||
|
server := &Server{stdin: stdin, stdout: stdout, monitorInterval: 5 * time.Millisecond}
|
||||||
|
|
||||||
|
done := make(chan error, 1)
|
||||||
|
go func() { done <- server.Run(context.Background()) }()
|
||||||
|
|
||||||
|
time.Sleep(300 * time.Millisecond) // many monitor ticks with no traffic
|
||||||
|
select {
|
||||||
|
case err := <-done:
|
||||||
|
t.Fatalf("Run ended during a pause with no traffic: %v", err)
|
||||||
|
default:
|
||||||
|
}
|
||||||
|
|
||||||
|
_, err := stdinWriter.WriteString(`{"jsonrpc":"2.0","id":7,"method":"tools/list"}` + "\n")
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
assert.Eventually(t, func() bool { return bytes.Contains([]byte(stdout.String()), []byte(`"id":7`)) }, 3*time.Second, 10*time.Millisecond,
|
||||||
|
"the request after the pause is answered")
|
||||||
|
select {
|
||||||
|
case err := <-done:
|
||||||
|
t.Fatalf("Run ended after answering: %v", err)
|
||||||
|
default:
|
||||||
|
}
|
||||||
|
}
|
||||||
+41
-8
@@ -30,6 +30,11 @@ type Server struct {
|
|||||||
version string
|
version string
|
||||||
writeMu sync.Mutex
|
writeMu sync.Mutex
|
||||||
lastActivity atomic.Int64
|
lastActivity atomic.Int64
|
||||||
|
// idleTimeout ends the server after this long without a message; 0 (the default) never does. A stdio server lives
|
||||||
|
// as long as its client: ending it while the client is still connected makes the client's next call fail.
|
||||||
|
idleTimeout time.Duration
|
||||||
|
// monitorInterval is how often the parent and idle checks run (30s when 0).
|
||||||
|
monitorInterval time.Duration
|
||||||
}
|
}
|
||||||
|
|
||||||
// NewServer creates a new MCP server that proxies to the worker HTTP API.
|
// NewServer creates a new MCP server that proxies to the worker HTTP API.
|
||||||
@@ -41,9 +46,24 @@ func NewServer(client *http.Client, workerURL, project, version string) *Server
|
|||||||
version: version,
|
version: version,
|
||||||
stdin: os.Stdin,
|
stdin: os.Stdin,
|
||||||
stdout: os.Stdout,
|
stdout: os.Stdout,
|
||||||
|
|
||||||
|
idleTimeout: idleTimeoutFromEnv(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// IdleTimeoutEnv names the environment variable that turns the idle shutdown on, as a Go duration such as "2h".
|
||||||
|
// It is off by default.
|
||||||
|
const IdleTimeoutEnv = "CLAUDE_MNEMONIC_MCP_IDLE_TIMEOUT"
|
||||||
|
|
||||||
|
// idleTimeoutFromEnv reads IdleTimeoutEnv; an unset, invalid or non-positive value means no idle shutdown.
|
||||||
|
func idleTimeoutFromEnv() time.Duration {
|
||||||
|
d, err := time.ParseDuration(os.Getenv(IdleTimeoutEnv))
|
||||||
|
if err != nil || d <= 0 {
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
return d
|
||||||
|
}
|
||||||
|
|
||||||
// Request represents a JSON-RPC request.
|
// Request represents a JSON-RPC request.
|
||||||
type Request struct {
|
type Request struct {
|
||||||
JSONRPC string `json:"jsonrpc"`
|
JSONRPC string `json:"jsonrpc"`
|
||||||
@@ -82,6 +102,13 @@ type Tool struct {
|
|||||||
|
|
||||||
// Run starts the MCP server loop.
|
// Run starts the MCP server loop.
|
||||||
func (s *Server) Run(ctx context.Context) error {
|
func (s *Server) Run(ctx context.Context) error {
|
||||||
|
// An idle shutdown cancels this context. Closing stdin does not do it: on a real (blocking) stdin the read stays
|
||||||
|
// blocked, the server lived on with a closed stdin, and the client's next message woke the read with "file already
|
||||||
|
// closed", which killed the connection mid-call.
|
||||||
|
ctx, cancel := context.WithCancel(ctx)
|
||||||
|
defer cancel()
|
||||||
|
var idleStopped atomic.Bool
|
||||||
|
|
||||||
scanner := bufio.NewScanner(s.stdin)
|
scanner := bufio.NewScanner(s.stdin)
|
||||||
buf := make([]byte, 0, 64*1024)
|
buf := make([]byte, 0, 64*1024)
|
||||||
scanner.Buffer(buf, 1024*1024) // 1MB max message size
|
scanner.Buffer(buf, 1024*1024) // 1MB max message size
|
||||||
@@ -100,11 +127,15 @@ func (s *Server) Run(ctx context.Context) error {
|
|||||||
scanErr <- scanner.Err()
|
scanErr <- scanner.Err()
|
||||||
}()
|
}()
|
||||||
|
|
||||||
// Monitor parent process liveness and idle timeout.
|
// Monitor parent process liveness and, when one is set, the idle timeout.
|
||||||
// If the parent dies (ppid changes) or no messages arrive for 30 minutes, shut down.
|
// If the parent dies (ppid changes) or no messages arrive for idleTimeout, shut down.
|
||||||
parentPID := os.Getppid()
|
parentPID := os.Getppid()
|
||||||
go func() {
|
go func() {
|
||||||
ticker := time.NewTicker(30 * time.Second)
|
interval := s.monitorInterval
|
||||||
|
if interval <= 0 {
|
||||||
|
interval = 30 * time.Second
|
||||||
|
}
|
||||||
|
ticker := time.NewTicker(interval)
|
||||||
defer ticker.Stop()
|
defer ticker.Stop()
|
||||||
for {
|
for {
|
||||||
select {
|
select {
|
||||||
@@ -118,11 +149,10 @@ func (s *Server) Run(ctx context.Context) error {
|
|||||||
}
|
}
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
if time.Since(time.Unix(s.lastActivity.Load(), 0)) > 30*time.Minute {
|
if s.idleTimeout > 0 && time.Since(time.Unix(s.lastActivity.Load(), 0)) > s.idleTimeout {
|
||||||
log.Info().Msg("MCP server idle timeout (30m), shutting down")
|
log.Info().Dur("idle", s.idleTimeout).Msg("MCP server idle timeout, shutting down")
|
||||||
if closer, ok := s.stdin.(io.Closer); ok {
|
idleStopped.Store(true)
|
||||||
_ = closer.Close()
|
cancel()
|
||||||
}
|
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -140,6 +170,9 @@ func (s *Server) Run(ctx context.Context) error {
|
|||||||
case <-ctx.Done():
|
case <-ctx.Done():
|
||||||
// Drain in-flight requests before returning.
|
// Drain in-flight requests before returning.
|
||||||
wg.Wait()
|
wg.Wait()
|
||||||
|
if idleStopped.Load() {
|
||||||
|
return nil // an idle shutdown the operator asked for is a clean exit
|
||||||
|
}
|
||||||
return ctx.Err()
|
return ctx.Err()
|
||||||
case line, ok := <-lines:
|
case line, ok := <-lines:
|
||||||
if !ok {
|
if !ok {
|
||||||
|
|||||||
+125
-107
@@ -5,6 +5,7 @@ import (
|
|||||||
"context"
|
"context"
|
||||||
"database/sql"
|
"database/sql"
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"sort"
|
||||||
"strconv"
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
"sync"
|
"sync"
|
||||||
@@ -139,9 +140,22 @@ func NewClient(cfg Config, embedSvc *embedding.Service) (*Client, error) {
|
|||||||
return c, nil
|
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 {
|
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
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -154,9 +168,11 @@ func (c *Client) AddDocuments(ctx context.Context, docs []Document) error {
|
|||||||
// Compute embeddings OUTSIDE the write lock for better concurrency.
|
// Compute embeddings OUTSIDE the write lock for better concurrency.
|
||||||
// Embedding is ONNX inference (slow, mutex-protected internally) — holding
|
// Embedding is ONNX inference (slow, mutex-protected internally) — holding
|
||||||
// writeMu during inference blocks all concurrent writes and reads.
|
// writeMu during inference blocks all concurrent writes and reads.
|
||||||
embeddings, err := c.embedSvc.EmbedBatchWithContext(ctx, texts)
|
var embeddings [][]float32
|
||||||
if err != nil {
|
if len(texts) > 0 {
|
||||||
return fmt.Errorf("generate embeddings: %w", err)
|
if embeddings, err = c.embedSvc.EmbedBatchWithContext(ctx, texts); err != nil {
|
||||||
|
return fmt.Errorf("generate embeddings: %w", err)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Acquire write lock for DB operations only
|
// Acquire write lock for DB operations only
|
||||||
@@ -165,7 +181,7 @@ func (c *Client) AddDocuments(ctx context.Context, docs []Document) error {
|
|||||||
|
|
||||||
// Insert into vectors table with model version tracking
|
// Insert into vectors table with model version tracking
|
||||||
const insertQuery = `
|
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 (?, ?, ?, ?, ?, ?, ?, ?)
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||||||
`
|
`
|
||||||
|
|
||||||
@@ -176,6 +192,8 @@ func (c *Client) AddDocuments(ctx context.Context, docs []Document) error {
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return fmt.Errorf("begin transaction: %w", err)
|
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() {
|
defer func() {
|
||||||
if err != nil {
|
if err != nil {
|
||||||
if rbErr := tx.Rollback(); rbErr != nil {
|
if rbErr := tx.Rollback(); rbErr != nil {
|
||||||
@@ -185,6 +203,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)
|
stmt, err := tx.PrepareContext(ctx, insertQuery)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return fmt.Errorf("prepare statement: %w", err)
|
return fmt.Errorf("prepare statement: %w", err)
|
||||||
@@ -193,29 +217,28 @@ func (c *Client) AddDocuments(ctx context.Context, docs []Document) error {
|
|||||||
|
|
||||||
for i, doc := range docs {
|
for i, doc := range docs {
|
||||||
// Serialize embedding to blob format
|
// Serialize embedding to blob format
|
||||||
embBlob, err := sqlite_vec.SerializeFloat32(embeddings[i])
|
var embBlob []byte
|
||||||
if err != nil {
|
if embBlob, err = sqlite_vec.SerializeFloat32(embeddings[i]); err != nil {
|
||||||
return fmt.Errorf("serialize embedding for %s: %w", doc.ID, err)
|
return fmt.Errorf("serialize embedding for %s: %w", doc.ID, err)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Extract metadata
|
// Extract metadata
|
||||||
sqliteID, _ := doc.Metadata["sqlite_id"].(int64)
|
docSQLiteID, _ := doc.Metadata["sqlite_id"].(int64)
|
||||||
docType, _ := doc.Metadata["doc_type"].(string)
|
docDocType, _ := doc.Metadata["doc_type"].(string)
|
||||||
fieldType, _ := doc.Metadata["field_type"].(string)
|
fieldType, _ := doc.Metadata["field_type"].(string)
|
||||||
project, _ := doc.Metadata["project"].(string)
|
project, _ := doc.Metadata["project"].(string)
|
||||||
scope, _ := doc.Metadata["scope"].(string)
|
scope, _ := doc.Metadata["scope"].(string)
|
||||||
|
|
||||||
_, err = stmt.ExecContext(ctx,
|
if _, err = stmt.ExecContext(ctx,
|
||||||
doc.ID,
|
doc.ID,
|
||||||
embBlob,
|
embBlob,
|
||||||
sqliteID,
|
docSQLiteID,
|
||||||
docType,
|
docDocType,
|
||||||
fieldType,
|
fieldType,
|
||||||
project,
|
project,
|
||||||
scope,
|
scope,
|
||||||
modelVersion,
|
modelVersion,
|
||||||
)
|
); err != nil {
|
||||||
if err != nil {
|
|
||||||
return fmt.Errorf("insert document %s: %w", doc.ID, err)
|
return fmt.Errorf("insert document %s: %w", doc.ID, err)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -292,67 +315,13 @@ func (c *Client) Query(ctx context.Context, query string, limit int, where map[s
|
|||||||
}
|
}
|
||||||
defer c.readMu.RUnlock()
|
defer c.readMu.RUnlock()
|
||||||
|
|
||||||
// Build query with filters
|
// vec0 answers equality filters on metadata columns, but not an OR between them, so a project filter
|
||||||
// vec0 supports WHERE clauses on metadata columns
|
// (the project's own documents, plus global ones) is asked as two searches and merged.
|
||||||
args := []any{queryBlob}
|
docType, _ := where["doc_type"].(string)
|
||||||
|
project, _ := where["project"].(string)
|
||||||
sqlQuery := `
|
results, err := c.nearest(ctx, queryBlob, docType, project, limit)
|
||||||
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...)
|
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("query vectors: %w", err)
|
return nil, 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)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// Cache the results
|
// Cache the results
|
||||||
@@ -507,60 +476,109 @@ func (c *Client) QueryMultiField(ctx context.Context, query string, limit int, d
|
|||||||
}
|
}
|
||||||
defer c.readMu.RUnlock()
|
defer c.readMu.RUnlock()
|
||||||
|
|
||||||
// Query with field type aggregation - get best match per document
|
// Each observation has several vector documents (one per field). Ask for enough of them that `limit`
|
||||||
sqlQuery := `
|
// distinct observations are likely to be among them, then keep the best match per observation.
|
||||||
WITH ranked_results AS (
|
found, err := c.nearest(ctx, queryBlob, docType, project, limit*multiFieldFetchFactor)
|
||||||
SELECT
|
if err != nil {
|
||||||
doc_id,
|
return nil, err
|
||||||
distance,
|
}
|
||||||
sqlite_id,
|
results := make([]QueryResult, 0, limit)
|
||||||
doc_type,
|
seen := make(map[float64]bool, limit)
|
||||||
field_type,
|
for _, r := range found { // best first
|
||||||
project,
|
id, _ := r.Metadata["sqlite_id"].(float64)
|
||||||
scope,
|
if seen[id] {
|
||||||
ROW_NUMBER() OVER (PARTITION BY sqlite_id ORDER BY distance ASC) as rn
|
continue
|
||||||
FROM vectors
|
}
|
||||||
WHERE embedding MATCH ?
|
seen[id] = true
|
||||||
AND doc_type = ?
|
results = append(results, r)
|
||||||
AND (project = ? OR scope = 'global')
|
if len(results) == limit {
|
||||||
)
|
break
|
||||||
SELECT doc_id, distance, sqlite_id, doc_type, field_type, project, scope
|
}
|
||||||
FROM ranked_results
|
}
|
||||||
WHERE rn = 1
|
return results, nil
|
||||||
ORDER BY distance
|
}
|
||||||
LIMIT ?
|
|
||||||
`
|
|
||||||
|
|
||||||
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) {
|
||||||
|
// #nosec G202 -- where is built from constant clauses with "?" placeholders only (see nearest); every value goes via args
|
||||||
|
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 {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("query vectors: %w", err)
|
return nil, fmt.Errorf("query vectors: %w", err)
|
||||||
}
|
}
|
||||||
defer rows.Close()
|
defer rows.Close()
|
||||||
|
|
||||||
// Pre-allocate with limit to avoid repeated slice growth
|
var results []QueryResult
|
||||||
results := make([]QueryResult, 0, limit)
|
|
||||||
for rows.Next() {
|
for rows.Next() {
|
||||||
var r QueryResult
|
var r QueryResult
|
||||||
var sqliteID int64
|
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)
|
return nil, fmt.Errorf("scan row: %w", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
r.Similarity = DistanceToSimilarity(r.Distance)
|
r.Similarity = DistanceToSimilarity(r.Distance)
|
||||||
r.Metadata = map[string]any{
|
r.Metadata = map[string]any{
|
||||||
"sqlite_id": float64(sqliteID),
|
"sqlite_id": float64(sqliteID), // Keep as float64 for compatibility
|
||||||
"doc_type": docTypeVal.String,
|
"doc_type": docType.String,
|
||||||
"field_type": fieldType.String,
|
"field_type": fieldType.String,
|
||||||
"project": projectVal.String,
|
"project": project.String,
|
||||||
"scope": scope.String,
|
"scope": scope.String,
|
||||||
}
|
}
|
||||||
|
|
||||||
results = append(results, r)
|
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.
|
// 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)
|
results, err := client.Query(context.Background(), "authentication security", 10, nil)
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
assert.NotEmpty(t, results, "Should find some results")
|
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) {
|
func TestClient_IsConnected(t *testing.T) {
|
||||||
@@ -1105,8 +1149,6 @@ func TestClient_QueryBatch_WithContextCancellation(t *testing.T) {
|
|||||||
// =============================================================================
|
// =============================================================================
|
||||||
|
|
||||||
func TestClient_QueryMultiField_Basic(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)
|
db, dbCleanup := testDB(t)
|
||||||
defer dbCleanup()
|
defer dbCleanup()
|
||||||
|
|
||||||
@@ -1173,8 +1215,6 @@ func TestClient_QueryMultiField_Basic(t *testing.T) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func TestClient_QueryMultiField_WithGlobalScope(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)
|
db, dbCleanup := testDB(t)
|
||||||
defer dbCleanup()
|
defer dbCleanup()
|
||||||
|
|
||||||
@@ -1217,7 +1257,37 @@ func TestClient_QueryMultiField_WithGlobalScope(t *testing.T) {
|
|||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
|
|
||||||
// Should include both project-scoped (matching project) and global
|
// 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")
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// =============================================================================
|
// =============================================================================
|
||||||
@@ -2152,3 +2222,67 @@ func TestGetOrComputeEmbedding_ContextCancelDuringSingleflight(t *testing.T) {
|
|||||||
// Wait for first call to finish
|
// Wait for first call to finish
|
||||||
<-firstDone
|
<-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"))
|
||||||
|
}
|
||||||
|
|||||||
@@ -22,11 +22,8 @@ func NewSync(client *Client) *Sync {
|
|||||||
// SyncObservation syncs a single observation to the vector store.
|
// SyncObservation syncs a single observation to the vector store.
|
||||||
func (s *Sync) SyncObservation(ctx context.Context, obs *models.Observation) error {
|
func (s *Sync) SyncObservation(ctx context.Context, obs *models.Observation) error {
|
||||||
docs := s.formatObservationDocs(obs)
|
docs := s.formatObservationDocs(obs)
|
||||||
if len(docs) == 0 {
|
// Replace, not add: a note that is synced again (after an edit, a re-scope or a re-import) already has documents.
|
||||||
return nil
|
if err := s.client.ReplaceDocuments(ctx, "observation", obs.ID, docs); err != nil {
|
||||||
}
|
|
||||||
|
|
||||||
if err := s.client.AddDocuments(ctx, docs); err != nil {
|
|
||||||
return fmt.Errorf("add observation docs: %w", err)
|
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.
|
// SyncSummary syncs a single session summary to the vector store.
|
||||||
func (s *Sync) SyncSummary(ctx context.Context, summary *models.SessionSummary) error {
|
func (s *Sync) SyncSummary(ctx context.Context, summary *models.SessionSummary) error {
|
||||||
docs := s.formatSummaryDocs(summary)
|
docs := s.formatSummaryDocs(summary)
|
||||||
if len(docs) == 0 {
|
if err := s.client.ReplaceDocuments(ctx, "session_summary", summary.ID, docs); err != nil {
|
||||||
return nil
|
|
||||||
}
|
|
||||||
|
|
||||||
if err := s.client.AddDocuments(ctx, docs); err != nil {
|
|
||||||
return fmt.Errorf("add summary docs: %w", err)
|
return fmt.Errorf("add summary docs: %w", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -346,3 +346,57 @@ func TestSync_FormatObservationDocs_EmptyScope(t *testing.T) {
|
|||||||
assert.Len(t, docs, 1)
|
assert.Len(t, docs, 1)
|
||||||
assert.Equal(t, "project", docs[0].Metadata["scope"])
|
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")
|
||||||
|
}
|
||||||
|
|||||||
@@ -0,0 +1,59 @@
|
|||||||
|
package worker
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/stretchr/testify/assert"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
)
|
||||||
|
|
||||||
|
func doUpdateRequest(t *testing.T, svc *Service, method, target string, body any) *httptest.ResponseRecorder {
|
||||||
|
t.Helper()
|
||||||
|
var buf bytes.Buffer
|
||||||
|
if body != nil {
|
||||||
|
require.NoError(t, json.NewEncoder(&buf).Encode(body))
|
||||||
|
}
|
||||||
|
rec := httptest.NewRecorder()
|
||||||
|
svc.router.ServeHTTP(rec, httptest.NewRequest(method, target, &buf))
|
||||||
|
return rec
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHandleUpdateObservation_EditsAndReturnsTheObservation(t *testing.T) {
|
||||||
|
svc, cleanup := testService(t)
|
||||||
|
defer cleanup()
|
||||||
|
id := createTestObservation(t, svc.observationStore, "proj_aaaaaa", "Old title", "Old narrative", []string{"cache"})
|
||||||
|
|
||||||
|
rec := doUpdateRequest(t, svc, http.MethodPut, fmt.Sprintf("/api/observations/%d", id),
|
||||||
|
map[string]any{"title": "New title", "facts": []string{"one", "two"}})
|
||||||
|
require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
|
||||||
|
var resp struct {
|
||||||
|
Observation struct {
|
||||||
|
Title string `json:"title"`
|
||||||
|
Narrative string `json:"narrative"`
|
||||||
|
Facts []string `json:"facts"`
|
||||||
|
} `json:"observation"`
|
||||||
|
}
|
||||||
|
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
||||||
|
assert.Equal(t, "New title", resp.Observation.Title)
|
||||||
|
assert.Equal(t, "Old narrative", resp.Observation.Narrative)
|
||||||
|
assert.Equal(t, []string{"one", "two"}, resp.Observation.Facts)
|
||||||
|
|
||||||
|
rec = doUpdateRequest(t, svc, http.MethodGet, fmt.Sprintf("/api/observations/%d", id), nil)
|
||||||
|
require.Equal(t, http.StatusOK, rec.Code)
|
||||||
|
assert.Contains(t, rec.Body.String(), "New title", "the edit is stored")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHandleUpdateObservation_Refusals(t *testing.T) {
|
||||||
|
svc, cleanup := testService(t)
|
||||||
|
defer cleanup()
|
||||||
|
id := createTestObservation(t, svc.observationStore, "proj_aaaaaa", "A title", "A narrative", nil)
|
||||||
|
|
||||||
|
assert.Equal(t, http.StatusNotFound, doUpdateRequest(t, svc, http.MethodPut, "/api/observations/999999", map[string]any{"title": "x"}).Code)
|
||||||
|
assert.Equal(t, http.StatusBadRequest, doUpdateRequest(t, svc, http.MethodPut, fmt.Sprintf("/api/observations/%d", id), map[string]any{"scope": "galaxy"}).Code)
|
||||||
|
assert.Equal(t, http.StatusBadRequest, doUpdateRequest(t, svc, http.MethodPut, "/api/observations/abc", map[string]any{"title": "x"}).Code)
|
||||||
|
}
|
||||||
@@ -51,7 +51,7 @@ func (p *Processor) checkVectorDeduplication(ctx context.Context, obs *models.Pa
|
|||||||
for _, r := range results {
|
for _, r := range results {
|
||||||
if r.Similarity >= cfg.DeduplicationThreshold {
|
if r.Similarity >= cfg.DeduplicationThreshold {
|
||||||
obsID := extractObservationIDFromVectorDoc(r)
|
obsID := extractObservationIDFromVectorDoc(r)
|
||||||
if obsID > 0 {
|
if obsID > 0 && p.canMergeInto(ctx, obsID) {
|
||||||
return &DeduplicationResult{
|
return &DeduplicationResult{
|
||||||
ExistingID: obsID,
|
ExistingID: obsID,
|
||||||
Similarity: r.Similarity,
|
Similarity: r.Similarity,
|
||||||
@@ -64,6 +64,20 @@ func (p *Processor) checkVectorDeduplication(ctx context.Context, obs *models.Pa
|
|||||||
return &DeduplicationResult{Action: "insert"}
|
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.
|
// buildObservationSearchText creates searchable text from a parsed observation.
|
||||||
func buildObservationSearchText(obs *models.ParsedObservation) string {
|
func buildObservationSearchText(obs *models.ParsedObservation) string {
|
||||||
var parts []string
|
var parts []string
|
||||||
|
|||||||
@@ -4,6 +4,7 @@ import (
|
|||||||
"context"
|
"context"
|
||||||
"testing"
|
"testing"
|
||||||
|
|
||||||
|
"github.com/lukaszraczylo/claude-mnemonic/internal/db/gorm"
|
||||||
"github.com/lukaszraczylo/claude-mnemonic/internal/vector/sqlitevec"
|
"github.com/lukaszraczylo/claude-mnemonic/internal/vector/sqlitevec"
|
||||||
"github.com/lukaszraczylo/claude-mnemonic/pkg/models"
|
"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)
|
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")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
+9
-9
@@ -13,7 +13,7 @@
|
|||||||
{
|
{
|
||||||
"name": "claude-mnemonic",
|
"name": "claude-mnemonic",
|
||||||
"description": "Persistent memory system for Claude Code - Go implementation with SQLite and ChromaDB vector search",
|
"description": "Persistent memory system for Claude Code - Go implementation with SQLite and ChromaDB vector search",
|
||||||
"version": "0.21.91",
|
"version": "0.21.105",
|
||||||
"author": {
|
"author": {
|
||||||
"name": "lukaszraczylo",
|
"name": "lukaszraczylo",
|
||||||
"url": "https://github.com/lukaszraczylo"
|
"url": "https://github.com/lukaszraczylo"
|
||||||
@@ -23,25 +23,25 @@
|
|||||||
"license": "MIT",
|
"license": "MIT",
|
||||||
"homepage": "https://github.com/lukaszraczylo/claude-mnemonic",
|
"homepage": "https://github.com/lukaszraczylo/claude-mnemonic",
|
||||||
"releases": {
|
"releases": {
|
||||||
"latest": "0.21.91",
|
"latest": "0.21.105",
|
||||||
"versions": {
|
"versions": {
|
||||||
"0.21.91": {
|
"0.21.105": {
|
||||||
"releaseDate": "2026-10-04",
|
"releaseDate": "2026-10-04",
|
||||||
"downloads": {
|
"downloads": {
|
||||||
"darwin-arm64": {
|
"darwin-arm64": {
|
||||||
"url": "https://github.com/lukaszraczylo/claude-mnemonic/releases/download/v0.21.91/claude-mnemonic_0.21.91_darwin_arm64.tar.gz",
|
"url": "https://github.com/lukaszraczylo/claude-mnemonic/releases/download/v0.21.105/claude-mnemonic_0.21.105_darwin_arm64.tar.gz",
|
||||||
"format": "tar.gz",
|
"format": "tar.gz",
|
||||||
"sha256": "96d212aaed9730dca9f2596676014825d3f2377353f0a9658cda014646610936"
|
"sha256": "162d3028d2340d6c2d5438607164ad62b99348d2eddb7186527b8c39642825fa"
|
||||||
},
|
},
|
||||||
"linux-amd64": {
|
"linux-amd64": {
|
||||||
"url": "https://github.com/lukaszraczylo/claude-mnemonic/releases/download/v0.21.91/claude-mnemonic_0.21.91_linux_amd64.tar.gz",
|
"url": "https://github.com/lukaszraczylo/claude-mnemonic/releases/download/v0.21.105/claude-mnemonic_0.21.105_linux_amd64.tar.gz",
|
||||||
"format": "tar.gz",
|
"format": "tar.gz",
|
||||||
"sha256": "5d6a0a6076e9f9f1d2371903df14c47722a5e03ef05e1db84a533f631d40f149"
|
"sha256": "8fc377726456685e4b5c21ea6720bc03b108e846c1daeb7c296ba9a040661308"
|
||||||
},
|
},
|
||||||
"windows-amd64": {
|
"windows-amd64": {
|
||||||
"url": "https://github.com/lukaszraczylo/claude-mnemonic/releases/download/v0.21.91/claude-mnemonic_0.21.91_windows_amd64.zip",
|
"url": "https://github.com/lukaszraczylo/claude-mnemonic/releases/download/v0.21.105/claude-mnemonic_0.21.105_windows_amd64.zip",
|
||||||
"format": "zip",
|
"format": "zip",
|
||||||
"sha256": "f2344e5e797ffa18468e4b65ef7d28c6949c405b254523af68ab8c8fd5fc828f"
|
"sha256": "555089ce26bdb3958e2aa98cdcd5ffe0216fa1dae88f1db5222bffe1cce99aae"
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user