214 lines
7.1 KiB
Go
214 lines
7.1 KiB
Go
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()
|
|
}
|