145 lines
5.3 KiB
Go
145 lines
5.3 KiB
Go
package memory
|
||
|
||
import (
|
||
"context"
|
||
"math"
|
||
"sort"
|
||
|
||
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/embedding"
|
||
)
|
||
|
||
// 本文件是记忆整理的**候选召回**（design/23 §5.1 T0-2）：分区内两两相似度
|
||
// ≥ recallThreshold 的对进入候选。
|
||
//
|
||
// ⚠️ 口径（不可走样）：相似度**只用来召回候选对**，永不直接当合并判据。
|
||
// design/23 §1.2 的实据是「必合」的 [12]/[15]=0.8330 与「必不合」的
|
||
// [3]/[6]=0.7100 只差 0.12，且全库最高相似度的一对（0.9319）恰是跨作用域
|
||
// ⇒ 单阈值双向失效，判定必须靠包含判定 + 规则否决（tidy_rules.go）。
|
||
|
||
// tidyPair 是一条被召回的候选对。
|
||
type tidyPair struct {
|
||
a, b MemoryItem
|
||
score float64
|
||
// mode 记录本轮相似度来源（embedding / lexical），进证据与日志。
|
||
mode string
|
||
}
|
||
|
||
// recallPairs 在分区内召回候选对。条目少于 2 条时为空操作（无候选、不报错）。
|
||
//
|
||
// 两两比对是 O(n²·dim)：设计稿每分区扫描上限 5000 条（design §9），但真正
|
||
// 参与比对的最多 maxTidyCompareItems 条（最新优先），超限记 WARN 并登记到
|
||
// stats.Truncated —— 截断绝不静默。
|
||
//
|
||
// 返回的对按「分数降序 + id」确定性排序，使每分区候选上限的截断可复现
|
||
// （同一批数据两次运行产出同一批候选 ⇒ 幂等，design §4.1 A）。
|
||
func (t *Tidier) recallPairs(ctx context.Context, p tidyPartition, items []MemoryItem, stats *TidyStats) []tidyPair {
|
||
if len(items) < 2 {
|
||
return nil
|
||
}
|
||
if len(items) > maxTidyCompareItems {
|
||
stats.Truncated = append(stats.Truncated, p.label()+"(compare)")
|
||
t.log.Warn("memory tidy compare truncated by per-iteration guard",
|
||
"partition", p.label(), "items", len(items), "limit", maxTidyCompareItems)
|
||
items = items[:maxTidyCompareItems]
|
||
}
|
||
stats.Compared += len(items)
|
||
|
||
sims, mode := t.similarityFunc(ctx, p, items)
|
||
pairs := make([]tidyPair, 0)
|
||
for i := 0; i < len(items); i++ {
|
||
for j := i + 1; j < len(items); j++ {
|
||
score := sims(i, j)
|
||
if score < t.cfg.RecallThreshold {
|
||
continue
|
||
}
|
||
pairs = append(pairs, tidyPair{a: items[i], b: items[j], score: score, mode: mode})
|
||
}
|
||
}
|
||
stats.Pairs += len(pairs)
|
||
sort.Slice(pairs, func(i, j int) bool {
|
||
if pairs[i].score != pairs[j].score {
|
||
return pairs[i].score > pairs[j].score
|
||
}
|
||
if pairs[i].a.ID != pairs[j].a.ID {
|
||
return pairs[i].a.ID < pairs[j].a.ID
|
||
}
|
||
return pairs[i].b.ID < pairs[j].b.ID
|
||
})
|
||
return pairs
|
||
}
|
||
|
||
// similarityFunc 返回分区内的两两相似度函数与来源标签。
|
||
//
|
||
// 接入 embedder 时用它批量向量化后取余弦（与检索侧同一口径：余弦夹到
|
||
// [0,1]）。embedder 缺席或整批 Embed 失败时回退到零依赖的词项余弦
|
||
// —— 该回退保证整理在离线/无 embedding 配置时仍可跑（宁可比语义路径更保守），
|
||
// 且让测试无需网络即可确定性运行。两个来源的相似度**不可混用阈值**：
|
||
// 阈值只有一个（design §13 开放点 2 的 0.80），回退路径命中更少 ⇒ 更保守。
|
||
func (t *Tidier) similarityFunc(ctx context.Context, p tidyPartition, items []MemoryItem) (func(i, j int) float64, string) {
|
||
if t.deps.Embedder != nil {
|
||
vecs, err := t.embedItems(ctx, items)
|
||
if err == nil && len(vecs) == len(items) {
|
||
return func(i, j int) float64 {
|
||
return clampUnit(cosineSim(vecs[i], vecs[j]))
|
||
}, "embedding"
|
||
}
|
||
t.log.Warn("memory tidy embedding recall unavailable, falling back to lexical",
|
||
"partition", p.label(), "error", err)
|
||
}
|
||
return func(i, j int) float64 {
|
||
return lexicalSimilarity(items[i].Content, items[j].Content)
|
||
}, "lexical"
|
||
}
|
||
|
||
// embedItems 分批把分区内条目向量化。单批失败即整体失败（调用方回退到
|
||
// 词项余弦），不半途改用混合表示 —— 混用会让召回范围不可预测。
|
||
func (t *Tidier) embedItems(ctx context.Context, items []MemoryItem) ([]embedding.Vector, error) {
|
||
out := make([]embedding.Vector, 0, len(items))
|
||
for start := 0; start < len(items); start += tidyEmbedBatch {
|
||
end := min(start+tidyEmbedBatch, len(items))
|
||
texts := make([]string, 0, end-start)
|
||
for _, it := range items[start:end] {
|
||
texts = append(texts, it.Content)
|
||
}
|
||
vecs, err := t.deps.Embedder.Embed(ctx, texts)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if len(vecs) != len(texts) {
|
||
return nil, ErrTidyEmbedCount
|
||
}
|
||
out = append(out, vecs...)
|
||
}
|
||
return out, nil
|
||
}
|
||
|
||
// lexicalSimilarity 是零依赖的相似度回退：两条内容词项集合的余弦。
|
||
// 词项来自 extractTidyFacts（CJK 二元组 + 拉丁词 + 数值 + emoji 符号），
|
||
// 因此「同义不同措辞」的中文近似表述仍能召回，而「红牛 vs 冰美式」
|
||
// 这类并列事实只共享极少词项（相似度低或中等），靠包含判定兜住。
|
||
func lexicalSimilarity(a, b string) float64 {
|
||
fa, fb := extractTidyFacts(a), extractTidyFacts(b)
|
||
if len(fa.terms) == 0 || len(fb.terms) == 0 {
|
||
return 0
|
||
}
|
||
shared := 0
|
||
for term := range fa.terms {
|
||
if _, ok := fb.terms[term]; ok {
|
||
shared++
|
||
}
|
||
}
|
||
return float64(shared) / math.Sqrt(float64(len(fa.terms))*float64(len(fb.terms)))
|
||
}
|
||
|
||
// clampUnit 把余弦夹到 [0,1]：相似度是「召回门」的输入，负相关对
|
||
// （余弦 < 0）与零相关对在召回上等价。
|
||
func clampUnit(v float64) float64 {
|
||
if v < 0 {
|
||
return 0
|
||
}
|
||
if v > 1 {
|
||
return 1
|
||
}
|
||
return v
|
||
}
|