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