feat(memory/vector): 向量检索引擎（SQLite BLOB 内积默认 + PG pgvector 真语义增强，配置后真生效）

2026-09-06 06:04:44 +08:00
1 parent ef79edd96a
commit fb880572a4
@@ -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()
+}