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