// 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 { // API-backed embedders do not know their vector dimension until the first // successful response. Probe once so the index is built with the real // dimension instead of erroring out and falling back to in-memory. if _, err := embedder.Embed(context.Background(), []string{"probe"}); err != nil { return nil, fmt.Errorf("memory/vector: probe embedder dim: %w", err) } 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") }