feat(memory): 长记忆 store（SQLite+persona 隔离）/工作记忆滑动窗口/检索注入/手工提取钩子
This commit is contained in:
Aska committed
1 parent
285e1e9911
commit
aa697c385f
10 files changed
+1705
No files matched your search
@@ -0,0 +1,291 @@
|
||||
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
|
||||
}
|
||||
Reference in new issue
Block a user