408 lines
16 KiB
Go
408 lines
16 KiB
Go
package store
|
||
|
||
import (
|
||
"context"
|
||
"database/sql"
|
||
"encoding/binary"
|
||
"errors"
|
||
"math"
|
||
"sort"
|
||
"strconv"
|
||
"strings"
|
||
|
||
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 {
|
||
return s.UpsertKnowledgeVectorExt(ctx, owner, docID, ref, "", "", vec)
|
||
}
|
||
|
||
// vectorLiteral 把 float32 数组编成 pgvector 字面量 '[a,b,c]'。
|
||
func vectorLiteral(vec []float32) string {
|
||
var sb strings.Builder
|
||
sb.WriteByte('[')
|
||
for i, v := range vec {
|
||
if i > 0 {
|
||
sb.WriteByte(',')
|
||
}
|
||
sb.WriteString(strconv.FormatFloat(float64(v), 'g', -1, 32))
|
||
}
|
||
sb.WriteByte(']')
|
||
return sb.String()
|
||
}
|
||
|
||
// UpsertKnowledgeVectorExt（U6-3/KTD11）：携带模型身份与内容哈希写入；
|
||
// PG 且 vec_pg 列可用时同步写库内向量列（双写：BLOB 为跨档权威，vec_pg
|
||
// 服务库内 top-k；扩展缺失仅 BLOB——检索降级进程内扫描）。
|
||
func (s *Store) UpsertKnowledgeVectorExt(ctx context.Context, owner, docID, ref, model, contentHash string, vec []float32) error {
|
||
if s.dialect.name() == "postgres" && s.ensureVecPG() {
|
||
_, err := s.db.ExecContext(ctx, `INSERT INTO knowledge_vectors(doc_id,owner,ref,dim,vec,model,content_hash,vec_pg,created_at) VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9)
|
||
ON CONFLICT(doc_id) DO UPDATE SET owner=excluded.owner, ref=excluded.ref, dim=excluded.dim, vec=excluded.vec, model=excluded.model, content_hash=excluded.content_hash, vec_pg=excluded.vec_pg, created_at=excluded.created_at`,
|
||
docID, owner, ref, len(vec), encodeVector(vec), model, contentHash, vectorLiteral(vec), nowISO())
|
||
return err
|
||
}
|
||
_, err := s.db.ExecContext(ctx, s.q(`INSERT INTO knowledge_vectors(doc_id,owner,ref,dim,vec,model,content_hash,created_at) VALUES(?,?,?,?,?,?,?,?)
|
||
ON CONFLICT(doc_id) DO UPDATE SET owner=excluded.owner, ref=excluded.ref, dim=excluded.dim, vec=excluded.vec, model=excluded.model, content_hash=excluded.content_hash, created_at=excluded.created_at`),
|
||
docID, owner, ref, len(vec), encodeVector(vec), model, contentHash, nowISO())
|
||
return err
|
||
}
|
||
|
||
func (s *Store) UpsertKnowledgeVectorAccountChecked(ctx context.Context, actor, docID, ref string, vec []float32, allowImplicitDefault bool) error {
|
||
return s.UpsertKnowledgeVectorExtAccountChecked(ctx, actor, docID, ref, "", "", vec, allowImplicitDefault)
|
||
}
|
||
|
||
// UpsertKnowledgeVectorExtAccountChecked requires an authoritative knowledge
|
||
// row owned by actor with the same ref. Conflict updates cannot take over a
|
||
// vector row whose owner/ref differs, even if legacy data is inconsistent.
|
||
func (s *Store) UpsertKnowledgeVectorExtAccountChecked(ctx context.Context, actor, docID, ref, model, contentHash string, vec []float32, allowImplicitDefault bool) error {
|
||
if err := ctx.Err(); err != nil {
|
||
return err
|
||
}
|
||
vecPG := s.dialect.name() == "postgres" && s.ensureVecPG()
|
||
return s.withEffectiveAccountWrite(actor, allowImplicitDefault, func(tx *sql.Tx, _ Account) error {
|
||
if err := ctx.Err(); err != nil {
|
||
return err
|
||
}
|
||
query := s.q(`SELECT owner,ref FROM knowledge WHERE id=? AND owner=? AND ref=?`)
|
||
if s.Dialect() == "postgres" {
|
||
query += ` FOR UPDATE`
|
||
}
|
||
var owner, storedRef string
|
||
if err := tx.QueryRowContext(ctx, query, docID, actor, ref).Scan(&owner, &storedRef); errors.Is(err, sql.ErrNoRows) {
|
||
return ErrNotFound
|
||
} else if err != nil {
|
||
return err
|
||
}
|
||
if owner != actor || storedRef != ref {
|
||
return ErrNotFound
|
||
}
|
||
var res sql.Result
|
||
var err error
|
||
if vecPG {
|
||
res, err = tx.ExecContext(ctx, s.q(`INSERT INTO knowledge_vectors(doc_id,owner,ref,dim,vec,model,content_hash,vec_pg,created_at) VALUES(?,?,?,?,?,?,?,?,?)
|
||
ON CONFLICT(doc_id) DO UPDATE SET owner=excluded.owner, ref=excluded.ref, dim=excluded.dim, vec=excluded.vec, model=excluded.model, content_hash=excluded.content_hash, vec_pg=excluded.vec_pg, created_at=excluded.created_at
|
||
WHERE knowledge_vectors.owner=excluded.owner AND knowledge_vectors.ref=excluded.ref`),
|
||
docID, actor, ref, len(vec), encodeVector(vec), model, contentHash, vectorLiteral(vec), nowISO())
|
||
} else {
|
||
res, err = tx.ExecContext(ctx, s.q(`INSERT INTO knowledge_vectors(doc_id,owner,ref,dim,vec,model,content_hash,created_at) VALUES(?,?,?,?,?,?,?,?)
|
||
ON CONFLICT(doc_id) DO UPDATE SET owner=excluded.owner, ref=excluded.ref, dim=excluded.dim, vec=excluded.vec, model=excluded.model, content_hash=excluded.content_hash, created_at=excluded.created_at
|
||
WHERE knowledge_vectors.owner=excluded.owner AND knowledge_vectors.ref=excluded.ref`),
|
||
docID, actor, ref, len(vec), encodeVector(vec), model, contentHash, nowISO())
|
||
}
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if n, err := res.RowsAffected(); err != nil {
|
||
return err
|
||
} else if n != 1 {
|
||
return ErrNotFound
|
||
}
|
||
return nil
|
||
})
|
||
}
|
||
|
||
// ensureVecPG 探测 vec_pg 列是否可用（pgvector 扩展在场；一次性缓存）。
|
||
func (s *Store) ensureVecPG() bool {
|
||
s.vecPGOnce.Do(func() {
|
||
var n int
|
||
if err := s.db.QueryRow(`SELECT COUNT(*) FROM information_schema.columns WHERE table_schema=current_schema() AND table_name='knowledge_vectors' AND column_name='vec_pg'`).Scan(&n); err == nil {
|
||
s.vecPG = n > 0
|
||
}
|
||
})
|
||
return s.vecPG
|
||
}
|
||
|
||
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
|
||
}
|
||
|
||
func (s *Store) DeleteKnowledgeVectorsAccountChecked(ctx context.Context, actor, ref string, allowImplicitDefault bool) error {
|
||
return s.withEffectiveAccountWrite(actor, allowImplicitDefault, func(tx *sql.Tx, _ Account) error {
|
||
if err := ctx.Err(); err != nil {
|
||
return err
|
||
}
|
||
_, err := tx.ExecContext(ctx, s.q(`DELETE FROM knowledge_vectors WHERE ref=? AND owner=?`), ref, actor)
|
||
return err
|
||
})
|
||
}
|
||
|
||
// SearchKnowledgeVectors 余弦扫描：候选 = 租户本人 + 共享（owner=”），dim 匹配，
|
||
// 再按 RetrievalQuery 的项目/来源过滤。
|
||
// SearchKnowledgeVectorsShared EXT-8 余项：他人 owner 的向量行 + 谓词后置判定
|
||
// （allow 由上层注入适配器；无 ProjectID 行拒绝）。结构同自检索（余弦扫描）。
|
||
func (s *Store) SearchKnowledgeVectorsShared(ctx context.Context, requester string, vec []float32, topK int, filter core.RetrievalQuery, allow func(docOwner, projectID core.ID) bool) ([]core.RetrievalHit, error) {
|
||
return s.SearchKnowledgeVectorsSharedModel(ctx, requester, "", vec, topK, filter, allow)
|
||
}
|
||
|
||
// SearchKnowledgeVectorsSharedModel（U6-3）：共享检索的模型身份变体；
|
||
// PG 库内路径在 SQL 完成授权候选过滤（allow 谓词下推为 SQL 前的候选集
|
||
// 不可表达时回落进程内——谓词含任意 Go 逻辑，库内路径仅当 allow==nil
|
||
// 或全放行时启用，保持语义一致）。
|
||
func (s *Store) SearchKnowledgeVectorsSharedModel(ctx context.Context, requester, model string, vec []float32, topK int, filter core.RetrievalQuery, allow func(docOwner, projectID core.ID) bool) ([]core.RetrievalHit, error) {
|
||
if topK <= 0 {
|
||
topK = 5
|
||
}
|
||
if s.dialect.name() == "postgres" && s.ensureVecPG() && allowAll(allow) {
|
||
if hits, err := s.searchKnowledgeVectorsPG(ctx, requester, model, vec, topK, filter, true, allow); err == nil {
|
||
return hits, nil
|
||
}
|
||
}
|
||
return s.SearchKnowledgeVectorsSharedAuthorized(ctx, requester, model, vec, topK, filter, func(d core.KnowledgeDoc) bool { return allow == nil || allow(d.Owner, d.ProjectID) })
|
||
}
|
||
|
||
// SearchKnowledgeVectorsSharedAuthorized applies per-document authorization to
|
||
// shared vector candidates before ranking, including their actual stored ID.
|
||
func (s *Store) SearchKnowledgeVectorsSharedAuthorized(ctx context.Context, requester, model string, vec []float32, topK int, filter core.RetrievalQuery, allow func(core.KnowledgeDoc) bool) ([]core.RetrievalHit, error) {
|
||
if topK <= 0 {
|
||
topK = 5
|
||
}
|
||
return s.searchKnowledgeVectorsSharedAuthorizedScan(ctx, requester, model, vec, topK, filter, allow)
|
||
}
|
||
|
||
// allowAll 判定谓词是否恒真（nil 或恒放行）——库内路径的启用条件。
|
||
func allowAll(allow func(docOwner, projectID core.ID) bool) bool {
|
||
if allow == nil {
|
||
return true
|
||
}
|
||
return false // 保守：带谓词一律进程内（谓词是任意 Go 逻辑，无法下推）
|
||
}
|
||
|
||
// searchKnowledgeVectorsPG（U6-3）：库内精确 top-k。shared 决定候选集
|
||
// （本人+共享 vs 他人共享），dim 过滤先行、模型身份精确匹配、
|
||
// 相似度 = 1 - (vec_pg <=> query)。空结果与错误都触发调用方降级判定
|
||
// （错误降级、空直接返回）。
|
||
func (s *Store) searchKnowledgeVectorsPG(ctx context.Context, owner, model string, vec []float32, topK int, filter core.RetrievalQuery, shared bool, allow func(docOwner, projectID core.ID) bool) ([]core.RetrievalHit, error) {
|
||
where := "v.owner=$1 OR v.owner=''"
|
||
if shared {
|
||
where = "v.owner<>$1 AND v.owner<>''"
|
||
}
|
||
q := `
|
||
SELECT k.id,k.owner,k.project_id,k.source,k.title,k.ref,k.chunk,k.created_at,
|
||
1 - (v.vec_pg <=> ($5)::vector) AS sim
|
||
FROM knowledge_vectors v JOIN knowledge k ON k.id = v.doc_id
|
||
WHERE (` + where + `) AND v.dim=$2`
|
||
args := []any{owner, len(vec), topK, model, vectorLiteral(vec)}
|
||
if model != "" {
|
||
q += ` AND v.model=$4`
|
||
} else {
|
||
q += ` AND v.model=''`
|
||
}
|
||
if filter.ProjectID != "" {
|
||
q += ` AND k.project_id='` + filter.ProjectID + `'`
|
||
}
|
||
q += ` ORDER BY v.vec_pg <=> ($5)::vector LIMIT $3`
|
||
rows, err := s.db.QueryContext(ctx, q, args...)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer rows.Close()
|
||
out := []core.RetrievalHit{}
|
||
for rows.Next() {
|
||
var d core.KnowledgeDoc
|
||
var own, ca string
|
||
var sim float64
|
||
if err := rows.Scan(&d.ID, &own, &d.ProjectID, &d.Source, &d.Title, &d.Ref, &d.Chunk, &ca, &sim); err != nil {
|
||
return nil, err
|
||
}
|
||
d.Owner = own
|
||
d.CreatedAt, _ = parseTime(ca)
|
||
if filter.ProjectID != "" && d.ProjectID != filter.ProjectID {
|
||
continue
|
||
}
|
||
if len(filter.Sources) > 0 && !containsStr(filter.Sources, d.Source) {
|
||
continue
|
||
}
|
||
out = append(out, core.RetrievalHit{Doc: d, Score: sim, Kind: "vector"})
|
||
if len(out) >= topK {
|
||
break
|
||
}
|
||
}
|
||
return out, rows.Err()
|
||
}
|
||
|
||
// searchKnowledgeVectorsSharedScan 原进程内共享扫描（降级/谓词路径）。
|
||
func (s *Store) searchKnowledgeVectorsSharedScan(ctx context.Context, requester, model string, vec []float32, topK int, filter core.RetrievalQuery, allow func(docOwner, projectID core.ID) bool) ([]core.RetrievalHit, error) {
|
||
return s.searchKnowledgeVectorsSharedAuthorizedScan(ctx, requester, model, vec, topK, filter, func(d core.KnowledgeDoc) bool { return allow == nil || allow(d.Owner, d.ProjectID) })
|
||
}
|
||
func (s *Store) searchKnowledgeVectorsSharedAuthorizedScan(ctx context.Context, requester, model string, vec []float32, topK int, filter core.RetrievalQuery, allow func(core.KnowledgeDoc) bool) ([]core.RetrievalHit, error) {
|
||
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<>? AND v.owner<>'') AND v.dim=?`
|
||
args := []any{requester, len(vec)}
|
||
if model != "" {
|
||
q += ` AND v.model=?`
|
||
args = append(args, model)
|
||
}
|
||
rows, err := s.db.QueryContext(ctx, s.q(q), args...)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer rows.Close()
|
||
type rawVectorCandidate struct {
|
||
doc core.KnowledgeDoc
|
||
blob []byte
|
||
}
|
||
var rawCands []rawVectorCandidate
|
||
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)
|
||
rawCands = append(rawCands, rawVectorCandidate{doc: d, blob: blob})
|
||
}
|
||
if err := rows.Err(); err != nil {
|
||
return nil, err
|
||
}
|
||
_ = rows.Close()
|
||
|
||
type cand struct {
|
||
hit core.RetrievalHit
|
||
sim float64
|
||
}
|
||
var cands []cand
|
||
for _, raw := range rawCands {
|
||
d := raw.doc
|
||
if d.ProjectID == "" || !allow(d) {
|
||
continue
|
||
}
|
||
if filter.ProjectID != "" && d.ProjectID != filter.ProjectID {
|
||
continue
|
||
}
|
||
if len(filter.Sources) > 0 && !containsStr(filter.Sources, d.Source) {
|
||
continue
|
||
}
|
||
sim := cosine(vec, decodeVector(raw.blob))
|
||
cands = append(cands, cand{hit: core.RetrievalHit{Doc: d, Score: sim, Kind: "vector"}, sim: sim})
|
||
}
|
||
sort.SliceStable(cands, func(i, j int) bool { return cands[i].sim > cands[j].sim })
|
||
out := []core.RetrievalHit{}
|
||
for _, c := range cands {
|
||
out = append(out, c.hit)
|
||
if len(out) >= topK {
|
||
break
|
||
}
|
||
}
|
||
return out, nil
|
||
}
|
||
|
||
func (s *Store) SearchKnowledgeVectors(ctx context.Context, owner string, vec []float32, topK int, filter core.RetrievalQuery) ([]core.RetrievalHit, error) {
|
||
return s.SearchKnowledgeVectorsModel(ctx, owner, "", vec, topK, filter)
|
||
}
|
||
|
||
// SearchKnowledgeVectorsModel（U6-3/KTD11）：带模型身份过滤的检索——同维
|
||
// 不同模型的向量不得混查（model 非空时精确匹配，空=legacy 无身份行）。
|
||
// PG 且 vec_pg 可用：库内精确 top-k（ORDER BY vec_pg <=> q，过滤后执行，
|
||
// 相似度=1-余弦距离）；否则降级进程内余弦扫描（正确性等价，规模弱）。
|
||
func (s *Store) SearchKnowledgeVectorsModel(ctx context.Context, owner, model string, vec []float32, topK int, filter core.RetrievalQuery) ([]core.RetrievalHit, error) {
|
||
if topK <= 0 {
|
||
topK = 5
|
||
}
|
||
if s.dialect.name() == "postgres" && s.ensureVecPG() {
|
||
hits, err := s.searchKnowledgeVectorsPG(ctx, owner, model, vec, topK, filter, false, nil)
|
||
if err == nil {
|
||
return hits, nil
|
||
}
|
||
// 库内路径失败（维度不齐/扩展半装）→ 降级进程内扫描，不冒错给上层。
|
||
}
|
||
return s.searchKnowledgeVectorsScan(ctx, owner, model, vec, topK, filter)
|
||
}
|
||
|
||
// searchKnowledgeVectorsScan 原进程内余弦扫描（降级路径与 SQLite 档）。
|
||
func (s *Store) searchKnowledgeVectorsScan(ctx context.Context, owner, model string, vec []float32, topK int, filter core.RetrievalQuery) ([]core.RetrievalHit, error) {
|
||
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=?`
|
||
args := []any{owner, len(vec)}
|
||
if model != "" {
|
||
q += ` AND v.model=?`
|
||
args = append(args, model)
|
||
}
|
||
rows, err := s.db.QueryContext(ctx, s.q(q), args...)
|
||
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
|
||
}
|