Cyrene_v2/internal/runtime/memory/retrieve.go

package memory

import (
	"context"
	"fmt"
	"math"
	"sort"
	"strings"
	"time"

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

// RetrieverOption 定制由 NewRetriever 构建的 Retriever。
type RetrieverOption func(*Retriever)

// WithEmbedder 设置用于内存内向量相似度计算的语义 embedder。
// 为 nil（默认）时，检索退化为仅关键词匹配
// 加衰减。若同时经 WithVectorIndex 接入了向量索引，则
// 索引优先：它是持久化语义路径（design §3
// 可选增强），此时 WithEmbedder 改为即时计算相似度。
func WithEmbedder(e embedding.Embedder) RetrieverOption {
	return func(r *Retriever) {
		r.embedder = e
	}
}

// WithVectorIndex 把持久化的 vector.Index 接入检索。设置后，
// 语义相关性由 idx.Search 在已存向量上作答，而非
// 内存内向量化，因此配置的 pgvector/SQLite 向量存储被
// 真正启用（design §3「真正启用向量检索」）。为 nil（默认）时，
// 除非设置了 WithEmbedder，否则语义路径跳过。调用方负责
// 让索引与已存条目保持同步（写入时 embed + add）。
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。
func NewRetriever(store *Store, opts ...RetrieverOption) *Retriever {
	r := &Retriever{
		store:           store,
		keywordWeight:   defaultKeywordWeight,
		semanticWeight:  defaultSemanticWeight,
		recencyHalfLife: time.Duration(defaultRecencyHalfLifeDays * float64(hoursPerDay*time.Hour)),
	}
	for _, opt := range opts {
		opt(r)
	}
	return r
}

// 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）。否则，在配置了 embedder
// 时，它在内存内对 query 与候选做向量化并使用夹取的余弦
// 相似度。两者都没有时返回空 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 {
		return r.vectorIndexScores(ctx, items, query)
	}
	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 {
		return scores
	}
	queryVectors, err := r.embedder.Embed(ctx, []string{query})
	if err != nil || len(queryVectors) == 0 {
		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]。由于
// idx.Search 在全部已存 id 上排名，不在索引中的条目没有
// 语义得分。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 {
		return scores
	}
	hits, err := r.vindex.Search(ctx, qvecs[0], len(items))
	if err != nil || len(hits) == 0 {
		return scores
	}
	max := hits[0].Score
	for _, h := range hits {
		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, " ")
}

// cosineSim 计算两个向量的余弦相似度，由调用方夹到 [0,1]。
// 长度为零或不一致的向量得 0。
func cosineSim(a, b embedding.Vector) float64 {
	if len(a) == 0 || len(b) == 0 || len(a) != len(b) {
		return 0
	}
	var dot, na, nb float64
	for i := range a {
		dot += float64(a[i]) * float64(b[i])
		na += float64(a[i]) * float64(a[i])
		nb += float64(b[i]) * float64(b[i])
	}
	if na == 0 || nb == 0 {
		return 0
	}
	return dot / (math.Sqrt(na) * math.Sqrt(nb))
}

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