Cyrene_v2/internal/runtime/memory/retrieve.go

package memory

import (
	"context"
	"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）：
	// 回填每进程最多一次，而语义回退的告警也只报一次，避免每个
	// Retrieve 都写日志。
	seedMu         sync.Mutex
	seeded         bool
	fallbackWarned bool
}

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

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

// keywordScores 用内存内 TF-IDF 倒排索引（embedding.Index）为每条条目
// 与 query 打分，复用 embedding 包的分词器。得分
// 以最大值为基准归一化到 [0,1]，以便与语义
// 得分可比。
func (r *Retriever) keywordScores(items []MemoryItem, query string) map[string]float64 {
	scores := make(map[string]float64, len(items))
	index := embedding.NewIndex[string]()
	for _, it := range items {
		index.Add(keywordText(it), it.ID)
	}
	hits := index.Search(query)
	if len(hits) == 0 {
		return scores
	}
	max := hits[0].Score
	if max <= 0 {
		return scores
	}
	for _, h := range hits {
		scores[h.ID] = h.Score / max
	}
	return scores
}

// semanticScores 计算每个条目对 query 的语义相关性。接入了向量
// 索引（WithVectorIndex）时，它经 idx.Search 从持久化索引
// 作答 —— 真正的向量路径（design §3）。索引没有给出任何候选命中时
// （索引空、索引报错或 query 为空），它**回退**到内存余弦：配置了
// 增强却仍拿到一片 0 分，等于增强既无效又遮住了可用的路径
// （审计 P0-3；design/03 §3「真正启用向量检索」、design/04 §7）。
// 两条路径都没有时返回空 map，因此语义相关性
// 不贡献任何分数。
func (r *Retriever) semanticScores(ctx context.Context, items []MemoryItem, query string) map[string]float64 {
	scores := make(map[string]float64, len(items))
	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
	}

	texts := make([]string, len(items))
	for i, it := range items {
		texts[i] = it.Content
	}
	vectors, err := r.embedder.Embed(ctx, texts)
	if err != nil {
		r.log.Warn("memory: semantic embed of candidates failed, semantic scores skipped", "error", err)
		return scores
	}
	queryVectors, err := r.embedder.Embed(ctx, []string{query})
	if err != nil || len(queryVectors) == 0 {
		r.log.Warn("memory: semantic embed of query failed, semantic scores skipped", "error", err)
		return scores
	}
	qvec := queryVectors[0]
	for i, vec := range vectors {
		if i >= len(items) {
			break
		}
		sim := cosineSim(qvec, vec)
		if sim > 0 {
			scores[items[i].ID] = sim
		}
	}
	return scores
}

// vectorIndexScores 从持久化向量索引作答语义相关性：
// 它用索引对 query 做向量化，检索最相似的若干 id，并
// 返回它们的余弦相似度，以候选命中的最高分为基准归一化到 [0,1]。
//
// 索引是**全局**的（承载所有角色/作用域的向量），而本次检索的候选
// 只是其中一个子集：不属于候选的命中不参与排名，也不得充当
// 归一化基准（也不得因占满 topK 名额而把候选挤出命中）。因此
// 这里只保留候选集内的命中。query 为空、索引报错或候选内
// 无命中时退化为空 map —— 调用方据此回退内存余弦。
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 项目索引各角色可查）。
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 {
		return nil
	}
	return p
}

// keywordText 是用于关键词匹配的索引文本：content 加 tags。
func keywordText(it MemoryItem) string {
	if len(it.Tags) == 0 {
		return it.Content
	}
	return it.Content + " " + strings.Join(it.Tags, " ")
}

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