Cyrene_v2/internal/runtime/memory/retrieve_test.go

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