491 lines
19 KiB
Go
491 lines
19 KiB
Go
package memory
|
||
|
||
import (
|
||
"context"
|
||
"errors"
|
||
"fmt"
|
||
"log/slog"
|
||
"sort"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
|
||
"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
|
||
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/embedding"
|
||
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/memory/projectindex"
|
||
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/memory/vector"
|
||
)
|
||
|
||
// 检索器默认值与可调常量。
|
||
const (
|
||
// defaultKeywordWeight 是关键词（FTS/TF-IDF）
|
||
// 相关性在综合得分中的相对权重。
|
||
defaultKeywordWeight = 1.0
|
||
|
||
// defaultSemanticWeight 是语义向量相关性
|
||
// 在综合得分中的相对权重。仅在配置了 embedder 时才参与计算。
|
||
defaultSemanticWeight = 1.0
|
||
|
||
// defaultRecencyHalfLifeDays 是记忆的
|
||
// 新近度权重减半所需的天数，用于偏好近期被触碰的记忆（design §2
|
||
// 衰减）。
|
||
defaultRecencyHalfLifeDays = 30.0
|
||
|
||
// defaultTopK 是 Retrieve 或 Inject 以
|
||
// 非正 topK 调用时返回的结果数。
|
||
defaultTopK = 10
|
||
|
||
// hoursPerDay 把天数换算为 time.Duration。
|
||
hoursPerDay = 24
|
||
)
|
||
|
||
// Result 是一次检索命中：记忆条目，外加解释其排名的
|
||
// 相关性 × 衰减综合得分。得分越高越相关。
|
||
type Result struct {
|
||
Item MemoryItem
|
||
Score float64
|
||
// SourceConversation 是条目提取自哪次会话，用于
|
||
// 溯源（design §2）。裁决 04.4.3 起为结构化 JSON（session_id/message_id，
|
||
// 见 ParseSourceConversation），由持久化层盖真实 id。它是
|
||
// Item.SourceConversation 的镜像，调用方无需深入
|
||
// Item 即可回溯一次命中。
|
||
SourceConversation string
|
||
// Project 是链接的项目索引条目（设计 §3 项目分层），仅当
|
||
// 检索器接入了项目索引存储、且该条目是携带 ProjectID 的
|
||
// 项目作用域记忆时才填充。它让调用方为一条检索结果
|
||
// 呈现「项目索引 → 项目描述 → 章节内容」链路。
|
||
Project *projectindex.Project
|
||
}
|
||
|
||
// Retriever 对某角色可见的长期记忆 —— 它自己的条目
|
||
// 加共享的用户/项目作用域（design §3/§4，见 visibleCandidates）——
|
||
// 按相关性（关键词，以及配置了 embedder 或向量索引时的
|
||
// 语义向量）加权衰减（decay_weight × 时间衰减）排序后返回
|
||
// 前 K 条。它并发安全（索引装配状态与一次性回填有锁保护）；排名在
|
||
// 每次调用时从存储重算，因此结果始终反映持久化的
|
||
// 状态。
|
||
type Retriever struct {
|
||
store *Store
|
||
embedder embedding.Embedder
|
||
vindex vector.Index
|
||
projectStore *projectindex.Store
|
||
keywordWeight float64
|
||
semanticWeight float64
|
||
recencyHalfLife time.Duration
|
||
log *slog.Logger
|
||
|
||
// seedMu 保护一次性索引回填状态与回退路径的一次性告警位
|
||
// （seeded / fallbackWarned / budgetWarned）：回填每进程最多一次，
|
||
// 语义回退与其预算超限的告警也只报一次，避免每个 Retrieve
|
||
// 都写日志。
|
||
seedMu sync.Mutex
|
||
seeded bool
|
||
fallbackWarned bool
|
||
budgetWarned bool
|
||
|
||
// fallbackBudget 是内存余弦回退把候选正文向量化时允许花掉的
|
||
// 时间上限（见 WithFallbackBudget）。非正时取 defaultFallbackBudget。
|
||
fallbackBudget time.Duration
|
||
|
||
// fallbackMu 保护回退路径的向量缓存指针，使它可按需初始化
|
||
// （缓存本体自带锁，见 vectorCache）。
|
||
fallbackMu sync.Mutex
|
||
fallbackCache *vectorCache
|
||
}
|
||
|
||
// RetrieverOption 定制由 NewRetriever 构建的 Retriever。
|
||
type RetrieverOption func(*Retriever)
|
||
|
||
// WithEmbedder 设置用于内存内向量相似度计算的语义 embedder。
|
||
// 为 nil（默认）时，检索退化为仅关键词匹配
|
||
// 加衰减。若同时经 WithVectorIndex 接入了向量索引，则
|
||
// 索引优先：它是持久化语义路径（design §3
|
||
// 可选增强）。索引没有作答时（索引空、无候选命中或索引报错）
|
||
// 语义分回退到此内存路径，而绝不因「接了索引」就交出一片
|
||
// 空语义分。
|
||
func WithEmbedder(e embedding.Embedder) RetrieverOption {
|
||
return func(r *Retriever) {
|
||
r.embedder = e
|
||
}
|
||
}
|
||
|
||
// WithRetrieverLogger 设置检索器的 logger，使语义路径的降级
|
||
// （索引报错、回退内存余弦）可被发现（AGENTS §10 日志必须标来源）。
|
||
// 未设置时默认用该 store 的 logger（来源 "memory"），因此没有
|
||
// 额外接线的调用方也不会静默降级。
|
||
func WithRetrieverLogger(log *slog.Logger) RetrieverOption {
|
||
return func(r *Retriever) {
|
||
r.log = logger.WithSource(log, "memory/retriever")
|
||
}
|
||
}
|
||
|
||
// WithVectorIndex 把持久化的 vector.Index 接入检索。设置后，
|
||
// 语义相关性由 idx.Search 在已存向量上作答，而非
|
||
// 内存内向量化，因此配置的 pgvector/SQLite 向量存储被
|
||
// 真正启用（design §3「真正启用向量检索」）。为 nil（默认）时，
|
||
// 除非设置了 WithEmbedder，否则语义路径跳过。
|
||
//
|
||
// 它同时把索引接到该 store 的**写入路径**（见 Store.indexItem），使
|
||
// 「写入时 embed + add」不再是调用方可能漏做的一步：条目写进
|
||
// memory_items 的同时也写进索引。索引中只有全局向量，检索侧
|
||
// 只认本次候选集内的命中，并对空索引做一次性回填（见
|
||
// Retriever.ensureIndexSeeded）。
|
||
func WithVectorIndex(idx vector.Index) RetrieverOption {
|
||
return func(r *Retriever) {
|
||
r.vindex = idx
|
||
}
|
||
}
|
||
|
||
// WithProjectIndex 接入 projectindex.Store，使检索能把链接的
|
||
// 项目索引条目（设计 §3 项目分层）附加到每条项目作用域命中上。为 nil
|
||
// （默认）时，Result.Project 不设置。
|
||
func WithProjectIndex(s *projectindex.Store) RetrieverOption {
|
||
return func(r *Retriever) {
|
||
r.projectStore = s
|
||
}
|
||
}
|
||
|
||
// WithKeywordWeight 设置关键词相关性的相对权重。非正值
|
||
// 被忽略，保留默认值。
|
||
func WithKeywordWeight(w float64) RetrieverOption {
|
||
return func(r *Retriever) {
|
||
if w > 0 {
|
||
r.keywordWeight = w
|
||
}
|
||
}
|
||
}
|
||
|
||
// WithSemanticWeight 设置语义相关性的相对权重。非正值
|
||
// 被忽略，保留默认值。
|
||
func WithSemanticWeight(w float64) RetrieverOption {
|
||
return func(r *Retriever) {
|
||
if w > 0 {
|
||
r.semanticWeight = w
|
||
}
|
||
}
|
||
}
|
||
|
||
// WithRecencyHalfLife 设置记忆的新近度权重减半所需的天数。
|
||
// 非正值被忽略，保留默认值。
|
||
func WithRecencyHalfLife(days float64) RetrieverOption {
|
||
return func(r *Retriever) {
|
||
if days > 0 {
|
||
r.recencyHalfLife = time.Duration(days * float64(hoursPerDay*time.Hour))
|
||
}
|
||
}
|
||
}
|
||
|
||
// WithFallbackBudget 设置内存余弦回退为「本次候选」向量化时允许花掉的
|
||
// 时间上限。非正值被忽略，保留默认值（defaultFallbackBudget）。
|
||
//
|
||
// 它是性能护栏而不是功能开关：持久索引正常作答时这条路径根本不执行，
|
||
// 一旦执行（索引确实缺候选向量），它把「整库候选依次向量化」的代价
|
||
// 封在预算内 —— 超预算就用手上已有的向量打分，绝不为了补齐语义分
|
||
// 让一个回合继续等下去（真机事故：302 条候选 = 39.3 秒，见
|
||
// retrieve_fallback.go 的文件注释）。测试用短预算来验证这条护栏。
|
||
func WithFallbackBudget(d time.Duration) RetrieverOption {
|
||
return func(r *Retriever) {
|
||
if d > 0 {
|
||
r.fallbackBudget = d
|
||
}
|
||
}
|
||
}
|
||
|
||
// NewRetriever 在 store 之上构建 Retriever。nil store 非法，会在
|
||
// 调用时 panic；调用方应传入非 nil 的 store。
|
||
//
|
||
// 接入向量索引（WithVectorIndex）时，它随即将该索引接到 store 的
|
||
// 写入路径，使此后每条落库的记忆都被向量化并写入索引 —— 索引与
|
||
// 条目因此不会只同步一半（审计 P0-3：装配了索引却没有任何
|
||
// Index.Add 调用，语义检索恒空）。logger 未显式给出时沿用 store 的
|
||
// logger，使降级可见。
|
||
func NewRetriever(store *Store, opts ...RetrieverOption) *Retriever {
|
||
r := &Retriever{
|
||
store: store,
|
||
keywordWeight: defaultKeywordWeight,
|
||
semanticWeight: defaultSemanticWeight,
|
||
recencyHalfLife: time.Duration(defaultRecencyHalfLifeDays * float64(hoursPerDay*time.Hour)),
|
||
log: retrieverLog(store),
|
||
}
|
||
for _, opt := range opts {
|
||
opt(r)
|
||
}
|
||
if store != nil {
|
||
store.attachVectorIndex(r.vindex)
|
||
// 项目索引同时接到写入路径：项目作用域记忆的 ProjectID 落的索引
|
||
// 条目因此始终存在（见 store_project.go），Retriever.projectFor 才
|
||
// 有东西可解析。
|
||
store.attachProjectIndex(r.projectStore)
|
||
}
|
||
return r
|
||
}
|
||
|
||
// retrieverLog 返回检索器默认使用的 logger：store 自带的那一个
|
||
// （NewStore 已把它标为来源 "memory"，且绝不为 nil），store 缺失时
|
||
// 用一个不回显的丢弃 logger，避免检索路径出现无来源裸日志。
|
||
func retrieverLog(store *Store) *slog.Logger {
|
||
if store == nil {
|
||
return logger.WithSource(discardLogger(nil), "memory")
|
||
}
|
||
return store.log
|
||
}
|
||
|
||
// Retrieve 为 personaID 返回最匹配 query 的最多 topK 条，按
|
||
// 相关性（关键词 + 可选语义）加权衰减（decay_weight ×
|
||
// 时间衰减）排序。候选集是 design/03 §3/§4 的可见域并集：该角色自己的条目
|
||
// （ListByPersona）∪ 全局共享的用户习惯域（ScopeUser）∪ 人人可读的
|
||
// 项目域（ScopeProject），按 id 去重。他角色的角色日常记忆（ScopePersona）
|
||
// 绝不进入候选集 —— 跨角色读只走「把 B 当子代理去查 B」的派发机制
|
||
// （design/22 §9），直接检索不越权。query 与候选内容都在
|
||
// ctx 下读取；ctx 被取消会中止工作。query 为空时返回
|
||
// 最近使用过的条目（仅按衰减顺序）。
|
||
func (r *Retriever) Retrieve(ctx context.Context, personaID, query string, topK int) ([]Result, error) {
|
||
if topK <= 0 {
|
||
topK = defaultTopK
|
||
}
|
||
items, err := r.visibleCandidates(ctx, personaID)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if len(items) == 0 {
|
||
return []Result{}, nil
|
||
}
|
||
|
||
kw := r.keywordScores(items, query)
|
||
sem := r.semanticScores(ctx, items, query)
|
||
|
||
now := time.Now().UTC()
|
||
// matched 为 true 表示至少有一条条目与 query 有相关性。这种
|
||
// 模式下相关性主导排名，新近度只在命中项之间
|
||
// 衰减；未命中的条目得分为 0 且排在后面。无任何命中
|
||
// （例如 query 为空）时，只按新近度排名。
|
||
matched := false
|
||
relevance := make(map[string]float64, len(items))
|
||
for _, it := range items {
|
||
rel := kw[it.ID] + sem[it.ID]
|
||
relevance[it.ID] = rel
|
||
if rel > 0 {
|
||
matched = true
|
||
}
|
||
}
|
||
|
||
results := make([]Result, 0, len(items))
|
||
for _, it := range items {
|
||
decay := r.decayFactor(it, now)
|
||
score := decay
|
||
if matched {
|
||
score = relevance[it.ID] * decay
|
||
}
|
||
res := Result{
|
||
Item: it,
|
||
Score: score,
|
||
SourceConversation: it.SourceConversation,
|
||
Project: r.projectFor(ctx, it),
|
||
}
|
||
results = append(results, res)
|
||
}
|
||
|
||
sort.SliceStable(results, func(i, j int) bool {
|
||
if results[i].Score != results[j].Score {
|
||
return results[i].Score > results[j].Score
|
||
}
|
||
if !results[i].Item.UpdatedAt.Equal(results[j].Item.UpdatedAt) {
|
||
return results[i].Item.UpdatedAt.After(results[j].Item.UpdatedAt)
|
||
}
|
||
return results[i].Item.ID < results[j].Item.ID
|
||
})
|
||
|
||
if len(results) > topK {
|
||
results = results[:topK]
|
||
}
|
||
// 命中即使用：把返回命中的衰减物化回写 decay_weight（design §2，
|
||
// 裁决 Q3）。失败只记日志，绝不影响本次检索结果。
|
||
r.persistDecay(ctx, results, relevance, now)
|
||
return results, nil
|
||
}
|
||
|
||
// visibleCandidates 载入 personaID 可见的候选记忆并集（design/03
|
||
// §3/§4）：本人条目（ListByPersona）∪ ScopeUser 全局共享域 ∪
|
||
// ScopeProject 人人可读域，按 id 去重。本人条目先入集，因此去重
|
||
// 保留的是同一行的唯一实例；他角色的 ScopePersona 条目（角色
|
||
// 日常）不在其列 —— 跨角色读走子代理派发（design/22 §9），
|
||
// 直接检索不得越权。任一载入失败都整体报错，绝不静默丢弃
|
||
// 某个共享域。
|
||
func (r *Retriever) visibleCandidates(ctx context.Context, personaID string) ([]MemoryItem, error) {
|
||
own, err := r.store.ListByPersona(ctx, personaID)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("memory: retrieve list: %w", err)
|
||
}
|
||
seen := make(map[string]struct{}, len(own))
|
||
for _, it := range own {
|
||
seen[it.ID] = struct{}{}
|
||
}
|
||
items := own
|
||
for _, scope := range []Scope{ScopeUser, ScopeProject} {
|
||
shared, err := r.store.ListScope(ctx, scope, personaID)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("memory: retrieve list %s scope: %w", scope, err)
|
||
}
|
||
for _, it := range shared {
|
||
if _, dup := seen[it.ID]; dup {
|
||
continue
|
||
}
|
||
seen[it.ID] = struct{}{}
|
||
items = append(items, it)
|
||
}
|
||
}
|
||
return items, nil
|
||
}
|
||
|
||
// semanticScores 计算每个条目对 query 的语义相关性。接入了向量
|
||
// 索引（WithVectorIndex）时，它经 idx.Search 从持久化索引
|
||
// 作答 —— 真正的向量路径（design §3）。索引没有给出任何候选命中时
|
||
// （索引空或索引报错），它**回退**到内存余弦：配置了
|
||
// 增强却仍拿到一片 0 分，等于增强既无效又遮住了可用的路径
|
||
// （审计 P0-3；design/03 §3「真正启用向量检索」、design/04 §7）。
|
||
// 两条路径都没有时返回空 map，因此语义相关性
|
||
// 不贡献任何分数。
|
||
//
|
||
// query 为空（无任何可向量化的文本）时整条语义路径**直接跳过**：
|
||
// 没有 query 就没有可比的语义方向，拿空串去向量化既问不出东西，
|
||
// 又要把整库候选逐条 embed 一遍 —— 真机实测这条弯路是 39.3 秒的来源
|
||
// （见 retrieve_fallback.go）。此时排名由关键词与新近度决定，
|
||
// 与 Retrieve 的既有口径一致。
|
||
func (r *Retriever) semanticScores(ctx context.Context, items []MemoryItem, query string) map[string]float64 {
|
||
scores := make(map[string]float64, len(items))
|
||
if !hasSemanticQuery(query) {
|
||
return scores
|
||
}
|
||
if r.vindex != nil {
|
||
// 持久索引为空时先做一次性回填（每进程最多一次），使既有
|
||
// 条目也进入语义路径，而不是只有新写的条目可被检索。
|
||
r.ensureIndexSeeded(ctx)
|
||
if indexed := r.vectorIndexScores(ctx, items, query); len(indexed) > 0 {
|
||
return indexed
|
||
}
|
||
if r.embedder != nil {
|
||
r.warnSemanticFallback()
|
||
}
|
||
}
|
||
if r.embedder == nil {
|
||
return scores
|
||
}
|
||
return r.inMemorySemanticScores(ctx, items, query)
|
||
}
|
||
|
||
// hasSemanticQuery 报告 query 是否含有可供向量化的语义内容：
|
||
// 非空且非纯空白。附件类回合（主人只发文件、正文为空）正是空 query，
|
||
// 它不该触发任何向量化工作。
|
||
func hasSemanticQuery(query string) bool {
|
||
return strings.TrimSpace(query) != ""
|
||
}
|
||
|
||
// vectorIndexScores 从持久化向量索引作答语义相关性：
|
||
// 它用索引对 query 做向量化，检索最相似的若干 id，并
|
||
// 返回它们的余弦相似度，以候选命中的最高分为基准归一化到 [0,1]。
|
||
//
|
||
// 索引是**全局**的（承载所有角色/作用域的向量），而本次检索的候选
|
||
// 只是其中一个子集：不属于候选的命中不参与排名，也不得充当
|
||
// 归一化基准（也不得因占满 topK 名额而把候选挤出命中）。因此
|
||
// 这里只保留候选集内的命中。索引报错或候选内无命中时退化为
|
||
// 空 map —— 调用方据此回退内存余弦。query 为空的早退只是防御：
|
||
// 语义路径的调用方（semanticScores）已经在更外层拦掉了空 query，
|
||
// 不会再把「没问过索引」误当成「索引答不出」。
|
||
func (r *Retriever) vectorIndexScores(ctx context.Context, items []MemoryItem, query string) map[string]float64 {
|
||
scores := make(map[string]float64, len(items))
|
||
if strings.TrimSpace(query) == "" {
|
||
return scores
|
||
}
|
||
qvecs, err := r.vindex.Embed(ctx, []string{query})
|
||
if err != nil || len(qvecs) == 0 {
|
||
r.log.Warn("memory: vector index query embed failed, falling back to in-memory semantic",
|
||
"error", err)
|
||
return scores
|
||
}
|
||
hits, err := r.vindex.Search(ctx, qvecs[0], len(items))
|
||
if err != nil {
|
||
r.log.Warn("memory: vector index search failed, falling back to in-memory semantic",
|
||
"error", err)
|
||
return scores
|
||
}
|
||
candidates := make(map[string]struct{}, len(items))
|
||
for _, it := range items {
|
||
candidates[it.ID] = struct{}{}
|
||
}
|
||
matched := make([]vector.Hit, 0, len(hits))
|
||
max := 0.0
|
||
for _, h := range hits {
|
||
if _, ok := candidates[h.ID]; !ok {
|
||
continue
|
||
}
|
||
matched = append(matched, h)
|
||
if h.Score > max {
|
||
max = h.Score
|
||
}
|
||
}
|
||
if len(matched) == 0 {
|
||
return scores
|
||
}
|
||
for _, h := range matched {
|
||
if max > 0 {
|
||
scores[h.ID] = h.Score / max
|
||
} else {
|
||
scores[h.ID] = h.Score
|
||
}
|
||
}
|
||
return scores
|
||
}
|
||
|
||
// projectFor 返回链接到项目作用域条目的项目索引条目；检索器
|
||
// 没有项目索引存储、条目不是
|
||
// 项目作用域、或条目不带 ProjectID 时返回 nil。加载失败按「无
|
||
// 项目」处理，而不是让整次检索失败（design §3 项目索引各角色可查）。
|
||
//
|
||
// 加载失败**必须留下带来源的日志**（审计 2026-10-05 §7 静默吞点）：
|
||
// 「项目行不存在」（ErrProjectNotFound，项目记忆的 project_id 悬空）是
|
||
// 数据状态，按 Debug 记；其余错误（存储故障等）按 Warn 记，使真实的
|
||
// 项目索引故障不会被当成「这条记忆本来就没项目」而无限静默。两条都不
|
||
// 影响检索结果，条目照常返回，只是不带 Project。
|
||
func (r *Retriever) projectFor(ctx context.Context, it MemoryItem) *projectindex.Project {
|
||
if r.projectStore == nil || it.Scope != ScopeProject || strings.TrimSpace(it.ProjectID) == "" {
|
||
return nil
|
||
}
|
||
p, err := r.projectStore.GetProject(ctx, it.ProjectID)
|
||
if err != nil {
|
||
if errors.Is(err, projectindex.ErrProjectNotFound) {
|
||
r.log.Debug("memory: project link not found, result returned without project",
|
||
"project_id", it.ProjectID)
|
||
} else {
|
||
r.log.Warn("memory: project link lookup failed, result returned without project",
|
||
"project_id", it.ProjectID, "error", err)
|
||
}
|
||
return nil
|
||
}
|
||
return p
|
||
}
|
||
|
||
// Inject 为 personaID 检索最多 topK 条相关记忆，并把它们格式化为
|
||
// 紧凑的上下文块，可供注入提示词。没有检索到任何内容时
|
||
// 返回空字符串，因此没有相关记忆时调用方可以忽略
|
||
// 结果。
|
||
func (r *Retriever) Inject(ctx context.Context, personaID, query string, topK int) (string, error) {
|
||
results, err := r.Retrieve(ctx, personaID, query, topK)
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
if len(results) == 0 {
|
||
return "", nil
|
||
}
|
||
|
||
var b strings.Builder
|
||
b.WriteString("相关记忆（按相关度 × 衰减排序）：")
|
||
for _, res := range results {
|
||
b.WriteByte('\n')
|
||
b.WriteString("- [")
|
||
b.WriteString(string(res.Item.Kind))
|
||
b.WriteString("] ")
|
||
b.WriteString(res.Item.Content)
|
||
}
|
||
return b.String(), nil
|
||
}
|