Cyrene_v2/internal/runtime/memory/tidy_recall.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）。
//
// 召回来源（embedding / 词项余弦）与「来源不可用时怎么办」被拆成两层：
// recallPairs 只报事实（不可用就回错误），**回退还是上报由调用方的
// tidyRecall 策略决定**（tidy.go 的生产轮回退、tidy_dryrun.go 的试算轮上报）。

const (
	// tidyRecallEmbedding 是本轮相似度来自语义 embedding。
	tidyRecallEmbedding = "embedding"
	// tidyRecallLexical 是本轮相似度来自零依赖的词项余弦。
	tidyRecallLexical = "lexical"
)

// tidyPair 是一条被召回的候选对。
type tidyPair struct {
	a, b  MemoryItem
	score float64
	// mode 记录本轮相似度来源（embedding / lexical），进证据与日志。
	mode string
}

// recallPairs 在分区内召回候选对。条目少于 2 条时为空操作（无候选、不报错）。
//
// 相似度来源不可用时（embedder 报错、向量条数错位）返回错误 —— 回退决策不在
// 这里：生产轮把它降级成词项余弦，试算轮如实上报（见 tidyRecall 策略）。
//
// 返回的对按「分数降序 + id」确定性排序，使每分区候选上限的截断可复现
// （同一批数据两次运行产出同一批候选 ⇒ 幂等，design §4.1 A）。
func (t *Tidier) recallPairs(ctx context.Context, p tidyPartition, items []MemoryItem, stats *TidyStats, budget tidyBudget) ([]tidyPair, error) {
	if len(items) < 2 {
		return nil, nil
	}
	if t.deps.Embedder == nil {
		return t.collectPairs(p, items, t.cfg.RecallThreshold,
			lexicalSimFunc(items), tidyRecallLexical, stats, budget), nil
	}
	vecs, err := t.embedItems(ctx, items)
	if err != nil {
		return nil, err
	}
	if len(vecs) != len(items) {
		return nil, ErrTidyEmbedCount
	}
	return t.collectPairs(p, items, t.cfg.RecallThreshold,
		vectorSimFunc(vecs), tidyRecallEmbedding, stats, budget), nil
}

// limitCompareItems 把参与比对的条目截断到性能护栏并计入 Compared。
// 两两比对是 O(n²·dim)：设计稿每分区扫描上限 5000 条（design §9），但真正参与
// 比对的最多 cfg.CompareLimit 条（最新优先，默认 2000，MEMORY_TIDY_COMPARE_LIMIT），
// 超限记 WARN 并登记到 stats.Truncated —— 截断绝不静默。
// 空/单条分区不计数（本来就没有对可比）。
func (t *Tidier) limitCompareItems(p tidyPartition, items []MemoryItem, stats *TidyStats) []MemoryItem {
	if len(items) < 2 {
		return items
	}
	if len(items) > t.cfg.CompareLimit {
		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", t.cfg.CompareLimit)
		items = items[:t.cfg.CompareLimit]
	}
	stats.Compared += len(items)
	return items
}

// collectTidyPairs 是召回内核：把相似度 ≥ 阈值的对收集起来，按分数降序
// （再按 id）确定性排序 —— 排序让「每分区候选上限」的截断可复现。
//
// 它是自由函数（入库前审查也复用同一个内核，design/23 §15 D3-3「不许另造一套
// 相似度实现」），因此**不认识分区与日志器**：cut=true 表示循环被时间预算提前
// 收手，由调用方决定怎么如实登记（整理侧见 Tidier.collectPairs）。
// budget 为零值（无 deadline）时行为与从前逐字一致。
//
// budget 用尽时已找到的对是**真结果**（不是编的），照常返回 —— 少算了多少由
// 调用方如实标记，绝不静默。
func collectTidyPairs(items []MemoryItem, threshold float64, sim func(i, j int) float64,
	mode string, stats *TidyStats, budget tidyBudget) ([]tidyPair, bool) {
	pairs := make([]tidyPair, 0)
	ticks := 0
	cut := false
pairLoop:
	for i := 0; i < len(items); i++ {
		for j := i + 1; j < len(items); j++ {
			if ticks++; ticks >= tidyBudgetCheckEveryPairs {
				ticks = 0
				if budget.exhausted() {
					cut = true
					break pairLoop
				}
			}
			score := sim(i, j)
			if score < threshold {
				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, cut
}

// collectPairs 是整理侧对内核对调用：它在内核之外补上「预算截断」的**诚实登记**
// （分区标签进响应的 truncated 数组 + 一条 WARN）。截断绝不静默是这一层的职责，
// 因此内核保持纯净（preinsert 审查不需要这套标签）。
func (t *Tidier) collectPairs(p tidyPartition, items []MemoryItem, threshold float64,
	sim func(i, j int) float64, mode string, stats *TidyStats, budget tidyBudget) []tidyPair {
	pairs, cut := collectTidyPairs(items, threshold, sim, mode, stats, budget)
	if cut {
		stats.Truncated = append(stats.Truncated, budgetTruncationLabel(p))
		t.log.Warn("memory tidy pair scan cut by time budget",
			"partition", p.label(), "items", len(items), "pairs", len(pairs))
	}
	return pairs
}

// lexicalSimFunc 是零依赖的相似度函数：两条内容词项集合的余弦。
func lexicalSimFunc(items []MemoryItem) func(i, j int) float64 {
	return func(i, j int) float64 {
		return lexicalSimilarity(items[i].Content, items[j].Content)
	}
}

// vectorSimFunc 是语义相似度函数（与检索侧同一口径：余弦夹到 [0,1]）。
func vectorSimFunc(vecs []embedding.Vector) func(i, j int) float64 {
	return func(i, j int) float64 {
		return clampUnit(cosineSim(vecs[i], vecs[j]))
	}
}

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

// tidyRecallMode 报告本轮相似度的来源，便于运维判断阈值是否在预期口径上。
func tidyRecallMode(hasEmbedder bool) string {
	if hasEmbedder {
		return tidyRecallEmbedding
	}
	return tidyRecallLexical
}