Cyrene_v2/internal/runtime/memory/vector/pg.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 驱动，注册名为 "pgx"
)

// pgMaxOpenConns 是 DevDocs/02-development-constraints.md §2.2 为
// pgvector 增强规定的 PostgreSQL 连接池大小。
const pgMaxOpenConns = 25

// PGVectorIndex 把向量存入 PostgreSQL 的 pgvector vector 列，并
// 用 SQL 中的余弦距离运算符 <=> 作答相似度。它是配置了
// DB_PG_URL 时选中的可选增强；启用时，
// 语义检索真正由 pgvector 支撑，不会退化为
// 关键词匹配（DevDocs/design/03-记忆.md §3）。
//
// schema 由 NewPG 准备（vector 扩展 + vec_items 表）。Search
// 发出 ORDER BY vector <=> $1 LIMIT n，返回精确的余弦距离
// 排名结果。没有 ANN 索引时，查询规划器做顺序扫描，
// 这正是 M4 里程碑（小数据集）期望的线性扫描行为。
// HNSW 索引（CREATE INDEX ... USING hnsw (vector vector_cosine_ops)）是
// 大数据集下既定的 ANN 升级方案，且无需改动 Search。
type PGVectorIndex struct {
	db       *sql.DB
	dim      int
	embedder embedding.Embedder
	log      *slog.Logger
}

// NewPG 用 pgx 驱动连接到 pgURL 的 PostgreSQL+pgvector 数据库
// 并准备 schema：vector 扩展与一张带固定维度 vector 列的
// vec_items 表。它刻意不创建 ANN 索引：
// M4 里程碑下 <=> 查询走精确顺序扫描，而 HNSW
// 索引是既定的大数据量可选升级。pgURL 为空
// 返回 ErrInvalidPGURL，dim 非正返回 ErrInvalidDim，两者都在
// 任何网络 I/O 之前。ctx 控制连接 ping 与 schema 准备。
//
// 返回的索引没有接入 embedder；用 New 构建完全接好的
// 索引，或在调用 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 幂等地创建 pgvector 扩展与 vec_items 表。
// M4 里程碑不创建 ANN 索引 —— <->/<=> 查询退化为
// 精确顺序扫描，与 SQLite 路径的线性扫描语义一致。
// 要扩容，可添加：
//
//	CREATE INDEX IF NOT EXISTS idx_vec_items_vector
//	    ON vec_items USING hnsw (vector vector_cosine_ops)
//
// 且无需改动 Search。每条语句在取输入处都参数化；
// 表 DDL 在维度替换后是静态的。
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 把产向量的 embedder 接入索引。New 使用它；
// 测试也可直接调用。
func (p *PGVectorIndex) withEmbedder(e embedding.Embedder) { p.embedder = e }

// Dim 返回固定向量维度。
func (p *PGVectorIndex) Dim() int { return p.dim }

// Embed 用接入的 embedder 把 texts 转成向量，每段输入文本
// 对应一个按索引对齐的 Vector。未接入 embedder 时返回
// 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 持久化 (id, vector) 对，id 重复时 upsert。向量以
// 绑定到 $n::vector 参数的向量字面量传给 Postgres。它
// 拒绝空 id（ErrMissingID）以及与索引维度不一致的
// 向量（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
}

// Vectors 返回与 ids 对齐的已存向量：命中按位置给出，未命中（含维度与
// 本索引不符的行、解析不了的行）为 nil。它是「增量向量化」的读口 ——
// 调用方据此只嵌缺的那些，而不是每轮重算整个分区。它只读，不写任何行。
//
// 向量列显式 ::text 取出：pgvector 的 OID 不在 database/sql 的标准类型里，
// 直接 scan 进 string 依赖驱动猜测格式；显式转型把「取什么格式」写在 SQL 里。
func (p *PGVectorIndex) Vectors(ctx context.Context, ids []string) ([]Vector, error) {
	if len(ids) == 0 {
		return []Vector{}, nil
	}
	found := make(map[string]Vector, len(ids))
	for start := 0; start < len(ids); start += vectorsQueryChunk {
		end := min(start+vectorsQueryChunk, len(ids))
		if err := p.lookupVectors(ctx, ids[start:end], found); err != nil {
			return nil, err
		}
	}
	return alignVectors(ids, found), nil
}

// lookupVectors 查一批 id 并把命中的向量写进 found。维度与本索引不符
// （或解析不了）的行当作未命中 —— 它不能参与本索引的余弦，重算一次才是对的。
func (p *PGVectorIndex) lookupVectors(ctx context.Context, ids []string, found map[string]Vector) error {
	query := `SELECT id, vector::text FROM vec_items WHERE id IN (` + pgPlaceholders(len(ids)) + `)`
	rows, err := p.db.QueryContext(ctx, query, idArgs(ids)...)
	if err != nil {
		return fmt.Errorf("memory/vector: lookup pg vectors: %w", err)
	}
	defer func() { _ = rows.Close() }()

	for rows.Next() {
		var (
			id   string
			text string
		)
		if err := rows.Scan(&id, &text); err != nil {
			return fmt.Errorf("memory/vector: scan pg vector: %w", err)
		}
		v, err := pgParseVector(text)
		if err != nil || len(v) != p.dim {
			continue
		}
		found[id] = v
	}
	if err := rows.Err(); err != nil {
		return fmt.Errorf("memory/vector: iterate pg vectors: %w", err)
	}
	return nil
}

// Search 返回与 q 最相似的最多 topK 个 id，按余弦相似度
// （1 - 余弦距离）降序排列，由 pgvector 在 SQL 中计算。topK 非正
// 返回空结果；长度与索引维度不一致的查询
// 返回 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 释放底层的 PostgreSQL 连接池。
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 把 v 渲染成 pgvector 的文本形式 "[a,b,c]"，
// 以便绑定为向量参数（$n::vector）。取值的格式化精度足以
// 往返一个 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()
}

// pgParseVector 把 pgvector 的文本形式 "[a,b,c]" 解析回 Vector ——
// pgVectorLiteral 的逆运算。它拒绝空串与任何非浮点的分量，使损坏的行以
// 错误暴露（调用方据此把它当未命中），而不是被静默读成一串零向量。
func pgParseVector(text string) (Vector, error) {
	trimmed := strings.TrimSpace(text)
	if len(trimmed) < 2 || trimmed[0] != '[' || trimmed[len(trimmed)-1] != ']' {
		return nil, fmt.Errorf("memory/vector: invalid pgvector literal of length %d", len(trimmed))
	}
	body := trimmed[1 : len(trimmed)-1]
	if body == "" {
		return Vector{}, nil
	}
	parts := strings.Split(body, ",")
	out := make(Vector, len(parts))
	for i, part := range parts {
		f, err := strconv.ParseFloat(strings.TrimSpace(part), 32)
		if err != nil {
			return nil, fmt.Errorf("memory/vector: parse pgvector component %d: %w", i, err)
		}
		out[i] = float32(f)
	}
	return out, nil
}