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

// vector 包为 Cyrene 的长期记忆层（DevDocs/design/03-记忆.md §3）
// 实现稠密向量相似度检索。
//
// 它在两种存储后端之上暴露 Index 抽象：
//
//   - SQLiteIndex（默认）：向量以 BLOB 存在 SQLite 的
//     vec_items 表中，相似度在 Go 侧通过对已持久化向量
//     线性扫描计算余弦相似度。这是 DB_PG_URL 为空时使用的
//     零额外依赖路径。
//   - PGVectorIndex（可选增强）：使用 PostgreSQL + pgvector，把
//     向量存入 vector 列，并用 SQL 中的余弦
//     距离运算符 <=> 作答相似度。仅当设置了 DB_PG_URL 时才选中；
//     此时语义检索真正由 pgvector 支撑，不会
//     静默退化为关键词匹配。
//
// 两个后端在 M4 里程碑都是线性扫描（数据集小，不做 ANN
// 训练）。随着数据增长，PostgreSQL 路径日后可升级为 HNSW 索引；
// Search SQL 不变。向量由经 New 接入的
// embedding.Embedder 产出。
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 是一个稠密向量：每维一个 float32。
type Vector []float32

// Hit 是一条相似度结果：一个已存 id 及其相似度得分。
// Score 越高越相似，因此结果按 Score 降序排列。
type Hit struct {
	ID    string
	Score float64
}

// Index 是建立在已存 (id, vector) 对上的稠密向量相似度索引。
//
// Embed 用接入的 embedder 把一批文本转成向量。Add
// 持久化一个 (id, vector) 对，id 重复时 upsert。Search 返回与 q
// 最相似的前 topK 个 id，按相似度降序排列。Dim 报告
// 固定向量维度。Close 释放底层数据库。
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
}

// 包级哨兵错误。每个都在调用点用 %w 包裹，调用方
// 可用 errors.Is 判断。
var (
	// ErrNoEmbedder 在索引被要求 Embed 却没有接入 embedder
	// 时返回（例如直接调用构造函数而非经 New）。
	ErrNoEmbedder = errors.New("memory/vector: no embedder configured")

	// ErrNilDB 在构造函数收到 nil *sql.DB 时返回。
	ErrNilDB = errors.New("memory/vector: nil db")

	// ErrInvalidDim 在维度非正时返回。
	ErrInvalidDim = errors.New("memory/vector: dimension must be positive")

	// ErrDimMismatch 在新增或查询的向量长度与
	// 索引维度不匹配时返回。
	ErrDimMismatch = errors.New("memory/vector: vector dimension mismatch")

	// ErrMissingID 在 id 为空时返回。
	ErrMissingID = errors.New("memory/vector: id is required")

	// ErrInvalidPGURL 在 PostgreSQL URL 为空时返回。
	ErrInvalidPGURL = errors.New("memory/vector: invalid postgres url")
)

// New 构建 cfg 选定的 Index：设置了 cfg.PGURL 时为 PGVectorIndex，
// 否则为 SQLiteIndex。embedder 必须非 nil 且 Dim() 为正；
// 它被接入返回的索引，使 Index.Embed 用它产向量。
// log 可为 nil；日志带 source "memory/vector" 标签。
//
// SQLiteIndex 建立在 cfg.SQLitePath 打开的连接上（":memory:"
// 用于隔离的测试数据库，或一个文件路径）。PGVectorIndex 用
// pgx 驱动连接 cfg.PGURL，要求已安装
// pgvector 扩展的 PostgreSQL 服务。
//
// New 用后台 context 做一次性准备，按设计不接收 ctx
// 参数；后续每个操作都接受 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 的 embedder 在首次成功响应前不知道自己的向量维度。
		// 探测一次，使索引以真实
		// 维度构建，而不是报错并回退到内存内实现。
		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 内部会记日志；交给它一个带来源的 logger，使其记录携带
	// 出处，而原始 log 传给 NewSQLite，由后者自行加来源。
	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 报告 cfg 是否应选择 PostgreSQL+pgvector 后端。
// 首尾空白视为未设置。
func usePG(pgURL string) bool {
	return strings.TrimSpace(pgURL) != ""
}

// sourcedLogger 返回 log，nil 时返回丢弃型 logger，并打上 source
// "memory/vector" 标签，使每条记录都携带出处。它在
// 输出上幂等：对已带来源的 logger 再次施加相同 source 值，
// 得到的输出属性相同。
func sourcedLogger(log *slog.Logger) *slog.Logger {
	if log == nil {
		log = slog.New(slog.NewTextHandler(io.Discard, nil))
	}
	return logger.WithSource(log, "memory/vector")
}