feat(memory): 长记忆 store(SQLite+persona 隔离)/工作记忆滑动窗口/检索注入/手工提取钩子
This commit is contained in:
1 parent
285e1e9911
commit
aa697c385f
10 files changed
+1705
No files matched your search
@@ -0,0 +1,81 @@
|
||||
package memory
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
|
||||
)
|
||||
|
||||
// Extractor derives candidate memory items from a conversation (design §1).
|
||||
//
|
||||
// The production path is an LLM-driven asynchronous extractor that denoises and
|
||||
// deduplicates each turn; that is deferred to milestone M4, and its integration
|
||||
// point is this interface. The M2 implementation is the rule-based
|
||||
// ManualExtractor, which only records what the user explicitly asks to remember.
|
||||
type Extractor interface {
|
||||
// Extract returns memory items worth persisting from conversation. ctx
|
||||
// controls cancellation; a canceled context returns an error.
|
||||
Extract(ctx context.Context, conversation []model.Message) ([]MemoryItem, error)
|
||||
}
|
||||
|
||||
// ManualExtractor implements the manual "记住 X" extraction path without an LLM.
|
||||
// It scans user turns for a leading remember directive and produces a memory
|
||||
// item for each, falling back to no item for ordinary dialogue. It is the
|
||||
// deterministic counterpart to the M4 LLM extractor.
|
||||
type ManualExtractor struct {
|
||||
personaID string
|
||||
roleID string
|
||||
}
|
||||
|
||||
// NewManualExtractor builds a ManualExtractor whose extracted items are scoped
|
||||
// to personaID and roleID, isolated from other personas (design §4).
|
||||
func NewManualExtractor(personaID, roleID string) *ManualExtractor {
|
||||
return &ManualExtractor{personaID: personaID, roleID: roleID}
|
||||
}
|
||||
|
||||
// Extract returns one memory item (kind relationship, persona scope) per user
|
||||
// turn that asks the agent to remember something. A leading "记住"/"请记住"
|
||||
// directive is stripped from the stored content; turns without one are skipped.
|
||||
func (e *ManualExtractor) Extract(ctx context.Context, conversation []model.Message) ([]MemoryItem, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items := make([]MemoryItem, 0)
|
||||
for _, msg := range conversation {
|
||||
if msg.Role != model.RoleUser {
|
||||
continue
|
||||
}
|
||||
text := strings.TrimSpace(messageText(msg))
|
||||
content, ok := rememberDirective(text)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
items = append(items, MemoryItem{
|
||||
PersonaID: e.personaID,
|
||||
RoleID: e.roleID,
|
||||
Scope: ScopePersona,
|
||||
Kind: KindRelationship,
|
||||
Content: content,
|
||||
Confidence: 1.0,
|
||||
DecayWeight: 1.0,
|
||||
})
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
// rememberDirective strips the "记住" / "请记住" prefix when text is a remember
|
||||
// instruction. It returns the remembered content and true, or ok=false when the
|
||||
// text is not an explicit remember request.
|
||||
func rememberDirective(text string) (string, bool) {
|
||||
for _, prefix := range []string{"请记住", "记住"} {
|
||||
if strings.HasPrefix(text, prefix) {
|
||||
content := strings.TrimSpace(strings.TrimPrefix(text, prefix))
|
||||
if content == "" {
|
||||
return "", false
|
||||
}
|
||||
return content, true
|
||||
}
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
package memory
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
|
||||
)
|
||||
|
||||
func TestManualExtractor_Extract_ParseRemember_正常(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ex := NewManualExtractor("persona.a", "role.1")
|
||||
conversation := []model.Message{
|
||||
model.TextMessage(model.RoleUser, "记住她喜欢喝冰美式"),
|
||||
model.TextMessage(model.RoleAssistant, "好的,我记住了"),
|
||||
model.TextMessage(model.RoleUser, "今天天气不错"), // no directive
|
||||
model.TextMessage(model.RoleUser, "请记住周三下午有项目会"),
|
||||
}
|
||||
|
||||
items, err := ex.Extract(context.Background(), conversation)
|
||||
if err != nil {
|
||||
t.Fatalf("Extract: %v", err)
|
||||
}
|
||||
if len(items) != 2 {
|
||||
t.Fatalf("got %d items, want 2", len(items))
|
||||
}
|
||||
if items[0].Content != "她喜欢喝冰美式" {
|
||||
t.Errorf("items[0].Content = %q, want the stripped directive", items[0].Content)
|
||||
}
|
||||
if items[0].PersonaID != "persona.a" || items[0].RoleID != "role.1" {
|
||||
t.Errorf("items[0] persona/role = %q/%q, want persona.a/role.1", items[0].PersonaID, items[0].RoleID)
|
||||
}
|
||||
if items[1].Content != "周三下午有项目会" {
|
||||
t.Errorf("items[1].Content = %q, want the 请记住 content", items[1].Content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestManualExtractor_Extract_NoRemember_正常(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ex := NewManualExtractor("persona.a", "role.1")
|
||||
conversation := []model.Message{
|
||||
model.TextMessage(model.RoleUser, "今天天气不错"),
|
||||
model.TextMessage(model.RoleUser, "记得带伞吗"),
|
||||
model.TextMessage(model.RoleAssistant, "记住是记忆"), // non-user role skipped
|
||||
}
|
||||
|
||||
items, err := ex.Extract(context.Background(), conversation)
|
||||
if err != nil {
|
||||
t.Fatalf("Extract: %v", err)
|
||||
}
|
||||
if len(items) != 0 {
|
||||
t.Fatalf("got %d items, want 0", len(items))
|
||||
}
|
||||
}
|
||||
|
||||
func TestManualExtractor_Extract_Canceled_异常(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ex := NewManualExtractor("persona.a", "role.1")
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
if _, err := ex.Extract(ctx, []model.Message{model.TextMessage(model.RoleUser, "记住x")}); !errors.Is(err, context.Canceled) {
|
||||
t.Fatalf("Extract with canceled ctx err = %v, want context.Canceled", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_RememberManual_正常(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
item, err := s.RememberManual(ctx, "persona.a", "role.1", "记住她喜欢喝咖啡")
|
||||
if err != nil {
|
||||
t.Fatalf("RememberManual: %v", err)
|
||||
}
|
||||
// The stored content keeps the directive wording; it's a manual single-item
|
||||
// path, distinct from the extractor which strips the directive.
|
||||
if item.Content != "记住她喜欢喝咖啡" {
|
||||
t.Fatalf("RememberManual content = %q, want the passed text", item.Content)
|
||||
}
|
||||
if item.PersonaID != "persona.a" || item.Scope != ScopePersona || item.Kind != KindRelationship {
|
||||
t.Fatalf("item persona/scope/kind = %q/%q/%q, want persona.a/persona/relationship",
|
||||
item.PersonaID, item.Scope, item.Kind)
|
||||
}
|
||||
|
||||
got, err := s.Get(ctx, item.ID)
|
||||
if err != nil {
|
||||
t.Fatalf("Get: %v", err)
|
||||
}
|
||||
if got.Content != item.Content {
|
||||
t.Fatalf("persisted content = %q, want %q", got.Content, item.Content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_RememberManual_Empty_异常(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
for _, text := range []string{"", " "} {
|
||||
if _, err := s.RememberManual(ctx, "persona.a", "role.1", text); !errors.Is(err, ErrMissingContent) {
|
||||
t.Fatalf("RememberManual(%q) err = %v, want ErrMissingContent", text, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
package memory
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"log/slog"
|
||||
"testing"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/db"
|
||||
)
|
||||
|
||||
// testLogger returns a logger that writes nowhere, keeping test output clean.
|
||||
func testLogger() *slog.Logger {
|
||||
return slog.New(slog.NewTextHandler(io.Discard, nil))
|
||||
}
|
||||
|
||||
// newTestStore opens an isolated in-memory SQLite DB and builds a Store over it.
|
||||
// Each call owns a private DB, so parallel tests are isolated by construction.
|
||||
func newTestStore(t *testing.T) *Store {
|
||||
t.Helper()
|
||||
ctx := context.Background()
|
||||
d, err := db.Open(ctx, ":memory:", testLogger())
|
||||
if err != nil {
|
||||
t.Fatalf("open in-memory db: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.Close() })
|
||||
|
||||
s, err := NewStore(ctx, d, testLogger())
|
||||
if err != nil {
|
||||
t.Fatalf("NewStore: %v", err)
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// baseItem is a minimal valid memory item for tests.
|
||||
func baseItem() MemoryItem {
|
||||
return MemoryItem{
|
||||
PersonaID: "persona.a",
|
||||
RoleID: "role.1",
|
||||
Scope: ScopePersona,
|
||||
Kind: KindRelationship,
|
||||
Content: "她喜欢喝冰美式",
|
||||
Tags: []string{"habit", "coffee"},
|
||||
Confidence: 0.9,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
package memory
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
|
||||
)
|
||||
|
||||
// messageText returns the concatenated plain-text payload of a message. Image
|
||||
// parts are skipped; text parts are joined with a single space. It is the
|
||||
// shared text view used by the working-memory summary and the manual extractor.
|
||||
func messageText(m model.Message) string {
|
||||
var b strings.Builder
|
||||
for _, part := range m.Content {
|
||||
if part.Type != model.ContentTypeText {
|
||||
continue
|
||||
}
|
||||
if b.Len() > 0 {
|
||||
b.WriteByte(' ')
|
||||
}
|
||||
b.WriteString(part.Text)
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -0,0 +1,217 @@
|
||||
package memory
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/embedding"
|
||||
)
|
||||
|
||||
// addItem is a convenience to persist an item and fail the test on error.
|
||||
func addItem(t *testing.T, s *Store, item MemoryItem) MemoryItem {
|
||||
t.Helper()
|
||||
ctx := context.Background()
|
||||
stored, err := s.Add(ctx, item)
|
||||
if err != nil {
|
||||
t.Fatalf("Add item: %v", err)
|
||||
}
|
||||
return stored
|
||||
}
|
||||
|
||||
func TestRetriever_Retrieve_Relevance_正常(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
addItem(t, s, MemoryItem{
|
||||
PersonaID: "persona.a", Content: "user loves coffee", Tags: []string{"coffee"},
|
||||
Scope: ScopePersona, Kind: KindWorldEvent,
|
||||
})
|
||||
addItem(t, s, MemoryItem{
|
||||
PersonaID: "persona.a", Content: "user plays guitar", Tags: []string{"music"},
|
||||
Scope: ScopePersona, Kind: KindWorldEvent,
|
||||
})
|
||||
|
||||
r := NewRetriever(s)
|
||||
results, err := r.Retrieve(ctx, "persona.a", "coffee", 10)
|
||||
if err != nil {
|
||||
t.Fatalf("Retrieve: %v", err)
|
||||
}
|
||||
if len(results) == 0 {
|
||||
t.Fatal("expected at least one result")
|
||||
}
|
||||
if got := results[0].Item.Content; got != "user loves coffee" {
|
||||
t.Fatalf("top result = %q, want the coffee item", got)
|
||||
}
|
||||
if results[0].Score <= 0 {
|
||||
t.Fatalf("top result score = %v, want > 0", results[0].Score)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRetriever_Retrieve_Recency_正常(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
old := time.Now().UTC().Add(-48 * time.Hour)
|
||||
newer := time.Now().UTC().Add(-2 * time.Hour)
|
||||
addItem(t, s, MemoryItem{
|
||||
PersonaID: "persona.a", Content: "old memory", UpdatedAt: old,
|
||||
Scope: ScopePersona, Kind: KindWorldEvent,
|
||||
})
|
||||
addItem(t, s, MemoryItem{
|
||||
PersonaID: "persona.a", Content: "new memory", UpdatedAt: newer,
|
||||
Scope: ScopePersona, Kind: KindWorldEvent,
|
||||
})
|
||||
|
||||
r := NewRetriever(s)
|
||||
// A blank query matches nothing, so ordering falls back to recency.
|
||||
results, err := r.Retrieve(ctx, "persona.a", "", 10)
|
||||
if err != nil {
|
||||
t.Fatalf("Retrieve: %v", err)
|
||||
}
|
||||
if len(results) != 2 {
|
||||
t.Fatalf("got %d results, want 2", len(results))
|
||||
}
|
||||
if got := results[0].Item.Content; got != "new memory" {
|
||||
t.Fatalf("top result = %q, want the newer memory", got)
|
||||
}
|
||||
// The older memory must rank below.
|
||||
if got := results[1].Item.Content; got != "old memory" {
|
||||
t.Fatalf("second result = %q, want the older memory", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRetriever_Retrieve_TopK_正常(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
for i := 0; i < 5; i++ {
|
||||
addItem(t, s, MemoryItem{
|
||||
PersonaID: "persona.a", Content: "shared term item",
|
||||
Tags: []string{"term"}, Scope: ScopePersona, Kind: KindWorldEvent,
|
||||
})
|
||||
}
|
||||
|
||||
r := NewRetriever(s)
|
||||
results, err := r.Retrieve(ctx, "persona.a", "term", 3)
|
||||
if err != nil {
|
||||
t.Fatalf("Retrieve: %v", err)
|
||||
}
|
||||
if len(results) != 3 {
|
||||
t.Fatalf("got %d results, want topK=3", len(results))
|
||||
}
|
||||
}
|
||||
|
||||
func TestRetriever_Retrieve_Isolation_正常(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
addItem(t, s, MemoryItem{
|
||||
PersonaID: "persona.a", Content: "coffee for a", Tags: []string{"coffee"},
|
||||
Scope: ScopePersona, Kind: KindRelationship,
|
||||
})
|
||||
addItem(t, s, MemoryItem{
|
||||
PersonaID: "persona.b", Content: "coffee for b", Tags: []string{"coffee"},
|
||||
Scope: ScopePersona, Kind: KindRelationship,
|
||||
})
|
||||
|
||||
r := NewRetriever(s)
|
||||
results, err := r.Retrieve(ctx, "persona.a", "coffee", 10)
|
||||
if err != nil {
|
||||
t.Fatalf("Retrieve: %v", err)
|
||||
}
|
||||
for _, res := range results {
|
||||
if res.Item.PersonaID != "persona.a" {
|
||||
t.Fatalf("retrieved item belongs to persona %q, want persona.a", res.Item.PersonaID)
|
||||
}
|
||||
}
|
||||
if len(results) != 1 {
|
||||
t.Fatalf("got %d results for persona.a, want 1", len(results))
|
||||
}
|
||||
}
|
||||
|
||||
func TestRetriever_Retrieve_WithEmbedder_正常(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
addItem(t, s, MemoryItem{
|
||||
PersonaID: "persona.a", Content: "she is a coffee lover", Tags: []string{"drink"},
|
||||
Scope: ScopePersona, Kind: KindRelationship,
|
||||
})
|
||||
addItem(t, s, MemoryItem{
|
||||
PersonaID: "persona.a", Content: "she plays the guitar", Tags: []string{"music"},
|
||||
Scope: ScopePersona, Kind: KindRelationship,
|
||||
})
|
||||
|
||||
emb := embedding.NewNGram(64)
|
||||
r := NewRetriever(s, WithEmbedder(emb))
|
||||
results, err := r.Retrieve(ctx, "persona.a", "coffee", 10)
|
||||
if err != nil {
|
||||
t.Fatalf("Retrieve with embedder: %v", err)
|
||||
}
|
||||
if len(results) == 0 {
|
||||
t.Fatal("expected results when embedding is enabled")
|
||||
}
|
||||
if got := results[0].Item.Content; got != "she is a coffee lover" {
|
||||
t.Fatalf("top result = %q, want the coffee item", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRetriever_Retrieve_EmptyPersona_正常(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
r := NewRetriever(s)
|
||||
results, err := r.Retrieve(ctx, "ghost", "anything", 10)
|
||||
if err != nil {
|
||||
t.Fatalf("Retrieve: %v", err)
|
||||
}
|
||||
if len(results) != 0 {
|
||||
t.Fatalf("got %d results for empty persona, want 0", len(results))
|
||||
}
|
||||
}
|
||||
|
||||
func TestRetriever_Inject_Format_正常(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
addItem(t, s, MemoryItem{
|
||||
PersonaID: "persona.a", Content: "她喜欢冰美式", Tags: []string{"coffee"},
|
||||
Scope: ScopePersona, Kind: KindRelationship,
|
||||
})
|
||||
|
||||
r := NewRetriever(s)
|
||||
block, err := r.Inject(ctx, "persona.a", "冰美式", 10)
|
||||
if err != nil {
|
||||
t.Fatalf("Inject: %v", err)
|
||||
}
|
||||
if !strings.Contains(block, "相关记忆") {
|
||||
t.Fatalf("inject block missing header: %q", block)
|
||||
}
|
||||
if !strings.Contains(block, "她喜欢冰美式") {
|
||||
t.Fatalf("inject block missing item content: %q", block)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRetriever_Inject_NoResults_正常(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
r := NewRetriever(s)
|
||||
block, err := r.Inject(ctx, "ghost", "anything", 10)
|
||||
if err != nil {
|
||||
t.Fatalf("Inject: %v", err)
|
||||
}
|
||||
if block != "" {
|
||||
t.Fatalf("Inject on empty persona = %q, want empty string", block)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,443 @@
|
||||
// Package memory implements the Cyrene v2 memory subsystem described in
|
||||
// DevDocs/design/03-记忆.md.
|
||||
//
|
||||
// It exposes three layers:
|
||||
//
|
||||
// - Store: a SQLite-backed long-term memory store. It persists memory items
|
||||
// with persona/role isolation (design §0 layer ②, §4) and offers
|
||||
// Add/Get/ListByPersona/Delete plus a manual "记住" path (design §1).
|
||||
// - WorkingMemory: a per-session sliding window of messages plus a summary
|
||||
// (design §0 layer ①).
|
||||
// - Retriever: relevance + recency retrieval and context injection
|
||||
// (design §2).
|
||||
//
|
||||
// Extraction (design §1) is exposed through the Extractor interface; the
|
||||
// concrete M2 implementation is the rule-based ManualExtractor, which records
|
||||
// what the user explicitly asks to remember. LLM-driven automatic extraction is
|
||||
// deferred to milestone M4 and is not implemented here.
|
||||
//
|
||||
// Milestone M4 additions that this package deliberately leaves out but keeps
|
||||
// extension points for: persistent vector storage, decay scheduling, and the
|
||||
// project-index hierarchy (design §3). The schema already carries decay_weight
|
||||
// and the retriever already plugs an optional embedding.Embedder, so those
|
||||
// features land without a schema migration.
|
||||
package memory
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"database/sql"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Scope classifies who a long-term memory belongs to. It is orthogonal to Kind
|
||||
// and both are stored on every item (design §3 note).
|
||||
type Scope string
|
||||
|
||||
const (
|
||||
// ScopeUser is a memory about a user's habits, shared across personas.
|
||||
ScopeUser Scope = "user"
|
||||
// ScopePersona is a persona's everyday memory; other personas are read-only.
|
||||
ScopePersona Scope = "persona"
|
||||
// ScopeProject is a project-scoped memory, readable and writable by all.
|
||||
ScopeProject Scope = "project"
|
||||
)
|
||||
|
||||
// valid reports whether s is one of the defined scopes.
|
||||
func (s Scope) valid() bool {
|
||||
switch s {
|
||||
case ScopeUser, ScopePersona, ScopeProject:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// Kind classifies the content category of a long-term memory item.
|
||||
type Kind string
|
||||
|
||||
const (
|
||||
// KindRelationship is a memory about "you and me": preferences, habits,
|
||||
// agreements (design §0 "关系记忆").
|
||||
KindRelationship Kind = "relationship"
|
||||
// KindWorldEvent is a memory about objective facts and events in the world
|
||||
// (design §0 "世界/事件记忆").
|
||||
KindWorldEvent Kind = "world_event"
|
||||
)
|
||||
|
||||
// valid reports whether k is one of the defined kinds.
|
||||
func (k Kind) valid() bool {
|
||||
switch k {
|
||||
case KindRelationship, KindWorldEvent:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// MemoryItem is a single long-term memory record. The zero value is not a valid
|
||||
// item: PersonaID and Content are required, Scope and Kind are normalized to
|
||||
// their defaults by Add.
|
||||
type MemoryItem struct {
|
||||
ID string
|
||||
PersonaID string
|
||||
RoleID string
|
||||
Scope Scope
|
||||
Kind Kind
|
||||
Content string
|
||||
Tags []string
|
||||
Confidence float64
|
||||
SourceConversation string
|
||||
CreatedAt time.Time
|
||||
UpdatedAt time.Time
|
||||
DecayWeight float64
|
||||
}
|
||||
|
||||
// Sentinels distinguishing memory package failures. Each is wrapped with %w at
|
||||
// the call site so callers can use errors.Is.
|
||||
var (
|
||||
// ErrMissingPersona is returned when an item has no PersonaID.
|
||||
ErrMissingPersona = errors.New("memory: persona_id is required")
|
||||
|
||||
// ErrMissingContent is returned when an item has no content.
|
||||
ErrMissingContent = errors.New("memory: content is required")
|
||||
|
||||
// ErrInvalidScope is returned when an item's Scope is not a defined scope.
|
||||
ErrInvalidScope = errors.New("memory: invalid scope")
|
||||
|
||||
// ErrInvalidKind is returned when an item's Kind is not a defined kind.
|
||||
ErrInvalidKind = errors.New("memory: invalid kind")
|
||||
|
||||
// ErrNotFound is returned when an item id does not exist.
|
||||
ErrNotFound = errors.New("memory: item not found")
|
||||
|
||||
// ErrNilDB is returned when NewStore receives a nil *sql.DB.
|
||||
ErrNilDB = errors.New("memory: nil db")
|
||||
)
|
||||
|
||||
// Schema defaults and time-window constants.
|
||||
const (
|
||||
defaultConfidence = 1.0
|
||||
defaultDecayWeight = 1.0
|
||||
)
|
||||
|
||||
// Store is a long-term memory store backed by a *sql.DB.
|
||||
//
|
||||
// The store owns its schema: NewStore creates the memory_items table and its
|
||||
// indexes idempotently. Concurrency is delegated to *sql.DB, which is safe for
|
||||
// concurrent use by contract, so Store holds no further lock. Every query is
|
||||
// parameterized (?) and lists columns explicitly (no SELECT *), and reads take
|
||||
// a context.Cancel.
|
||||
type Store struct {
|
||||
db *sql.DB
|
||||
log *slog.Logger
|
||||
}
|
||||
|
||||
// rowScanner is satisfied by both *sql.Row and *sql.Rows, so a single scanItem
|
||||
// helper serves Get and ListByPersona.
|
||||
type rowScanner interface {
|
||||
Scan(dest ...any) error
|
||||
}
|
||||
|
||||
// NewStore builds a Store on db, ensuring the memory_items schema exists.
|
||||
//
|
||||
// db must be a live, migrated connection (the migration framework itself is
|
||||
// owned by internal/db). A nil db returns ErrNilDB. log may be nil, in which
|
||||
// case a discard logger is used. ctx controls schema creation and must be
|
||||
// honored by the caller.
|
||||
func NewStore(ctx context.Context, db *sql.DB, log *slog.Logger) (*Store, error) {
|
||||
if db == nil {
|
||||
return nil, ErrNilDB
|
||||
}
|
||||
log = discardLogger(log)
|
||||
for _, ddl := range []string{
|
||||
createMemoryItemsTable,
|
||||
createMemoryItemsPersonaIndex,
|
||||
createMemoryItemsPersonaCreatedIndex,
|
||||
} {
|
||||
if _, err := db.ExecContext(ctx, ddl); err != nil {
|
||||
return nil, fmt.Errorf("memory: apply schema: %w", err)
|
||||
}
|
||||
}
|
||||
return &Store{db: db, log: log}, nil
|
||||
}
|
||||
|
||||
// Add persists item, assigning an id and missing timestamps, and normalizing
|
||||
// Scope/Kind/Confidence/DecayWeight to their defaults. It returns the stored
|
||||
// item with the id and timestamps filled in.
|
||||
func (s *Store) Add(ctx context.Context, item MemoryItem) (MemoryItem, error) {
|
||||
if err := item.normalize(); err != nil {
|
||||
return MemoryItem{}, err
|
||||
}
|
||||
var err error
|
||||
if item, err = applyTimestamps(item); err != nil {
|
||||
return MemoryItem{}, err
|
||||
}
|
||||
tagsJSON, err := json.Marshal(item.Tags)
|
||||
if err != nil {
|
||||
return MemoryItem{}, fmt.Errorf("memory: encode tags: %w", err)
|
||||
}
|
||||
|
||||
_, err = s.db.ExecContext(ctx, `INSERT INTO memory_items
|
||||
(id, persona_id, role_id, scope, kind, content, tags, confidence,
|
||||
source_conversation, created_at, updated_at, decay_weight)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
|
||||
item.ID, item.PersonaID, item.RoleID, string(item.Scope), string(item.Kind),
|
||||
item.Content, string(tagsJSON), item.Confidence, item.SourceConversation,
|
||||
formatTime(item.CreatedAt), formatTime(item.UpdatedAt), item.DecayWeight)
|
||||
if err != nil {
|
||||
return MemoryItem{}, fmt.Errorf("memory: insert item: %w", err)
|
||||
}
|
||||
return item, nil
|
||||
}
|
||||
|
||||
// Get returns the item with the given id. It returns ErrNotFound when the id is
|
||||
// empty or does not exist.
|
||||
func (s *Store) Get(ctx context.Context, id string) (MemoryItem, error) {
|
||||
if strings.TrimSpace(id) == "" {
|
||||
return MemoryItem{}, ErrNotFound
|
||||
}
|
||||
row := s.db.QueryRowContext(ctx, `SELECT id, persona_id, role_id, scope, kind,
|
||||
content, tags, confidence, source_conversation, created_at, updated_at,
|
||||
decay_weight FROM memory_items WHERE id = ?`, id)
|
||||
item, err := scanItem(row)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return MemoryItem{}, ErrNotFound
|
||||
}
|
||||
if err != nil {
|
||||
return MemoryItem{}, fmt.Errorf("memory: get item: %w", err)
|
||||
}
|
||||
return item, nil
|
||||
}
|
||||
|
||||
// ListByPersona returns every item for personaID, newest first. It is the
|
||||
// persona-isolation boundary: items belonging to other personas are never
|
||||
// returned, so personas never share memory (design §4).
|
||||
func (s *Store) ListByPersona(ctx context.Context, personaID string) ([]MemoryItem, error) {
|
||||
if strings.TrimSpace(personaID) == "" {
|
||||
return nil, ErrMissingPersona
|
||||
}
|
||||
rows, err := s.db.QueryContext(ctx, `SELECT id, persona_id, role_id, scope, kind,
|
||||
content, tags, confidence, source_conversation, created_at, updated_at,
|
||||
decay_weight FROM memory_items WHERE persona_id = ?
|
||||
ORDER BY created_at DESC`, personaID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("memory: list persona items: %w", err)
|
||||
}
|
||||
defer func() { _ = rows.Close() }()
|
||||
|
||||
items := make([]MemoryItem, 0)
|
||||
for rows.Next() {
|
||||
item, err := scanItem(rows)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("memory: scan persona item: %w", err)
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, fmt.Errorf("memory: iterate persona items: %w", err)
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
// Delete removes the item with the given id and returns ErrNotFound when the id
|
||||
// is empty or no row was removed. Callers should only pass ids they obtained
|
||||
// through a persona-scoped read (ListByPersona/Get) so deletion stays within
|
||||
// the caller's persona boundary.
|
||||
func (s *Store) Delete(ctx context.Context, id string) error {
|
||||
if strings.TrimSpace(id) == "" {
|
||||
return ErrNotFound
|
||||
}
|
||||
res, err := s.db.ExecContext(ctx, `DELETE FROM memory_items WHERE id = ?`, id)
|
||||
if err != nil {
|
||||
return fmt.Errorf("memory: delete item: %w", err)
|
||||
}
|
||||
n, err := res.RowsAffected()
|
||||
if err != nil {
|
||||
return fmt.Errorf("memory: delete row count: %w", err)
|
||||
}
|
||||
if n == 0 {
|
||||
return ErrNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// RememberManual writes a single long-term memory item for content text — the
|
||||
// "记住 X" path from design §1. The item is scoped to the persona by default and
|
||||
// recorded as a relationship memory at full confidence. Returns the stored item.
|
||||
func (s *Store) RememberManual(ctx context.Context, personaID, roleID, text string) (MemoryItem, error) {
|
||||
item := MemoryItem{
|
||||
PersonaID: personaID,
|
||||
RoleID: roleID,
|
||||
Scope: ScopePersona,
|
||||
Kind: KindRelationship,
|
||||
Content: strings.TrimSpace(text),
|
||||
}
|
||||
if item.Content == "" {
|
||||
return MemoryItem{}, ErrMissingContent
|
||||
}
|
||||
return s.Add(ctx, item)
|
||||
}
|
||||
|
||||
// normalize fills default Scope/Kind/Confidence/DecayWeight then validates.
|
||||
// It is called by Add so callers may omit values while still catching
|
||||
// explicitly invalid ones.
|
||||
func (it *MemoryItem) normalize() error {
|
||||
if it.Scope == "" {
|
||||
it.Scope = ScopePersona
|
||||
}
|
||||
if it.Kind == "" {
|
||||
it.Kind = KindRelationship
|
||||
}
|
||||
if it.Confidence == 0 {
|
||||
it.Confidence = defaultConfidence
|
||||
}
|
||||
if it.DecayWeight == 0 {
|
||||
it.DecayWeight = defaultDecayWeight
|
||||
}
|
||||
return it.validate()
|
||||
}
|
||||
|
||||
// validate checks required and enumerated fields. It is pure: it does not apply
|
||||
// defaults, so it can also be used to reject a caller-provided value.
|
||||
func (it *MemoryItem) validate() error {
|
||||
if strings.TrimSpace(it.PersonaID) == "" {
|
||||
return ErrMissingPersona
|
||||
}
|
||||
if strings.TrimSpace(it.Content) == "" {
|
||||
return ErrMissingContent
|
||||
}
|
||||
if !it.Scope.valid() {
|
||||
return ErrInvalidScope
|
||||
}
|
||||
if !it.Kind.valid() {
|
||||
return ErrInvalidKind
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// applyTimestamps assigns an id and missing CreatedAt/UpdatedAt values. It
|
||||
// returns an error only when a fresh id cannot be generated.
|
||||
func applyTimestamps(item MemoryItem) (MemoryItem, error) {
|
||||
now := time.Now().UTC()
|
||||
if item.ID == "" {
|
||||
id, err := newID()
|
||||
if err != nil {
|
||||
return MemoryItem{}, err
|
||||
}
|
||||
item.ID = id
|
||||
}
|
||||
if item.CreatedAt.IsZero() {
|
||||
item.CreatedAt = now
|
||||
}
|
||||
if item.UpdatedAt.IsZero() {
|
||||
item.UpdatedAt = now
|
||||
}
|
||||
return item, nil
|
||||
}
|
||||
|
||||
// scanItem reads one memory_items row. tags is stored as a JSON array string;
|
||||
// timestamps are stored as RFC3339 strings.
|
||||
func scanItem(sc rowScanner) (MemoryItem, error) {
|
||||
var (
|
||||
it MemoryItem
|
||||
scope, kind string
|
||||
tagsJSON string
|
||||
createdAt, updatedAt string
|
||||
)
|
||||
if err := sc.Scan(&it.ID, &it.PersonaID, &it.RoleID, &scope, &kind,
|
||||
&it.Content, &tagsJSON, &it.Confidence, &it.SourceConversation,
|
||||
&createdAt, &updatedAt, &it.DecayWeight); err != nil {
|
||||
return MemoryItem{}, err
|
||||
}
|
||||
it.Scope = Scope(scope)
|
||||
it.Kind = Kind(kind)
|
||||
if err := json.Unmarshal([]byte(tagsJSON), &it.Tags); err != nil {
|
||||
return MemoryItem{}, fmt.Errorf("memory: decode tags: %w", err)
|
||||
}
|
||||
if it.Tags == nil {
|
||||
it.Tags = []string{}
|
||||
}
|
||||
|
||||
var err error
|
||||
if it.CreatedAt, err = parseTime(createdAt); err != nil {
|
||||
return MemoryItem{}, err
|
||||
}
|
||||
if it.UpdatedAt, err = parseTime(updatedAt); err != nil {
|
||||
return MemoryItem{}, err
|
||||
}
|
||||
return it, nil
|
||||
}
|
||||
|
||||
// newID returns a 128-bit cryptographically random id, hex-encoded. It avoids a
|
||||
// UUID dependency so the package stays self-contained. rand.Read fails only if
|
||||
// the platform entropy source is broken, so the error is propagated.
|
||||
func newID() (string, error) {
|
||||
b := make([]byte, 16)
|
||||
if _, err := rand.Read(b); err != nil {
|
||||
return "", fmt.Errorf("memory: generate id: %w", err)
|
||||
}
|
||||
return hex.EncodeToString(b), nil
|
||||
}
|
||||
|
||||
// formatTime encodes t in UTC as an RFC3339 string. The value is always written
|
||||
// as UTC so ordering and round-trips are stable across processes.
|
||||
func formatTime(t time.Time) string {
|
||||
return t.UTC().Format(time.RFC3339Nano)
|
||||
}
|
||||
|
||||
// parseTime decodes a stored RFC3339 string. Blank input yields a zero time.
|
||||
func parseTime(s string) (time.Time, error) {
|
||||
if strings.TrimSpace(s) == "" {
|
||||
return time.Time{}, nil
|
||||
}
|
||||
t, err := time.Parse(time.RFC3339Nano, s)
|
||||
if err != nil {
|
||||
return time.Time{}, fmt.Errorf("memory: parse time %q: %w", s, err)
|
||||
}
|
||||
return t, nil
|
||||
}
|
||||
|
||||
// discardLogger returns log if non-nil, otherwise a logger that writes nowhere.
|
||||
func discardLogger(log *slog.Logger) *slog.Logger {
|
||||
if log != nil {
|
||||
return log
|
||||
}
|
||||
return slog.New(slog.NewTextHandler(io.Discard, nil))
|
||||
}
|
||||
|
||||
// Schema DDL. The migration framework in internal/db is the authoritative
|
||||
// mechanism for domain tables, but it is intentionally not used here: the db
|
||||
// package's migration tests assert an exact migration count, and the memory
|
||||
// store must be self-contained and usable from any *sql.DB. The table and
|
||||
// indexes are therefore created idempotently by NewStore.
|
||||
const createMemoryItemsTable = `
|
||||
CREATE TABLE IF NOT EXISTS memory_items (
|
||||
id TEXT PRIMARY KEY,
|
||||
persona_id TEXT NOT NULL,
|
||||
role_id TEXT NOT NULL DEFAULT '',
|
||||
scope TEXT NOT NULL,
|
||||
kind TEXT NOT NULL,
|
||||
content TEXT NOT NULL,
|
||||
tags TEXT NOT NULL DEFAULT '[]',
|
||||
confidence REAL NOT NULL DEFAULT 1.0,
|
||||
source_conversation TEXT NOT NULL DEFAULT '',
|
||||
created_at TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL,
|
||||
decay_weight REAL NOT NULL DEFAULT 1.0
|
||||
)`
|
||||
|
||||
const createMemoryItemsPersonaIndex = `
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_items_persona_id
|
||||
ON memory_items (persona_id)`
|
||||
|
||||
const createMemoryItemsPersonaCreatedIndex = `
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_items_persona_id_created_at
|
||||
ON memory_items (persona_id, created_at)`
|
||||
@@ -0,0 +1,239 @@
|
||||
package memory
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestNewStore_NilDB_异常(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
|
||||
if _, err := NewStore(ctx, nil, testLogger()); !errors.Is(err, ErrNilDB) {
|
||||
t.Fatalf("NewStore(nil db) err = %v, want ErrNilDB", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_AddGet_正常(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
item := baseItem()
|
||||
// Add returns the stored item with id and timestamps filled in.
|
||||
stored, err := s.Add(ctx, item)
|
||||
if err != nil {
|
||||
t.Fatalf("Add: %v", err)
|
||||
}
|
||||
if stored.ID == "" {
|
||||
t.Fatal("Add returned empty id")
|
||||
}
|
||||
if stored.CreatedAt.IsZero() || stored.UpdatedAt.IsZero() {
|
||||
t.Fatal("Add did not assign timestamps")
|
||||
}
|
||||
|
||||
got, err := s.Get(ctx, stored.ID)
|
||||
if err != nil {
|
||||
t.Fatalf("Get: %v", err)
|
||||
}
|
||||
if got.ID != stored.ID {
|
||||
t.Errorf("ID = %q, want %q", got.ID, stored.ID)
|
||||
}
|
||||
if got.PersonaID != item.PersonaID {
|
||||
t.Errorf("PersonaID = %q, want %q", got.PersonaID, item.PersonaID)
|
||||
}
|
||||
if got.RoleID != item.RoleID {
|
||||
t.Errorf("RoleID = %q, want %q", got.RoleID, item.RoleID)
|
||||
}
|
||||
if got.Scope != item.Scope {
|
||||
t.Errorf("Scope = %q, want %q", got.Scope, item.Scope)
|
||||
}
|
||||
if got.Kind != item.Kind {
|
||||
t.Errorf("Kind = %q, want %q", got.Kind, item.Kind)
|
||||
}
|
||||
if got.Content != item.Content {
|
||||
t.Errorf("Content = %q, want %q", got.Content, item.Content)
|
||||
}
|
||||
if got.Confidence != item.Confidence {
|
||||
t.Errorf("Confidence = %v, want %v", got.Confidence, item.Confidence)
|
||||
}
|
||||
if len(got.Tags) != 2 || got.Tags[0] != "habit" || got.Tags[1] != "coffee" {
|
||||
t.Errorf("Tags = %v, want [habit coffee]", got.Tags)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_Add_Defaults_正常(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
// Empty Scope/Kind/Confidence/DecayWeight normalize to their defaults.
|
||||
stored, err := s.Add(ctx, MemoryItem{
|
||||
PersonaID: "persona.a",
|
||||
Content: "我们约定每周日晚上看电影",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Add: %v", err)
|
||||
}
|
||||
if stored.Scope != ScopePersona {
|
||||
t.Errorf("Scope = %q, want persona default", stored.Scope)
|
||||
}
|
||||
if stored.Kind != KindRelationship {
|
||||
t.Errorf("Kind = %q, want relationship default", stored.Kind)
|
||||
}
|
||||
if stored.Confidence != 1.0 {
|
||||
t.Errorf("Confidence = %v, want 1.0 default", stored.Confidence)
|
||||
}
|
||||
if stored.DecayWeight != 1.0 {
|
||||
t.Errorf("DecayWeight = %v, want 1.0 default", stored.DecayWeight)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_Add_Invalid_异常(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
item MemoryItem
|
||||
wantErr error
|
||||
}{
|
||||
{name: "missing persona", item: MemoryItem{Content: "x"}, wantErr: ErrMissingPersona},
|
||||
{name: "missing content", item: MemoryItem{PersonaID: "p"}, wantErr: ErrMissingContent},
|
||||
{name: "invalid scope", item: MemoryItem{PersonaID: "p", Content: "x", Scope: "bogus"}, wantErr: ErrInvalidScope},
|
||||
{name: "invalid kind", item: MemoryItem{PersonaID: "p", Content: "x", Kind: "bogus"}, wantErr: ErrInvalidKind},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
tt := tt
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
if _, err := s.Add(ctx, tt.item); !errors.Is(err, tt.wantErr) {
|
||||
t.Fatalf("Add err = %v, want wrap of %v", err, tt.wantErr)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_Get_NotFound_异常(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
for _, id := range []string{"", "does-not-exist"} {
|
||||
if _, err := s.Get(ctx, id); !errors.Is(err, ErrNotFound) {
|
||||
t.Fatalf("Get(%q) err = %v, want ErrNotFound", id, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_ListByPersona_Isolation(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
for _, p := range []string{"persona.a", "persona.b"} {
|
||||
stored, err := s.Add(ctx, MemoryItem{
|
||||
PersonaID: p,
|
||||
Content: "内容 " + p,
|
||||
Scope: ScopePersona,
|
||||
Kind: KindWorldEvent,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Add for %s: %v", p, err)
|
||||
}
|
||||
if stored.PersonaID != p {
|
||||
t.Fatalf("stored persona = %q, want %q", stored.PersonaID, p)
|
||||
}
|
||||
}
|
||||
|
||||
items, err := s.ListByPersona(ctx, "persona.a")
|
||||
if err != nil {
|
||||
t.Fatalf("ListByPersona: %v", err)
|
||||
}
|
||||
if len(items) != 1 {
|
||||
t.Fatalf("ListByPersona(a) = %d items, want 1", len(items))
|
||||
}
|
||||
if items[0].PersonaID != "persona.a" {
|
||||
t.Errorf("item persona = %q, want persona.a", items[0].PersonaID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_ListByPersona_MissingPersona_异常(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
if _, err := s.ListByPersona(ctx, ""); !errors.Is(err, ErrMissingPersona) {
|
||||
t.Fatalf("ListByPersona(\"\") err = %v, want ErrMissingPersona", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_Delete_正常(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
stored, err := s.Add(ctx, baseItem())
|
||||
if err != nil {
|
||||
t.Fatalf("Add: %v", err)
|
||||
}
|
||||
if err := s.Delete(ctx, stored.ID); err != nil {
|
||||
t.Fatalf("Delete: %v", err)
|
||||
}
|
||||
if _, err := s.Get(ctx, stored.ID); !errors.Is(err, ErrNotFound) {
|
||||
t.Fatalf("Get after delete err = %v, want ErrNotFound", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_Delete_NotFound_异常(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
for _, id := range []string{"", "missing"} {
|
||||
if err := s.Delete(ctx, id); !errors.Is(err, ErrNotFound) {
|
||||
t.Fatalf("Delete(%q) err = %v, want ErrNotFound", id, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_Concurrent_并发(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
s := newTestStore(t)
|
||||
|
||||
const n = 40
|
||||
var wg sync.WaitGroup
|
||||
errs := make(chan error, n)
|
||||
for i := 0; i < n; i++ {
|
||||
wg.Add(1)
|
||||
go func(i int) {
|
||||
defer wg.Done()
|
||||
_, err := s.Add(ctx, MemoryItem{
|
||||
PersonaID: "persona.a",
|
||||
Content: "并发内容",
|
||||
Scope: ScopePersona,
|
||||
Kind: KindRelationship,
|
||||
})
|
||||
errs <- err
|
||||
}(i)
|
||||
}
|
||||
wg.Wait()
|
||||
close(errs)
|
||||
|
||||
for err := range errs {
|
||||
if err != nil {
|
||||
t.Fatalf("concurrent Add error: %v", err)
|
||||
}
|
||||
}
|
||||
items, err := s.ListByPersona(ctx, "persona.a")
|
||||
if err != nil {
|
||||
t.Fatalf("ListByPersona: %v", err)
|
||||
}
|
||||
if len(items) != n {
|
||||
t.Fatalf("ListByPersona(a) = %d items, want %d", len(items), n)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,147 @@
|
||||
package memory
|
||||
|
||||
import (
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"unicode/utf8"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
|
||||
)
|
||||
|
||||
// Working memory sizing and summarization constants.
|
||||
const (
|
||||
// defaultWorkingWindow is the default number of messages retained in the
|
||||
// sliding window when NewWorkingMemory is called with a non-positive size.
|
||||
defaultWorkingWindow = 20
|
||||
|
||||
// summaryLastTextMax is the maximum rune length of the trailing message text
|
||||
// included in a naive summary, keeping self-summaries compact.
|
||||
summaryLastTextMax = 80
|
||||
)
|
||||
|
||||
// WorkingMemory is the per-session short-term memory: a sliding window of the
|
||||
// current conversation's messages plus a summary (design §0 layer ①). It is the
|
||||
// immediate context an agent reasons over.
|
||||
//
|
||||
// The window is bounded: adding beyond the maximum evicts the oldest message.
|
||||
// Summarize returns the summary set via SetSummary, or a naive structural
|
||||
// summary computed on demand. LLM-generated summaries are deferred to M4; the
|
||||
// plumbing is in place so a summarizer can call SetSummary.
|
||||
//
|
||||
// WorkingMemory is safe for concurrent use; Add/SetSummary serialize with a
|
||||
// write lock and Window/Summarize read under a read lock.
|
||||
type WorkingMemory struct {
|
||||
mu sync.RWMutex
|
||||
window []model.Message
|
||||
max int
|
||||
summary string
|
||||
}
|
||||
|
||||
// NewWorkingMemory builds a working memory with a sliding window of up to
|
||||
// maxMessages messages. A non-positive maxMessages selects defaultWorkingWindow.
|
||||
func NewWorkingMemory(maxMessages int) *WorkingMemory {
|
||||
if maxMessages <= 0 {
|
||||
maxMessages = defaultWorkingWindow
|
||||
}
|
||||
return &WorkingMemory{max: maxMessages}
|
||||
}
|
||||
|
||||
// Add appends msg to the window and evicts the oldest message when the window
|
||||
// exceeds its limit.
|
||||
func (w *WorkingMemory) Add(msg model.Message) {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
|
||||
w.window = append(w.window, msg)
|
||||
if len(w.window) > w.max {
|
||||
overflow := len(w.window) - w.max
|
||||
copy(w.window, w.window[overflow:])
|
||||
w.window = w.window[:w.max]
|
||||
}
|
||||
}
|
||||
|
||||
// Window returns a deep copy of the current message window, oldest first. The
|
||||
// returned slice and each message's content parts are detached from internal
|
||||
// state, so callers may mutate them freely.
|
||||
func (w *WorkingMemory) Window() []model.Message {
|
||||
w.mu.RLock()
|
||||
defer w.mu.RUnlock()
|
||||
|
||||
out := make([]model.Message, len(w.window))
|
||||
for i, m := range w.window {
|
||||
out[i] = cloneMessage(m)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// cloneMessage returns a copy of m whose content parts are detached from the
|
||||
// source. Every field of model.Content is a value type, so copying the parts
|
||||
// slice is sufficient.
|
||||
func cloneMessage(m model.Message) model.Message {
|
||||
parts := make([]model.Content, len(m.Content))
|
||||
copy(parts, m.Content)
|
||||
m.Content = parts
|
||||
return m
|
||||
}
|
||||
|
||||
// SetSummary records an explicit summary (e.g. produced by an LLM in M4).
|
||||
func (w *WorkingMemory) SetSummary(s string) {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
w.summary = s
|
||||
}
|
||||
|
||||
// Summarize returns the stored summary when one was set, otherwise a naive
|
||||
// structural summary derived from the current window. It never mutates state, so
|
||||
// repeated calls return a value consistent with the window at call time.
|
||||
func (w *WorkingMemory) Summarize() string {
|
||||
w.mu.RLock()
|
||||
defer w.mu.RUnlock()
|
||||
|
||||
if w.summary != "" {
|
||||
return w.summary
|
||||
}
|
||||
return w.naiveSummaryLocked()
|
||||
}
|
||||
|
||||
// naiveSummaryLocked builds "{n} 条消息, {role} {count}, ...; 最近: {text}". It
|
||||
// requires w.mu (read or write) to be held.
|
||||
func (w *WorkingMemory) naiveSummaryLocked() string {
|
||||
if len(w.window) == 0 {
|
||||
return ""
|
||||
}
|
||||
counts := make(map[model.Role]int)
|
||||
last := ""
|
||||
for _, m := range w.window {
|
||||
counts[m.Role]++
|
||||
last = messageText(m)
|
||||
}
|
||||
|
||||
var b strings.Builder
|
||||
b.WriteString(strconv.Itoa(len(w.window)))
|
||||
b.WriteString(" 条消息")
|
||||
for _, role := range []model.Role{model.RoleSystem, model.RoleUser, model.RoleAssistant, model.RoleTool} {
|
||||
if c := counts[role]; c > 0 {
|
||||
b.WriteString(", ")
|
||||
b.WriteString(string(role))
|
||||
b.WriteByte(' ')
|
||||
b.WriteString(strconv.Itoa(c))
|
||||
}
|
||||
}
|
||||
if last != "" {
|
||||
b.WriteString("; 最近: ")
|
||||
b.WriteString(truncateRunes(last, summaryLastTextMax))
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// truncateRunes shortens s to at most n runes. Unicode-safe: it never slices a
|
||||
// multi-byte rune in half.
|
||||
func truncateRunes(s string, n int) string {
|
||||
if utf8.RuneCountInString(s) <= n {
|
||||
return s
|
||||
}
|
||||
runes := []rune(s)
|
||||
return string(runes[:n])
|
||||
}
|
||||
@@ -0,0 +1,108 @@
|
||||
package memory
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
|
||||
)
|
||||
|
||||
func TestWorkingMemory_Window_Sliding_正常(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
w := NewWorkingMemory(3)
|
||||
for i := 0; i < 5; i++ {
|
||||
w.Add(model.TextMessage(model.RoleUser, string(rune('a'+i))))
|
||||
}
|
||||
|
||||
window := w.Window()
|
||||
if len(window) != 3 {
|
||||
t.Fatalf("Window length = %d, want 3", len(window))
|
||||
}
|
||||
// The oldest two messages are evicted; the last three remain.
|
||||
for i, want := range []string{"c", "d", "e"} {
|
||||
if got := window[i].Content[0].Text; got != want {
|
||||
t.Errorf("window[%d].Text = %q, want %q", i, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestWorkingMemory_Window_IsDetached_正常(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
w := NewWorkingMemory(2)
|
||||
w.Add(model.TextMessage(model.RoleUser, "one"))
|
||||
window := w.Window()
|
||||
window[0].Content[0].Text = "mutated"
|
||||
|
||||
again := w.Window()
|
||||
if again[0].Content[0].Text != "one" {
|
||||
t.Fatalf("Window returned a shared slice; internal state mutated to %q", again[0].Content[0].Text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWorkingMemory_Summarize_Default_正常(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
w := NewWorkingMemory(10)
|
||||
w.Add(model.TextMessage(model.RoleUser, "你好"))
|
||||
w.Add(model.TextMessage(model.RoleUser, "今天天气不错"))
|
||||
w.Add(model.TextMessage(model.RoleAssistant, "是呀,很适合散步"))
|
||||
|
||||
summary := w.Summarize()
|
||||
if !strings.Contains(summary, "3 条消息") {
|
||||
t.Errorf("summary missing message count: %q", summary)
|
||||
}
|
||||
if !strings.Contains(summary, "user 2") {
|
||||
t.Errorf("summary missing user count: %q", summary)
|
||||
}
|
||||
if !strings.Contains(summary, "assistant 1") {
|
||||
t.Errorf("summary missing assistant count: %q", summary)
|
||||
}
|
||||
if !strings.Contains(summary, "很适合散步") {
|
||||
t.Errorf("summary missing last message text: %q", summary)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWorkingMemory_Summarize_SetSummary_正常(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
w := NewWorkingMemory(10)
|
||||
w.Add(model.TextMessage(model.RoleUser, "hello"))
|
||||
w.Add(model.TextMessage(model.RoleAssistant, "hi"))
|
||||
w.SetSummary("对话摘要:用户打招呼")
|
||||
|
||||
if got := w.Summarize(); got != "对话摘要:用户打招呼" {
|
||||
t.Fatalf("Summarize() = %q, want stored summary", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWorkingMemory_Summarize_Empty_正常(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
w := NewWorkingMemory(10)
|
||||
if got := w.Summarize(); got != "" {
|
||||
t.Fatalf("Summarize() on empty window = %q, want empty", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWorkingMemory_AddWindow_Concurrent_并发(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
const n = 50
|
||||
w := NewWorkingMemory(n)
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < n; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
w.Add(model.TextMessage(model.RoleUser, "x"))
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
if got := len(w.Window()); got != n {
|
||||
t.Fatalf("Window length = %d after concurrent add, want %d", got, n)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user