292 lines
8.6 KiB
Go
292 lines
8.6 KiB
Go
package memory
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
"math"
|
||
"sort"
|
||
"strings"
|
||
"time"
|
||
|
||
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/embedding"
|
||
)
|
||
|
||
// 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 × recency score explaining its ranking. Higher is more relevant.
|
||
type Result struct {
|
||
Item MemoryItem
|
||
Score float64
|
||
}
|
||
|
||
// Retriever ranks a persona's long-term memories by relevance (keyword and,
|
||
// when an embedder is configured, semantic vectors) weighted by recency, 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
|
||
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. When nil
|
||
// (the default), retrieval falls back to keyword matching plus recency only.
|
||
func WithEmbedder(e embedding.Embedder) RetrieverOption {
|
||
return func(r *Retriever) {
|
||
r.embedder = e
|
||
}
|
||
}
|
||
|
||
// 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 recency. Query and
|
||
// candidate contents are both read under ctx; a canceled context aborts the
|
||
// work. A blank query returns the most recently used items (recency only).
|
||
func (r *Retriever) Retrieve(ctx context.Context, personaID, query string, topK int) ([]Result, error) {
|
||
if topK <= 0 {
|
||
topK = defaultTopK
|
||
}
|
||
items, err := r.store.ListByPersona(ctx, personaID)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("memory: retrieve list: %w", 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 {
|
||
recency := r.recencyFactor(it, now)
|
||
score := recency
|
||
if matched {
|
||
score = relevance[it.ID] * recency
|
||
}
|
||
results = append(results, Result{Item: it, Score: score})
|
||
}
|
||
|
||
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
|
||
}
|
||
|
||
// 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 embeds the query and each candidate content with the configured
|
||
// embedder and returns the clamped cosine similarity per item. With no embedder
|
||
// 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.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
|
||
}
|
||
|
||
// recencyFactor combines the stored decay_weight with the time since the item
|
||
// was last touched, halving every recencyHalfLife days.
|
||
func (r *Retriever) recencyFactor(it MemoryItem, now time.Time) float64 {
|
||
age := now.Sub(it.UpdatedAt)
|
||
if age < 0 {
|
||
age = 0
|
||
}
|
||
halfLives := float64(age) / float64(r.recencyHalfLife)
|
||
return it.DecayWeight * math.Pow(0.5, halfLives)
|
||
}
|
||
|
||
// 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
|
||
}
|