118 lines
3.5 KiB
Go
118 lines
3.5 KiB
Go
package store
|
||
|
||
import (
|
||
core "agent-system-core"
|
||
)
|
||
|
||
// searchKnowledgePG：PG 档知识检索（设计 D16）。
|
||
// 单查询混合：tsvector（'simple' 配置，词级）OR trigram LIKE 子串（中文兜底；
|
||
// LIKE 由 pg_trgm GIN 索引加速）。排序：ts_rank 优先，命中时间次之。
|
||
func (s *Store) searchKnowledgePG(owner string, q core.RetrievalQuery) ([]core.RetrievalHit, error) {
|
||
topK := q.TopK
|
||
if topK <= 0 {
|
||
topK = 5
|
||
}
|
||
text := q.Text
|
||
rows, err := s.db.Query(s.q(`
|
||
SELECT k.id,k.owner,k.project_id,k.source,k.title,k.ref,k.chunk,k.created_at,
|
||
ts_rank(k.tsv, websearch_to_tsquery('simple', ?)) AS r
|
||
FROM knowledge k
|
||
WHERE (k.owner=? OR k.owner='')
|
||
AND (? = '' OR k.tsv @@ websearch_to_tsquery('simple', ?) OR k.chunk LIKE ? OR k.title LIKE ?)
|
||
ORDER BY r DESC, k.created_at DESC
|
||
LIMIT ?`), text, owner, text, text, "%"+text+"%", "%"+text+"%", topK*2)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer rows.Close()
|
||
out := []core.RetrievalHit{}
|
||
for rows.Next() {
|
||
var d core.KnowledgeDoc
|
||
var own, ca string
|
||
var rank float64
|
||
if err := rows.Scan(&d.ID, &own, &d.ProjectID, &d.Source, &d.Title, &d.Ref, &d.Chunk, &ca, &rank); err != nil {
|
||
return nil, err
|
||
}
|
||
d.Owner = own
|
||
d.CreatedAt, _ = parseTime(ca)
|
||
if !knowledgeVisible(owner, q, d) {
|
||
continue
|
||
}
|
||
score := rank
|
||
if score == 0 {
|
||
score = 0.5 // 纯子串命中（trigram/中文兜底）给保底分
|
||
}
|
||
out = append(out, core.RetrievalHit{Doc: d, Score: score, Kind: "fulltext"})
|
||
if len(out) >= topK {
|
||
break
|
||
}
|
||
}
|
||
return out, rows.Err()
|
||
}
|
||
|
||
// searchKnowledgePGShared EXT-8 切片 2 共享版：他人 owner 的行 + 谓词后置判定
|
||
// （allow 回调由上层注入；排序与自检索同构）。
|
||
func (s *Store) searchKnowledgePGShared(requester string, q core.RetrievalQuery, pred knowledgePred) ([]core.RetrievalHit, error) {
|
||
return s.searchKnowledgePGSharedCandidates(requester, q, pred, true)
|
||
}
|
||
|
||
func (s *Store) searchKnowledgePGSharedCandidates(requester string, q core.RetrievalQuery, pred knowledgePred, bounded bool) ([]core.RetrievalHit, error) {
|
||
topK := q.TopK
|
||
if topK <= 0 {
|
||
topK = 5
|
||
}
|
||
text := q.Text
|
||
query := `
|
||
SELECT k.id,k.owner,k.project_id,k.source,k.title,k.ref,k.chunk,k.created_at,
|
||
ts_rank(k.tsv, websearch_to_tsquery('simple', ?)) AS r
|
||
FROM knowledge k
|
||
WHERE (k.owner<>? AND k.owner<>'')
|
||
AND (? = '' OR k.tsv @@ websearch_to_tsquery('simple', ?) OR k.chunk LIKE ? OR k.title LIKE ?)
|
||
ORDER BY r DESC, k.created_at DESC`
|
||
args := []any{text, requester, text, text, "%" + text + "%", "%" + text + "%"}
|
||
if bounded {
|
||
query += " LIMIT ?"
|
||
args = append(args, topK*2)
|
||
}
|
||
rows, err := s.db.Query(s.q(query), args...)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer rows.Close()
|
||
type rawPGCandidate struct {
|
||
doc core.KnowledgeDoc
|
||
rank float64
|
||
}
|
||
var rawCands []rawPGCandidate
|
||
for rows.Next() {
|
||
var d core.KnowledgeDoc
|
||
var own, ca string
|
||
var rank float64
|
||
if err := rows.Scan(&d.ID, &own, &d.ProjectID, &d.Source, &d.Title, &d.Ref, &d.Chunk, &ca, &rank); err != nil {
|
||
return nil, err
|
||
}
|
||
d.Owner = own
|
||
d.CreatedAt, _ = parseTime(ca)
|
||
rawCands = append(rawCands, rawPGCandidate{doc: d, rank: rank})
|
||
}
|
||
if err := rows.Err(); err != nil {
|
||
return nil, err
|
||
}
|
||
_ = rows.Close()
|
||
out := []core.RetrievalHit{}
|
||
for _, c := range rawCands {
|
||
if !pred(c.doc) {
|
||
continue
|
||
}
|
||
score := c.rank
|
||
if score == 0 {
|
||
score = 0.5
|
||
}
|
||
out = append(out, core.RetrievalHit{Doc: c.doc, Score: score, Kind: "fulltext"})
|
||
if len(out) >= topK {
|
||
break
|
||
}
|
||
}
|
||
return out, nil
|
||
}
|