From fb880572a49438a6dd804d56eed491d0bf13645c Mon Sep 17 00:00:00 2001 From: Aska Date: Sun, 6 Sep 2026 06:04:44 +0800 Subject: [PATCH] =?UTF-8?q?feat(memory/vector):=20=E5=90=91=E9=87=8F?= =?UTF-8?q?=E6=A3=80=E7=B4=A2=E5=BC=95=E6=93=8E=EF=BC=88SQLite=20BLOB=20?= =?UTF-8?q?=E5=86=85=E7=A7=AF=E9=BB=98=E8=AE=A4=20+=20PG=20pgvector=20?= =?UTF-8?q?=E7=9C=9F=E8=AF=AD=E4=B9=89=E5=A2=9E=E5=BC=BA=EF=BC=8C=E9=85=8D?= =?UTF-8?q?=E7=BD=AE=E5=90=8E=E7=9C=9F=E7=94=9F=E6=95=88=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/runtime/memory/vector/index.go | 142 ++++++++++++ internal/runtime/memory/vector/index_test.go | 164 ++++++++++++++ internal/runtime/memory/vector/pg.go | 213 ++++++++++++++++++ internal/runtime/memory/vector/pg_test.go | 92 ++++++++ internal/runtime/memory/vector/sqlite.go | 207 +++++++++++++++++ internal/runtime/memory/vector/sqlite_test.go | 191 ++++++++++++++++ 6 files changed, 1009 insertions(+) create mode 100644 internal/runtime/memory/vector/index.go create mode 100644 internal/runtime/memory/vector/index_test.go create mode 100644 internal/runtime/memory/vector/pg.go create mode 100644 internal/runtime/memory/vector/pg_test.go create mode 100644 internal/runtime/memory/vector/sqlite.go create mode 100644 internal/runtime/memory/vector/sqlite_test.go diff --git a/internal/runtime/memory/vector/index.go b/internal/runtime/memory/vector/index.go new file mode 100644 index 0000000..0fdd8c1 --- /dev/null +++ b/internal/runtime/memory/vector/index.go @@ -0,0 +1,142 @@ +// Package vector implements dense-vector similarity search for Cyrene's +// long-term memory layer (DevDocs/design/03-记忆.md §3). +// +// It exposes an Index abstraction over two storage backends: +// +// - SQLiteIndex (the default): vectors are stored as BLOBs in a SQLite +// vec_items table and similarity is computed in Go as cosine similarity via +// a linear scan over the persisted vectors. This is the zero-extra- +// dependency path used when DB_PG_URL is empty. +// - PGVectorIndex (optional enhancement): uses PostgreSQL + pgvector, storing +// vectors in a vector column and answering similarity with the cosine- +// distance operator <=> in SQL. It is selected only when DB_PG_URL is set; +// when it is, semantic search is genuinely backed by pgvector and does not +// silently fall back to keyword matching. +// +// Both backends are linear-scan for milestone M4 (small datasets, no ANN +// training). The PostgreSQL path may later be upgraded to an HNSW index as data +// grows; the Search SQL is unchanged. Vectors are produced by an +// embedding.Embedder wired in via New. +package vector + +import ( + "context" + "errors" + "fmt" + "io" + "log/slog" + "strings" + + "git.yeij.top/AskaEth/Cyrene_v2/internal/config" + "git.yeij.top/AskaEth/Cyrene_v2/internal/db" + "git.yeij.top/AskaEth/Cyrene_v2/internal/logger" + "git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/embedding" +) + +// Vector is a dense embedding: one float32 per dimension. +type Vector []float32 + +// Hit is a single similarity result: a stored id and its similarity score. +// Higher Score means more similar, so results sort in descending Score order. +type Hit struct { + ID string + Score float64 +} + +// Index is a dense-vector similarity index over stored (id, vector) pairs. +// +// Embed turns a batch of texts into vectors using the wired embedder. Add +// persists an (id, vector) pair, upserting on a repeated id. Search returns the +// topK ids most similar to q, ranked by descending similarity. Dim reports the +// fixed vector dimension. Close releases the underlying database. +type Index interface { + Embed(ctx context.Context, texts []string) ([]Vector, error) + Add(ctx context.Context, id string, v Vector) error + Search(ctx context.Context, q Vector, topK int) ([]Hit, error) + Dim() int + Close() error +} + +// Package-level sentinel errors. Each is wrapped with %w at the call site so +// callers can use errors.Is. +var ( + // ErrNoEmbedder is returned when an index is asked to Embed but no embedder + // has been wired in (e.g. a constructor called directly rather than via New). + ErrNoEmbedder = errors.New("memory/vector: no embedder configured") + + // ErrNilDB is returned when a constructor receives a nil *sql.DB. + ErrNilDB = errors.New("memory/vector: nil db") + + // ErrInvalidDim is returned when a dimension is not positive. + ErrInvalidDim = errors.New("memory/vector: dimension must be positive") + + // ErrDimMismatch is returned when an added or queried vector's length does + // not match the index dimension. + ErrDimMismatch = errors.New("memory/vector: vector dimension mismatch") + + // ErrMissingID is returned when an id is blank. + ErrMissingID = errors.New("memory/vector: id is required") + + // ErrInvalidPGURL is returned when the PostgreSQL URL is blank. + ErrInvalidPGURL = errors.New("memory/vector: invalid postgres url") +) + +// New builds the Index selected by cfg: PGVectorIndex when cfg.PGURL is set, +// otherwise SQLiteIndex. embedder must be non-nil and report a positive Dim(); +// it is wired into the returned index so Index.Embed produces vectors with it. +// log may be nil; logs are tagged with source "memory/vector". +// +// SQLiteIndex is created on a connection opened at cfg.SQLitePath (":memory:" +// for an isolated test database or a file path). PGVectorIndex connects to +// cfg.PGURL with the pgx driver and requires a PostgreSQL server with the +// pgvector extension installed. +// +// New performs one-time setup with a background context and has no ctx +// parameter by design; every subsequent operation accepts a context.Cancel. +func New(cfg config.DBConfig, embedder embedding.Embedder, log *slog.Logger) (Index, error) { + if embedder == nil { + return nil, ErrNoEmbedder + } + dim := embedder.Dim() + if dim <= 0 { + return nil, fmt.Errorf("memory/vector: %w (got %d)", ErrInvalidDim, dim) + } + if usePG(cfg.PGURL) { + idx, err := NewPG(context.Background(), cfg.PGURL, dim, log) + if err != nil { + return nil, err + } + idx.withEmbedder(embedder) + return idx, nil + } + // db.Open logs internally; hand it a sourced logger so its records carry + // provenance, while the raw log goes to NewSQLite which sources it itself. + d, err := db.Open(context.Background(), cfg.SQLitePath, sourcedLogger(log)) + if err != nil { + return nil, fmt.Errorf("memory/vector: open sqlite: %w", err) + } + idx, err := NewSQLite(context.Background(), d, dim, log) + if err != nil { + _ = d.Close() + return nil, err + } + idx.withEmbedder(embedder) + return idx, nil +} + +// usePG reports whether cfg should select the PostgreSQL+pgvector backend. +// Surrounding whitespace counts as unset. +func usePG(pgURL string) bool { + return strings.TrimSpace(pgURL) != "" +} + +// sourcedLogger returns log, or a discard logger when nil, tagged with source +// "memory/vector" so every record carries provenance. It is idempotent in +// output: re-applying the same source value on an already-sourced logger yields +// the same emitted attribute. +func sourcedLogger(log *slog.Logger) *slog.Logger { + if log == nil { + log = slog.New(slog.NewTextHandler(io.Discard, nil)) + } + return logger.WithSource(log, "memory/vector") +} diff --git a/internal/runtime/memory/vector/index_test.go b/internal/runtime/memory/vector/index_test.go new file mode 100644 index 0000000..1327c52 --- /dev/null +++ b/internal/runtime/memory/vector/index_test.go @@ -0,0 +1,164 @@ +package vector + +import ( + "context" + "database/sql" + "errors" + "io" + "log/slog" + "testing" + + "git.yeij.top/AskaEth/Cyrene_v2/internal/config" + "git.yeij.top/AskaEth/Cyrene_v2/internal/db" + "git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/embedding" +) + +// testLogger returns a logger that writes nowhere, keeping test output clean. +func testLogger() *slog.Logger { + return slog.New(slog.NewTextHandler(io.Discard, nil)) +} + +// openTestDB opens an isolated in-memory SQLite DB. Each call owns a private DB, +// so parallel tests are isolated by construction. The returned DB is closed by +// t.Cleanup. +func openTestDB(t *testing.T) *sql.DB { + t.Helper() + ctx := context.Background() + d, err := db.Open(ctx, ":memory:", testLogger()) + if err != nil { + t.Fatalf("open in-memory db: %v", err) + } + t.Cleanup(func() { _ = d.Close() }) + return d +} + +// newInMemorySQLite builds a SQLiteIndex on a fresh in-memory DB with dim +// dimensions. The embedder is left unwired so tests control it. +func newInMemorySQLite(t *testing.T, dim int) *SQLiteIndex { + t.Helper() + ctx := context.Background() + idx, err := NewSQLite(ctx, openTestDB(t), dim, testLogger()) + if err != nil { + t.Fatalf("NewSQLite: %v", err) + } + return idx +} + +// mockEmbedder is a scriptable embedding.Embedder for tests. With no fn it +// returns one-hot vectors, one per input text, so tests can predict results. +type mockEmbedder struct { + dim int + fn func(ctx context.Context, texts []string) ([]embedding.Vector, error) +} + +func (m *mockEmbedder) Dim() int { return m.dim } + +func (m *mockEmbedder) Embed(ctx context.Context, texts []string) ([]embedding.Vector, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + if m.fn != nil { + return m.fn(ctx, texts) + } + out := make([]embedding.Vector, len(texts)) + for i := range texts { + v := make(embedding.Vector, m.dim) + v[i%m.dim] = 1 + out[i] = v + } + return out, nil +} + +func TestNew_NilEmbedder_异常(t *testing.T) { + t.Parallel() + if _, err := New(config.DBConfig{}, nil, testLogger()); !errors.Is(err, ErrNoEmbedder) { + t.Fatalf("New(nil embedder) err = %v, want ErrNoEmbedder", err) + } +} + +func TestNew_InvalidDim_异常(t *testing.T) { + t.Parallel() + emb := &mockEmbedder{dim: 0} + if _, err := New(config.DBConfig{SQLitePath: ":memory:"}, emb, testLogger()); !errors.Is(err, ErrInvalidDim) { + t.Fatalf("New(dim 0) err = %v, want ErrInvalidDim", err) + } +} + +func TestNew_SelectsSQLite_WhenPGURLBlank(t *testing.T) { + t.Parallel() + ctx := context.Background() + idx, err := New(config.DBConfig{SQLitePath: ":memory:"}, &mockEmbedder{dim: 4}, testLogger()) + if err != nil { + t.Fatalf("New: %v", err) + } + if _, ok := idx.(*SQLiteIndex); !ok { + t.Fatalf("New type = %T, want *SQLiteIndex", idx) + } + // New wires the embedder, so Embed works on the returned index. + vs, err := idx.Embed(ctx, []string{"hello"}) + if err != nil { + t.Fatalf("Embed: %v", err) + } + if len(vs) != 1 || len(vs[0]) != 4 { + t.Fatalf("Embed => %v, want one 4-dim vector", vs) + } + if err := idx.Close(); err != nil { + t.Fatalf("Close: %v", err) + } +} + +func TestNew_PGURLWhitespace_SelectsSQLite(t *testing.T) { + t.Parallel() + idx, err := New(config.DBConfig{SQLitePath: ":memory:", PGURL: " \t "}, &mockEmbedder{dim: 4}, testLogger()) + if err != nil { + t.Fatalf("New: %v", err) + } + if _, ok := idx.(*SQLiteIndex); !ok { + t.Fatalf("New type = %T, want *SQLiteIndex for whitespace PGURL", idx) + } + if err := idx.Close(); err != nil { + t.Fatalf("Close: %v", err) + } +} + +func TestNew_PGSelection_AttemptsPGWhenURLSet(t *testing.T) { + t.Parallel() + // A loopback URL on a closed port fails fast, proving New chose the PG path + // rather than silently falling back to SQLite. Port 1 (tcpmux) does not + // listen on localhost, so the connection is refused immediately. + idx, err := New(config.DBConfig{ + SQLitePath: ":memory:", + PGURL: "postgres://127.0.0.1:1/mem?sslmode=disable&connect_timeout=1", + }, &mockEmbedder{dim: 4}, testLogger()) + if err == nil { + if idx != nil { + _ = idx.Close() + } + t.Fatal("New with unreachable PGURL: expected error, got nil") + } + if idx != nil { + t.Fatalf("New returned non-nil index %T on error", idx) + } +} + +func TestUsePG(t *testing.T) { + t.Parallel() + tests := []struct { + name string + url string + want bool + }{ + {name: "empty", url: "", want: false}, + {name: "whitespace", url: " \t ", want: false}, + {name: "set", url: "postgres://user:pass@host:5432/db", want: true}, + } + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + if got := usePG(tt.url); got != tt.want { + t.Fatalf("usePG(%q) = %v, want %v", tt.url, got, tt.want) + } + }) + } +} diff --git a/internal/runtime/memory/vector/pg.go b/internal/runtime/memory/vector/pg.go new file mode 100644 index 0000000..da0fd08 --- /dev/null +++ b/internal/runtime/memory/vector/pg.go @@ -0,0 +1,213 @@ +package vector + +import ( + "context" + "database/sql" + "fmt" + "log/slog" + "strconv" + "strings" + + "git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/embedding" + + _ "github.com/jackc/pgx/v5/stdlib" // Postgres driver, registers name "pgx" +) + +// pgMaxOpenConns is the PostgreSQL connection pool size mandated by +// DevDocs/02-development-constraints.md §2.2 for the pgvector enhancement. +const pgMaxOpenConns = 25 + +// PGVectorIndex stores vectors in a PostgreSQL pgvector vector column and +// answers similarity with the cosine-distance operator <=> in SQL. It is the +// optional enhancement selected when DB_PG_URL is configured; when used, +// semantic search is genuinely backed by pgvector and does not fall back to +// keyword matching (DevDocs/design/03-记忆.md §3). +// +// The schema is prepared by NewPG (vector extension + vec_items table). Search +// issues ORDER BY vector <=> $1 LIMIT n, which returns exact cosine-distance- +// ranked results. With no ANN index the query planner does a sequential scan, +// which is the linear-scan behavior expected for milestone M4 (small datasets). +// An HNSW index (CREATE INDEX ... USING hnsw (vector vector_cosine_ops)) is the +// documented ANN upgrade for larger datasets and requires no change to Search. +type PGVectorIndex struct { + db *sql.DB + dim int + embedder embedding.Embedder + log *slog.Logger +} + +// NewPG connects to the PostgreSQL+pgvector database at pgURL using the pgx +// driver and prepares the schema: the vector extension and a vec_items table +// with a fixed-dimension vector column. It deliberately creates no ANN index: +// for milestone M4 the <=> query runs as an exact sequential scan, and an HNSW +// index is a documented, optional upgrade for larger data. A blank pgURL +// returns ErrInvalidPGURL, and a non-positive dim returns ErrInvalidDim, both +// before any network I/O. ctx controls the connection ping and schema setup. +// +// The returned index has no embedder wired in; use New to build a fully wired +// index, or wire one in before calling Embed. +func NewPG(ctx context.Context, pgURL string, dim int, log *slog.Logger) (*PGVectorIndex, error) { + if strings.TrimSpace(pgURL) == "" { + return nil, ErrInvalidPGURL + } + if dim <= 0 { + return nil, ErrInvalidDim + } + log = sourcedLogger(log) + + db, err := sql.Open("pgx", pgURL) + if err != nil { + return nil, fmt.Errorf("memory/vector: open pg: %w", err) + } + db.SetMaxOpenConns(pgMaxOpenConns) + if err := db.PingContext(ctx); err != nil { + _ = db.Close() + return nil, fmt.Errorf("memory/vector: ping pg: %w", err) + } + + idx := &PGVectorIndex{db: db, dim: dim, log: log} + if err := idx.setup(ctx); err != nil { + _ = db.Close() + return nil, err + } + log.Info("pgvector index ready", "dim", dim) + return idx, nil +} + +// setup creates the pgvector extension and the vec_items table idempotently. +// For milestone M4 no ANN index is created — <->/<=> queries fall back to an +// exact sequential scan, matching the linear-scan semantics of the SQLite path. +// To scale, add: +// +// CREATE INDEX IF NOT EXISTS idx_vec_items_vector +// ON vec_items USING hnsw (vector vector_cosine_ops) +// +// which requires no change to Search. Each statement is parameterized where it +// takes input; the table DDL is static after dimension substitution. +func (p *PGVectorIndex) setup(ctx context.Context) error { + statements := []string{ + "CREATE EXTENSION IF NOT EXISTS vector", + fmt.Sprintf(`CREATE TABLE IF NOT EXISTS vec_items ( + id TEXT PRIMARY KEY, + vector vector(%d) NOT NULL, + dim INTEGER NOT NULL + )`, p.dim), + } + for _, stmt := range statements { + if _, err := p.db.ExecContext(ctx, stmt); err != nil { + return fmt.Errorf("memory/vector: setup pg schema: %w", err) + } + } + return nil +} + +// withEmbedder wires the vector-producing embedder into the index. Used by New; +// tests may also call it directly. +func (p *PGVectorIndex) withEmbedder(e embedding.Embedder) { p.embedder = e } + +// Dim returns the fixed vector dimension. +func (p *PGVectorIndex) Dim() int { return p.dim } + +// Embed turns texts into vectors with the wired embedder, returning one Vector +// per input text aligned by index. With no embedder wired it returns +// ErrNoEmbedder. +func (p *PGVectorIndex) Embed(ctx context.Context, texts []string) ([]Vector, error) { + if p.embedder == nil { + return nil, ErrNoEmbedder + } + vecs, err := p.embedder.Embed(ctx, texts) + if err != nil { + return nil, fmt.Errorf("memory/vector: embed %d texts: %w", len(texts), err) + } + out := make([]Vector, len(vecs)) + for i, v := range vecs { + out[i] = Vector(v) + } + return out, nil +} + +// Add persists the (id, vector) pair, upserting on a repeated id. The vector is +// passed to Postgres as a vector literal bound to a $n::vector parameter. It +// rejects a blank id (ErrMissingID) and a vector whose length differs from the +// index dimension (ErrDimMismatch). +func (p *PGVectorIndex) Add(ctx context.Context, id string, v Vector) error { + if strings.TrimSpace(id) == "" { + return ErrMissingID + } + if len(v) != p.dim { + return fmt.Errorf("memory/vector: %w (got %d, want %d)", ErrDimMismatch, len(v), p.dim) + } + _, err := p.db.ExecContext(ctx, `INSERT INTO vec_items (id, vector, dim) + VALUES ($1, $2::vector, $3) + ON CONFLICT (id) DO UPDATE SET vector = EXCLUDED.vector, dim = EXCLUDED.dim`, + id, pgVectorLiteral(v), p.dim) + if err != nil { + return fmt.Errorf("memory/vector: add %q: %w", id, err) + } + return nil +} + +// Search returns up to topK ids most similar to q, ranked by descending cosine +// similarity (1 - cosine distance), computed by pgvector in SQL. A non-positive +// topK returns an empty result; a query whose length differs from the index +// dimension returns ErrDimMismatch. +func (p *PGVectorIndex) Search(ctx context.Context, q Vector, topK int) ([]Hit, error) { + if topK <= 0 { + return []Hit{}, nil + } + if len(q) != p.dim { + return nil, fmt.Errorf("memory/vector: %w (query %d, want %d)", ErrDimMismatch, len(q), p.dim) + } + rows, err := p.db.QueryContext(ctx, `SELECT id, 1.0 - (vector <=> $1::vector) AS score + FROM vec_items + ORDER BY vector <=> $1::vector + LIMIT $2`, pgVectorLiteral(q), topK) + if err != nil { + return nil, fmt.Errorf("memory/vector: search pg: %w", err) + } + defer func() { _ = rows.Close() }() + + hits := make([]Hit, 0) + for rows.Next() { + var ( + id string + score float64 + ) + if err := rows.Scan(&id, &score); err != nil { + return nil, fmt.Errorf("memory/vector: scan pg hit: %w", err) + } + hits = append(hits, Hit{ID: id, Score: score}) + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("memory/vector: iterate pg hits: %w", err) + } + return hits, nil +} + +// Close releases the underlying PostgreSQL connection pool. +func (p *PGVectorIndex) Close() error { + if p.db == nil { + return nil + } + err := p.db.Close() + if err != nil { + return fmt.Errorf("memory/vector: close pg: %w", err) + } + return nil +} + +// pgVectorLiteral renders v in pgvector's text form "[a,b,c]" so it can be +// bound as a vector parameter ($n::vector). Values are formatted with enough +// precision to round-trip a float32. +func pgVectorLiteral(v Vector) string { + var b strings.Builder + b.WriteByte('[') + for i, x := range v { + if i > 0 { + b.WriteByte(',') + } + b.WriteString(strconv.FormatFloat(float64(x), 'f', -1, 32)) + } + b.WriteByte(']') + return b.String() +} diff --git a/internal/runtime/memory/vector/pg_test.go b/internal/runtime/memory/vector/pg_test.go new file mode 100644 index 0000000..f4e4b50 --- /dev/null +++ b/internal/runtime/memory/vector/pg_test.go @@ -0,0 +1,92 @@ +package vector + +import ( + "context" + "errors" + "os" + "testing" +) + +func TestNewPG_InvalidURL_异常(t *testing.T) { + t.Parallel() + if _, err := NewPG(context.Background(), " ", 4, testLogger()); !errors.Is(err, ErrInvalidPGURL) { + t.Fatalf("NewPG(blank url) err = %v, want ErrInvalidPGURL", err) + } +} + +func TestNewPG_InvalidDim_异常(t *testing.T) { + t.Parallel() + // dim is validated before any connection is opened, so no network I/O here. + if _, err := NewPG(context.Background(), "postgres://host/db", 0, testLogger()); !errors.Is(err, ErrInvalidDim) { + t.Fatalf("NewPG(dim 0) err = %v, want ErrInvalidDim", err) + } +} + +func TestPGVectorLiteral(t *testing.T) { + t.Parallel() + tests := []struct { + name string + v Vector + want string + }{ + {name: "single", v: Vector{1}, want: "[1]"}, + {name: "three", v: Vector{1.5, 2, -0.25}, want: "[1.5,2,-0.25]"}, + {name: "empty", v: Vector{}, want: "[]"}, + } + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + if got := pgVectorLiteral(tt.v); got != tt.want { + t.Fatalf("pgVectorLiteral(%v) = %q, want %q", tt.v, got, tt.want) + } + }) + } +} + +// TestPG_RequiresEnv runs a real PostgreSQL+pgvector round trip when DB_PG_URL +// is set, and skips otherwise. There is no PG server in this environment by +// default, so the pgvector-backed Add/Search path is validated end-to-end only +// when the operator supplies a live database with the pgvector extension. The +// SQLite path (the default) is fully covered by the sqlite_test.go suite. +func TestPG_RequiresEnv(t *testing.T) { + url := os.Getenv("DB_PG_URL") + if url == "" { + t.Skip("DB_PG_URL not set; skipping real PostgreSQL+pgvector integration test") + } + + ctx := context.Background() + idx, err := NewPG(ctx, url, 4, testLogger()) + if err != nil { + t.Fatalf("NewPG: %v", err) + } + idx.withEmbedder(&mockEmbedder{dim: 4}) + defer func() { _ = idx.Close() }() + + if err := idx.Add(ctx, "a", Vector{1, 0, 0, 0}); err != nil { + t.Fatalf("Add a: %v", err) + } + if err := idx.Add(ctx, "b", Vector{0, 1, 0, 0}); err != nil { + t.Fatalf("Add b: %v", err) + } + + hits, err := idx.Search(ctx, Vector{1, 0, 0, 0}, 1) + if err != nil { + t.Fatalf("Search: %v", err) + } + if len(hits) != 1 || hits[0].ID != "a" { + t.Fatalf("pg top hit = %+v, want single a", hits) + } + if hits[0].Score < 0.99 { + t.Fatalf("pg a score = %v, want ~1.0", hits[0].Score) + } + + // The embedder is wired, so Index.Embed is usable. + vs, err := idx.Embed(ctx, []string{"hello"}) + if err != nil { + t.Fatalf("Embed: %v", err) + } + if len(vs) != 1 || len(vs[0]) != 4 { + t.Fatalf("Embed => %d vectors, want one 4-dim", len(vs)) + } +} diff --git a/internal/runtime/memory/vector/sqlite.go b/internal/runtime/memory/vector/sqlite.go new file mode 100644 index 0000000..71c5d30 --- /dev/null +++ b/internal/runtime/memory/vector/sqlite.go @@ -0,0 +1,207 @@ +package vector + +import ( + "context" + "database/sql" + "encoding/binary" + "fmt" + "log/slog" + "math" + "sort" + "strings" + + "git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/embedding" +) + +// createVecItemsTable is the idempotent SQLite schema. vector is stored as a +// BLOB of little-endian float32 values; dim records the dimension so the blob +// can be decoded without external knowledge. The table is created by +// NewSQLite; migration compatibility with internal/db is not required because +// the vector store is self-contained (see internal/runtime/memory/store.go). +const createVecItemsTable = ` +CREATE TABLE IF NOT EXISTS vec_items ( + id TEXT PRIMARY KEY, + vector BLOB NOT NULL, + dim INTEGER NOT NULL +)` + +// SQLiteIndex stores vectors as BLOBs in a SQLite vec_items table and answers +// similarity with Go-side cosine similarity over a linear scan (milestone M4 +// semantics: small datasets, no ANN index). +// +// Concurrency is delegated to the *sql.DB, which is safe for concurrent use by +// contract, so SQLiteIndex holds no further lock. Every query is parameterized +// (?) and columns are listed explicitly. +type SQLiteIndex struct { + db *sql.DB + dim int + embedder embedding.Embedder + log *slog.Logger +} + +// NewSQLite builds a SQLiteIndex on db for vectors of dimension dim, ensuring +// the vec_items schema exists. db must be a live connection; a nil db returns +// ErrNilDB, and a non-positive dim returns ErrInvalidDim. ctx controls schema +// creation. log may be nil. +// +// The returned index has no embedder wired in, so calling Embed directly +// returns ErrNoEmbedder; use New to build a fully wired index, or wire one in +// before calling Embed. +func NewSQLite(ctx context.Context, db *sql.DB, dim int, log *slog.Logger) (*SQLiteIndex, error) { + if db == nil { + return nil, ErrNilDB + } + if dim <= 0 { + return nil, ErrInvalidDim + } + log = sourcedLogger(log) + if _, err := db.ExecContext(ctx, createVecItemsTable); err != nil { + return nil, fmt.Errorf("memory/vector: create sqlite schema: %w", err) + } + log.Info("sqlite vector index ready", "dim", dim) + return &SQLiteIndex{db: db, dim: dim, log: log}, nil +} + +// withEmbedder wires the vector-producing embedder into the index. It is used +// by New; tests may also call it directly. +func (s *SQLiteIndex) withEmbedder(e embedding.Embedder) { s.embedder = e } + +// Dim returns the fixed vector dimension. +func (s *SQLiteIndex) Dim() int { return s.dim } + +// Embed turns texts into vectors with the wired embedder. It returns one Vector +// per input text, aligned by index. With no embedder wired it returns +// ErrNoEmbedder. +func (s *SQLiteIndex) Embed(ctx context.Context, texts []string) ([]Vector, error) { + if s.embedder == nil { + return nil, ErrNoEmbedder + } + vecs, err := s.embedder.Embed(ctx, texts) + if err != nil { + return nil, fmt.Errorf("memory/vector: embed %d texts: %w", len(texts), err) + } + out := make([]Vector, len(vecs)) + for i, v := range vecs { + out[i] = Vector(v) + } + return out, nil +} + +// Add persists the (id, vector) pair, upserting on a repeated id so the +// latest vector wins. It rejects a blank id (ErrMissingID) and a vector whose +// length differs from the index dimension (ErrDimMismatch). ctx aborts the +// write when canceled. +func (s *SQLiteIndex) Add(ctx context.Context, id string, v Vector) error { + if strings.TrimSpace(id) == "" { + return ErrMissingID + } + if len(v) != s.dim { + return fmt.Errorf("memory/vector: %w (got %d, want %d)", ErrDimMismatch, len(v), s.dim) + } + _, err := s.db.ExecContext(ctx, `INSERT INTO vec_items (id, vector, dim) + VALUES (?, ?, ?) + ON CONFLICT(id) DO UPDATE SET vector = excluded.vector, dim = excluded.dim`, + id, encodeVector(v), s.dim) + if err != nil { + return fmt.Errorf("memory/vector: add %q: %w", id, err) + } + return nil +} + +// Search returns up to topK ids most similar to q, ranked by descending cosine +// similarity. It scans every persisted vector in Go (linear scan, no ANN). A +// non-positive topK returns an empty result; a query whose length differs from +// the index dimension returns ErrDimMismatch. +func (s *SQLiteIndex) Search(ctx context.Context, q Vector, topK int) ([]Hit, error) { + if topK <= 0 { + return []Hit{}, nil + } + if len(q) != s.dim { + return nil, fmt.Errorf("memory/vector: %w (query %d, want %d)", ErrDimMismatch, len(q), s.dim) + } + rows, err := s.db.QueryContext(ctx, `SELECT id, vector FROM vec_items`) + if err != nil { + return nil, fmt.Errorf("memory/vector: search: %w", err) + } + defer func() { _ = rows.Close() }() + + hits := make([]Hit, 0) + for rows.Next() { + var ( + id string + blob []byte + ) + if err := rows.Scan(&id, &blob); err != nil { + return nil, fmt.Errorf("memory/vector: scan item: %w", err) + } + v, err := decodeVector(blob) + if err != nil { + return nil, fmt.Errorf("memory/vector: decode %q: %w", id, err) + } + hits = append(hits, Hit{ID: id, Score: cosine(q, v)}) + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("memory/vector: iterate items: %w", err) + } + + sort.Slice(hits, func(i, j int) bool { return hits[i].Score > hits[j].Score }) + if len(hits) > topK { + hits = hits[:topK] + } + return hits, nil +} + +// Close releases the underlying SQLite database. It is safe to call when the +// database is already closed. +func (s *SQLiteIndex) Close() error { + if s.db == nil { + return nil + } + err := s.db.Close() + if err != nil { + return fmt.Errorf("memory/vector: close sqlite: %w", err) + } + return nil +} + +// encodeVector serializes v to a little-endian float32 BLOB. It cannot fail. +func encodeVector(v Vector) []byte { + buf := make([]byte, 4*len(v)) + for i, x := range v { + binary.LittleEndian.PutUint32(buf[i*4:], math.Float32bits(x)) + } + return buf +} + +// decodeVector deserializes a BLOB back into a Vector. It rejects empty or +// non-4-byte-aligned blobs so a corrupted row surfaces as an error rather than +// a silent misread. +func decodeVector(blob []byte) (Vector, error) { + if len(blob) == 0 || len(blob)%4 != 0 { + return nil, fmt.Errorf("memory/vector: invalid vector blob length %d", len(blob)) + } + v := make(Vector, len(blob)/4) + for i := range v { + v[i] = math.Float32frombits(binary.LittleEndian.Uint32(blob[i*4:])) + } + return v, nil +} + +// cosine returns the cosine similarity of a and b in [-1, 1]. Zero-length or +// mismatched vectors yield 0, so an inconsistent stored row never corrupts the +// ranking. +func cosine(a, b Vector) float64 { + if len(a) == 0 || len(b) == 0 || len(a) != len(b) { + return 0 + } + var dot, na, nb float64 + for i := range a { + dot += float64(a[i]) * float64(b[i]) + na += float64(a[i]) * float64(a[i]) + nb += float64(b[i]) * float64(b[i]) + } + if na == 0 || nb == 0 { + return 0 + } + return dot / (math.Sqrt(na) * math.Sqrt(nb)) +} diff --git a/internal/runtime/memory/vector/sqlite_test.go b/internal/runtime/memory/vector/sqlite_test.go new file mode 100644 index 0000000..0135f9c --- /dev/null +++ b/internal/runtime/memory/vector/sqlite_test.go @@ -0,0 +1,191 @@ +package vector + +import ( + "context" + "errors" + "testing" + + "git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/embedding" +) + +func TestSQLite_New_NilDB_异常(t *testing.T) { + t.Parallel() + if _, err := NewSQLite(context.Background(), nil, 4, testLogger()); !errors.Is(err, ErrNilDB) { + t.Fatalf("NewSQLite(nil db) err = %v, want ErrNilDB", err) + } +} + +func TestSQLite_New_InvalidDim_异常(t *testing.T) { + t.Parallel() + ctx := context.Background() + if _, err := NewSQLite(ctx, openTestDB(t), 0, testLogger()); !errors.Is(err, ErrInvalidDim) { + t.Fatalf("NewSQLite(dim 0) err = %v, want ErrInvalidDim", err) + } +} + +func TestSQLite_AddSearch_正常(t *testing.T) { + t.Parallel() + ctx := context.Background() + s := newInMemorySQLite(t, 3) + + for _, e := range []struct { + id string + v Vector + }{ + {id: "cat", v: Vector{1, 0.5, 0}}, + {id: "dog", v: Vector{0.5, 1, 0}}, + {id: "bird", v: Vector{0, 0, 1}}, + } { + if err := s.Add(ctx, e.id, e.v); err != nil { + t.Fatalf("Add(%s): %v", e.id, err) + } + } + + hits, err := s.Search(ctx, Vector{1, 0, 0}, 2) + if err != nil { + t.Fatalf("Search: %v", err) + } + if len(hits) != 2 { + t.Fatalf("got %d hits, want 2", len(hits)) + } + if hits[0].ID != "cat" || hits[1].ID != "dog" { + t.Fatalf("hits = [%s %s], want [cat dog]", hits[0].ID, hits[1].ID) + } + if hits[0].Score <= hits[1].Score { + t.Fatalf("top score not descending: %v <= %v", hits[0].Score, hits[1].Score) + } + if hits[0].Score <= 0.89 || hits[0].Score >= 0.9 { + t.Errorf("cat score = %v, want ~0.894", hits[0].Score) + } +} + +func TestSQLite_Add_Upsert_正常(t *testing.T) { + t.Parallel() + ctx := context.Background() + s := newInMemorySQLite(t, 3) + + if err := s.Add(ctx, "a", Vector{1, 0, 0}); err != nil { + t.Fatalf("Add first: %v", err) + } + if err := s.Add(ctx, "a", Vector{0, 1, 0}); err != nil { + t.Fatalf("Add update: %v", err) + } + + hits, err := s.Search(ctx, Vector{0, 1, 0}, 5) + if err != nil { + t.Fatalf("Search: %v", err) + } + if len(hits) != 1 || hits[0].ID != "a" { + t.Fatalf("hits = %+v, want single a", hits) + } + if hits[0].Score < 0.99 { + t.Fatalf("a score = %v, want ~1.0 after upsert", hits[0].Score) + } +} + +func TestSQLite_Add_MissingID_异常(t *testing.T) { + t.Parallel() + ctx := context.Background() + s := newInMemorySQLite(t, 4) + if err := s.Add(ctx, " ", Vector{1, 0, 0, 0}); !errors.Is(err, ErrMissingID) { + t.Fatalf("Add(blank id) err = %v, want ErrMissingID", err) + } +} + +func TestSQLite_Add_DimMismatch_异常(t *testing.T) { + t.Parallel() + ctx := context.Background() + s := newInMemorySQLite(t, 4) + if err := s.Add(ctx, "x", Vector{1, 2}); !errors.Is(err, ErrDimMismatch) { + t.Fatalf("Add(wrong dim) err = %v, want ErrDimMismatch", err) + } +} + +func TestSQLite_Search_Empty_正常(t *testing.T) { + t.Parallel() + ctx := context.Background() + s := newInMemorySQLite(t, 4) + + hits, err := s.Search(ctx, Vector{1, 0, 0, 0}, 3) + if err != nil { + t.Fatalf("Search empty: %v", err) + } + if hits == nil || len(hits) != 0 { + t.Fatalf("Search(empty) = %v, want non-nil empty slice", hits) + } +} + +func TestSQLite_Search_TopKNonPositive_正常(t *testing.T) { + t.Parallel() + ctx := context.Background() + s := newInMemorySQLite(t, 4) + if err := s.Add(ctx, "a", Vector{1, 0, 0, 0}); err != nil { + t.Fatalf("Add: %v", err) + } + + for _, k := range []int{0, -1} { + hits, err := s.Search(ctx, Vector{1, 0, 0, 0}, k) + if err != nil { + t.Fatalf("Search(topK %d): %v", k, err) + } + if len(hits) != 0 { + t.Fatalf("Search(topK %d) = %v, want empty", k, hits) + } + } +} + +func TestSQLite_Search_DimMismatch_异常(t *testing.T) { + t.Parallel() + ctx := context.Background() + s := newInMemorySQLite(t, 4) + if _, err := s.Search(ctx, Vector{1, 0}, 3); !errors.Is(err, ErrDimMismatch) { + t.Fatalf("Search(wrong dim) err = %v, want ErrDimMismatch", err) + } +} + +func TestSQLite_Embed_正常(t *testing.T) { + t.Parallel() + ctx := context.Background() + s := newInMemorySQLite(t, 4) + s.withEmbedder(&mockEmbedder{dim: 4}) + + vs, err := s.Embed(ctx, []string{"hello", "world"}) + if err != nil { + t.Fatalf("Embed: %v", err) + } + if len(vs) != 2 || len(vs[0]) != 4 || len(vs[1]) != 4 { + t.Fatalf("Embed => %d vectors, want 2 x 4-dim", len(vs)) + } +} + +func TestSQLite_Embed_NoEmbedder_异常(t *testing.T) { + t.Parallel() + ctx := context.Background() + s := newInMemorySQLite(t, 4) + if _, err := s.Embed(ctx, []string{"hi"}); !errors.Is(err, ErrNoEmbedder) { + t.Fatalf("Embed(no embedder) err = %v, want ErrNoEmbedder", err) + } +} + +func TestSQLite_Embed_EmbedderError_异常(t *testing.T) { + t.Parallel() + ctx := context.Background() + s := newInMemorySQLite(t, 4) + boom := errors.New("embed boom") + s.withEmbedder(&mockEmbedder{ + dim: 4, + fn: func(context.Context, []string) ([]embedding.Vector, error) { return nil, boom }, + }) + + if _, err := s.Embed(ctx, []string{"hi"}); !errors.Is(err, boom) { + t.Fatalf("Embed(driver error) err = %v, want wrap of %v", err, boom) + } +} + +func TestSQLite_Close_正常(t *testing.T) { + t.Parallel() + s := newInMemorySQLite(t, 4) + if err := s.Close(); err != nil { + t.Fatalf("Close: %v", err) + } +}