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

// Retriever defaults and tuning constants.
const (
	// defaultKeywordWeight is the relative weight of keyword (FTS/TF-IDF)
	// relevance in the combined score.
	defaultKeywordWeight = 1.0

	// defaultSemanticWeight is the relative weight of semantic-vector relevance
	// in the combined score. It only contributes when an embedder is configured.
	defaultSemanticWeight = 1.0

	// defaultRecencyHalfLifeDays is the number of days over which a memory's
	// recency weight halves, used to favor recently touched memories (design §2
	// attenuation).
	defaultRecencyHalfLifeDays = 30.0

	// defaultTopK is the number of results returned when Retrieve or Inject is
	// called with a non-positive topK.
	defaultTopK = 10

	// hoursPerDay converts a day count into a time.Duration.
	hoursPerDay = 24
)

// Result is a single retrieval hit: the memory item plus its combined
// relevance × decay score explaining its ranking. Higher is more relevant.
type Result struct {
	Item  MemoryItem
	Score float64
	// SourceConversation is the conversation the item was extracted from, for
	// 溯源 (design §2). 裁决 04.4.3 起为结构化 JSON（session_id/message_id，
	// 见 ParseSourceConversation），由持久化层盖真实 id。It mirrors
	// Item.SourceConversation so callers can trace a hit without reaching into
	// Item.
	SourceConversation string
	// Project is the linked project-index entry (设计 §3 项目分层), populated only
	// when the retriever is wired with a project-index store and the item is a
	// project-scoped memory carrying a ProjectID. It lets callers surface the
	// 项目索引 → 项目描述 → 章节内容 chain for a retrieved item.
	Project *projectindex.Project
}

// Retriever ranks the long-term memories visible to a persona — its own items
// plus the shared user/project scopes (design §3/§4, see visibleCandidates) —
// by relevance (keyword and, when an embedder or vector index is configured,
// semantic vectors) weighted by decay (decay_weight × time decay), then returns
// the top K. It is stateless and safe for concurrent use; the ranking is
// recomputed from the store on every call, so results always reflect persisted
// state.
type Retriever struct {
	store           *Store
	embedder        embedding.Embedder
	vindex          vector.Index
	projectStore    *projectindex.Store
	keywordWeight   float64
	semanticWeight  float64
	recencyHalfLife time.Duration
}

// RetrieverOption customizes a Retriever built by NewRetriever.
type RetrieverOption func(*Retriever)

// WithEmbedder sets the semantic embedder used for vector similarity computed
// in-memory. When nil (the default), retrieval falls back to keyword matching
// plus decay only. When a vector index is also wired via WithVectorIndex, the
// index takes precedence: it is the persistent semantic path (design §3
// 可选增强), and WithEmbedder computes similarity on the fly instead.
func WithEmbedder(e embedding.Embedder) RetrieverOption {
	return func(r *Retriever) {
		r.embedder = e
	}
}

// WithVectorIndex wires a persistent vector.Index into retrieval. When set,
// semantic relevance is answered by idx.Search over the stored vectors rather
// than by in-memory embedding, so a configured pgvector/SQLite vector store is
// genuinely used (design §3 "真正启用向量检索"). When nil (the default), the
// semantic path is skipped unless WithEmbedder is set. The caller is responsible
// for keeping the index in sync with stored items (embed + add on write).
func WithVectorIndex(idx vector.Index) RetrieverOption {
	return func(r *Retriever) {
		r.vindex = idx
	}
}

// WithProjectIndex wires a projectindex.Store so retrieval can attach the linked
// project-index entry (设计 §3 项目分层) to each project-scoped hit. When nil
// (the default), Result.Project is left unset.
func WithProjectIndex(s *projectindex.Store) RetrieverOption {
	return func(r *Retriever) {
		r.projectStore = s
	}
}

// WithKeywordWeight sets the relative weight of keyword relevance. Non-positive
// values are ignored and the default kept.
func WithKeywordWeight(w float64) RetrieverOption {
	return func(r *Retriever) {
		if w > 0 {
			r.keywordWeight = w
		}
	}
}

// WithSemanticWeight sets the relative weight of semantic relevance. Non-positive
// values are ignored and the default kept.
func WithSemanticWeight(w float64) RetrieverOption {
	return func(r *Retriever) {
		if w > 0 {
			r.semanticWeight = w
		}
	}
}

// WithRecencyHalfLife sets how many days a memory's recency weight takes to
// halve. Non-positive values are ignored and the default kept.
func WithRecencyHalfLife(days float64) RetrieverOption {
	return func(r *Retriever) {
		if days > 0 {
			r.recencyHalfLife = time.Duration(days * float64(hoursPerDay*time.Hour))
		}
	}
}

// NewRetriever builds a Retriever over store. A nil store is invalid and panics
// at call time; callers should pass a non-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 returns up to topK items for personaID that best match query, ranked
// by relevance (keyword + optional semantic) weighted by decay (decay_weight ×
// time decay). 候选集是 design/03 §3/§4 的可见域并集：该角色自己的条目
// （ListByPersona）∪ 全局共享的用户习惯域（ScopeUser）∪ 人人可读的
// 项目域（ScopeProject），按 id 去重。他角色的角色日常记忆（ScopePersona）
// 绝不进入候选集 —— 跨角色读只走「把 B 当子代理去查 B」的派发机制
// （design/22 §9），直接检索不越权。Query and candidate contents are both
// read under ctx; a canceled context aborts the work. A blank query returns the
// most recently used items (decay order only).
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 is true when at least one item shares relevance with the query. In
	// that mode relevance dominates ranking and recency only attenuates among
	// matches; unmatched items get score 0 and sort below. When nothing matches
	// (e.g. a blank query), we rank by recency alone.
	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 scores each item against query with an in-memory TF-IDF inverted
// index (embedding.Index), reusing the embedding package's tokenizer. Scores are
// normalized to [0,1] by the maximum so they are comparable with the semantic
// scores.
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 computes per-item semantic relevance for query. When a vector
// index is wired (WithVectorIndex) it answers from the persistent index via
// idx.Search — the real vector path (design §3). Otherwise, with an embedder
// configured, it embeds query and candidates in-memory and uses clamped cosine
// similarity. With neither it returns an empty map, so semantic relevance
// contributes nothing.
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 answers semantic relevance from the persistent vector index:
// it embeds the query with the index, searches for the top similar ids, and
// returns their cosine similarity, normalized to [0,1] by the top hit. Because
// idx.Search ranks over all stored ids, an item absent from the index carries no
// semantic score. A blank query or an index error degrades to an empty 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 returns the project-index entry linked to a project-scoped item, or
// nil when the retriever has no project-index store, the item is not
// project-scoped, or it carries no ProjectID. A load failure is treated as "no
// project" rather than failing the whole retrieval (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 is the text indexed for keyword matching: content plus tags.
func keywordText(it MemoryItem) string {
	if len(it.Tags) == 0 {
		return it.Content
	}
	return it.Content + " " + strings.Join(it.Tags, " ")
}

// cosineSim computes the cosine similarity of two vectors, clamped to [0,1] by
// the caller. Zero-length or mismatched vectors yield 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 retrieves up to topK relevant memories for personaID and formats them
// as a compact context block ready to be injected into the prompt. It returns an
// empty string when nothing is retrieved, so callers can ignore the result when
// there is no relevant memory.
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
}