Cyrene_v2/internal/runtime/memory/vector/index.go

// 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")
}