weichen
4fdbfbc9a2
Details (original message): Add the vector RAG pipeline: paragraph chunker, OpenAI-compatible embeddings and rerank clients, dual-tier vector store with cosine search, and RRF hybrid retrieval
113 lines
3.4 KiB
Go
113 lines
3.4 KiB
Go
package store
|
||
|
||
import (
|
||
"context"
|
||
"encoding/binary"
|
||
"math"
|
||
|
||
core "agent-system-core"
|
||
)
|
||
|
||
// ---- 向量索引（4.12 向量管线）----
|
||
// 存储形态双档一致：BLOB（little-endian float32 数组），检索为进程内余弦扫描。
|
||
// 规模注记：当前实现面向 ≤ 数万块/实例；PG 档 pgvector+HNSW 是既定优化路径
|
||
//（换实现不动管线与契约，见设计 D13）。
|
||
|
||
func encodeVector(vec []float32) []byte {
|
||
buf := make([]byte, 4*len(vec))
|
||
for i, v := range vec {
|
||
binary.LittleEndian.PutUint32(buf[i*4:], math.Float32bits(v))
|
||
}
|
||
return buf
|
||
}
|
||
|
||
func decodeVector(blob []byte) []float32 {
|
||
n := len(blob) / 4
|
||
out := make([]float32, n)
|
||
for i := 0; i < n; i++ {
|
||
out[i] = math.Float32frombits(binary.LittleEndian.Uint32(blob[i*4:]))
|
||
}
|
||
return out
|
||
}
|
||
|
||
// cosine 余弦相似度（零向量返回 0）。
|
||
func cosine(a, b []float32) float64 {
|
||
if len(a) != len(b) {
|
||
return 0
|
||
}
|
||
var dot, na, nb float64
|
||
for i := range a {
|
||
dot += float64(a[i]) * float64(b[i])
|
||
na += float64(a[i]) * float64(a[i])
|
||
nb += float64(b[i]) * float64(b[i])
|
||
}
|
||
if na == 0 || nb == 0 {
|
||
return 0
|
||
}
|
||
return dot / (math.Sqrt(na) * math.Sqrt(nb))
|
||
}
|
||
|
||
func (s *Store) UpsertKnowledgeVector(ctx context.Context, owner, docID, ref string, vec []float32) error {
|
||
_, err := s.db.ExecContext(ctx, s.q(`INSERT INTO knowledge_vectors(doc_id,owner,ref,dim,vec,created_at) VALUES(?,?,?,?,?,?)
|
||
ON CONFLICT(doc_id) DO UPDATE SET owner=excluded.owner, ref=excluded.ref, dim=excluded.dim, vec=excluded.vec, created_at=excluded.created_at`),
|
||
docID, owner, ref, len(vec), encodeVector(vec), nowISO())
|
||
return err
|
||
}
|
||
|
||
func (s *Store) DeleteKnowledgeVectors(ctx context.Context, owner, ref string) error {
|
||
_, err := s.db.ExecContext(ctx, s.q(`DELETE FROM knowledge_vectors WHERE ref=? AND owner=?`), ref, owner)
|
||
return err
|
||
}
|
||
|
||
// SearchKnowledgeVectors 余弦扫描：候选 = 租户本人 + 共享（owner=”），dim 匹配，
|
||
// 再按 RetrievalQuery 的项目/来源过滤。
|
||
func (s *Store) SearchKnowledgeVectors(ctx context.Context, owner string, vec []float32, topK int, filter core.RetrievalQuery) ([]core.RetrievalHit, error) {
|
||
if topK <= 0 {
|
||
topK = 5
|
||
}
|
||
rows, err := s.db.QueryContext(ctx, s.q(`
|
||
SELECT k.id,k.owner,k.project_id,k.source,k.title,k.ref,k.chunk,k.created_at, v.vec
|
||
FROM knowledge_vectors v JOIN knowledge k ON k.id = v.doc_id
|
||
WHERE (v.owner=? OR v.owner='') AND v.dim=?`), owner, len(vec))
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer rows.Close()
|
||
type cand struct {
|
||
hit core.RetrievalHit
|
||
sim float64
|
||
}
|
||
var cands []cand
|
||
for rows.Next() {
|
||
var d core.KnowledgeDoc
|
||
var own, ca string
|
||
var blob []byte
|
||
if err := rows.Scan(&d.ID, &own, &d.ProjectID, &d.Source, &d.Title, &d.Ref, &d.Chunk, &ca, &blob); err != nil {
|
||
return nil, err
|
||
}
|
||
d.Owner = own
|
||
d.CreatedAt, _ = parseTime(ca)
|
||
if !knowledgeVisible(owner, filter, d) {
|
||
continue
|
||
}
|
||
sim := cosine(vec, decodeVector(blob))
|
||
cands = append(cands, cand{hit: core.RetrievalHit{Doc: d, Score: sim, Kind: "vector"}, sim: sim})
|
||
}
|
||
if err := rows.Err(); err != nil {
|
||
return nil, err
|
||
}
|
||
// 简单选择排序取 TopK（候选规模有限；HNSW 归 PG 档优化）。
|
||
out := make([]core.RetrievalHit, 0, topK)
|
||
for len(out) < topK && len(cands) > 0 {
|
||
best := 0
|
||
for i := 1; i < len(cands); i++ {
|
||
if cands[i].sim > cands[best].sim {
|
||
best = i
|
||
}
|
||
}
|
||
out = append(out, cands[best].hit)
|
||
cands = append(cands[:best], cands[best+1:]...)
|
||
}
|
||
return out, nil
|
||
}
|