Cyrene_v2/internal/runtime/memory/retrieve_m4_test.go

package memory

import (
	"context"
	"testing"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/db"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/memory/projectindex"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/memory/vector"
)

// mockVectorIndex 是测试用的可脚本化 vector.Index。未设置 fns 时它
// 发出 one-hot 向量且不返回检索命中。
type mockVectorIndex struct {
	dim      int
	embedFn  func(ctx context.Context, texts []string) ([]vector.Vector, error)
	searchFn func(ctx context.Context, q vector.Vector, topK int) ([]vector.Hit, error)
}

func (m *mockVectorIndex) Dim() int {
	if m.dim > 0 {
		return m.dim
	}
	return 4
}

func (m *mockVectorIndex) Add(_ context.Context, _ string, _ vector.Vector) error { return nil }

// Vectors 默认一律未命中：这些用例验的是检索/写入接线，不是增量复用。
func (m *mockVectorIndex) Vectors(_ context.Context, ids []string) ([]vector.Vector, error) {
	return make([]vector.Vector, len(ids)), nil
}

func (m *mockVectorIndex) Close() error { return nil }

func (m *mockVectorIndex) Embed(ctx context.Context, texts []string) ([]vector.Vector, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if m.embedFn != nil {
		return m.embedFn(ctx, texts)
	}
	out := make([]vector.Vector, len(texts))
	for i := range texts {
		v := make(vector.Vector, m.Dim())
		v[i%m.Dim()] = 1
		out[i] = v
	}
	return out, nil
}

func (m *mockVectorIndex) Search(ctx context.Context, _ vector.Vector, topK int) ([]vector.Hit, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if m.searchFn != nil {
		return m.searchFn(ctx, nil, topK)
	}
	return []vector.Hit{}, nil
}

func TestRetriever_Retrieve_WithVectorIndex_正常(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	s := newTestStore(t)

	// 两条条目都是英文而 query 是中文，因此关键词路径
	// 不共享任何 token、毫无贡献。排名因此完全由
	// 向量索引驱动，证明走的是持久化向量路径。
	item1 := addItem(t, s, MemoryItem{
		PersonaID: "persona.a", Content: "she is a morning person",
		Scope: ScopePersona, Kind: KindRelationship,
	})
	item2 := addItem(t, s, MemoryItem{
		PersonaID: "persona.a", Content: "she likes jazz",
		Scope: ScopePersona, Kind: KindRelationship,
	})

	idx := &mockVectorIndex{
		searchFn: func(_ context.Context, _ vector.Vector, _ int) ([]vector.Hit, error) {
			return []vector.Hit{
				{ID: item1.ID, Score: 0.9},
				{ID: item2.ID, Score: 0.2},
			}, nil
		},
	}
	r := NewRetriever(s, WithVectorIndex(idx))

	results, err := r.Retrieve(ctx, "persona.a", "咖啡", 10)
	if err != nil {
		t.Fatalf("Retrieve: %v", err)
	}
	if len(results) == 0 {
		t.Fatal("expected results with a vector index wired")
	}
	if got := results[0].Item.Content; got != "she is a morning person" {
		t.Fatalf("top result = %q, want the item the vector index ranked highest", got)
	}
	if results[0].Score <= 0 {
		t.Fatalf("top result score = %v, want > 0", results[0].Score)
	}
}

func TestRetriever_Retrieve_WithVectorIndex_PartialHits_正常(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	s := newTestStore(t)

	item1 := addItem(t, s, MemoryItem{
		PersonaID: "persona.a", Content: "she is a morning person",
		Scope: ScopePersona, Kind: KindRelationship,
	})
	item2 := addItem(t, s, MemoryItem{
		PersonaID: "persona.a", Content: "she likes jazz",
		Scope: ScopePersona, Kind: KindRelationship,
	})

	// 向量索引只知道 item1；item2 不在其中，因此它
	// 没有语义得分。在关键词路径无法匹配的中文 query 下，
	// 只有 item1 相关：它必须以正得分排在
	// 首位，而 item2 得 0 分，证明语义贡献
	// 严格来自向量索引。
	idx := &mockVectorIndex{
		searchFn: func(_ context.Context, _ vector.Vector, _ int) ([]vector.Hit, error) {
			return []vector.Hit{{ID: item1.ID, Score: 0.9}}, nil
		},
	}
	r := NewRetriever(s, WithVectorIndex(idx))

	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 != "she is a morning person" {
		t.Fatalf("top result = %q, want the vector-indexed item", got)
	}
	if results[0].Score <= 0 {
		t.Fatalf("top result score = %v, want > 0", results[0].Score)
	}
	// item2 既无向量命中也无关键词重叠，因此得 0 分（已匹配
	// 模式）并排在后面。
	if results[1].Item.ID != item2.ID {
		t.Fatalf("second result id = %q, want item2", results[1].Item.ID)
	}
	if results[1].Score != 0 {
		t.Fatalf("second result score = %v, want 0 (no vector hit)", results[1].Score)
	}
}

func TestRetriever_Retrieve_Source_正常(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	s := newTestStore(t)

	addItem(t, s, MemoryItem{
		PersonaID:          "persona.a",
		Content:            "她喜欢冰美式",
		Tags:               []string{"coffee"},
		SourceConversation: "conv.42",
		Scope:              ScopePersona,
		Kind:               KindRelationship,
	})

	r := NewRetriever(s)
	results, err := r.Retrieve(ctx, "persona.a", "冰美式", 10)
	if err != nil {
		t.Fatalf("Retrieve: %v", err)
	}
	if len(results) == 0 {
		t.Fatal("expected a result")
	}
	if results[0].SourceConversation != "conv.42" {
		t.Fatalf("Result.SourceConversation = %q, want conv.42 (溯源)", results[0].SourceConversation)
	}
	if results[0].Item.SourceConversation != "conv.42" {
		t.Fatalf("Item.SourceConversation = %q, want conv.42", results[0].Item.SourceConversation)
	}
}

func TestRetriever_Retrieve_ProjectLink_正常(t *testing.T) {
	t.Parallel()
	ctx := context.Background()

	// 独立的项目索引 DB，因此与记忆存储不共享连接。
	projDB, err := db.Open(ctx, ":memory:", testLogger())
	if err != nil {
		t.Fatalf("open project db: %v", err)
	}
	t.Cleanup(func() { _ = projDB.Close() })
	ps, err := projectindex.New(ctx, projDB, testLogger())
	if err != nil {
		t.Fatalf("projectindex.New: %v", err)
	}
	if err := ps.PutProject(ctx, projectindex.Project{
		ID: "project.alpha", Title: "Alpha 项目", Desc: "Alpha 项目描述",
	}); err != nil {
		t.Fatalf("PutProject: %v", err)
	}

	ms := newTestStore(t)
	addItem(t, ms, MemoryItem{
		PersonaID: "persona.a",
		Content:   "Alpha 项目的一个事实",
		Scope:     ScopeProject,
		ProjectID: "project.alpha",
		Kind:      KindWorldEvent,
	})

	r := NewRetriever(ms, WithProjectIndex(ps))
	results, err := r.Retrieve(ctx, "persona.a", "Alpha", 10)
	if err != nil {
		t.Fatalf("Retrieve: %v", err)
	}
	if len(results) == 0 {
		t.Fatal("expected a result")
	}
	if results[0].Item.Scope != ScopeProject {
		t.Fatalf("result scope = %q, want project", results[0].Item.Scope)
	}
	if results[0].Project == nil {
		t.Fatal("Result.Project is nil, want the linked project-index entry")
	}
	if results[0].Project.ID != "project.alpha" {
		t.Fatalf("Result.Project.ID = %q, want project.alpha", results[0].Project.ID)
	}
}

func TestRetriever_Retrieve_ProjectLink_NoStore_正常(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	ms := newTestStore(t)

	addItem(t, ms, MemoryItem{
		PersonaID: "persona.a", Content: "Alpha 项目的一个事实",
		Scope: ScopeProject, ProjectID: "project.alpha", Kind: KindWorldEvent,
	})

	// 未接入项目索引存储时，命中仍会返回，Project
	// 只是不设置（不报错）。
	r := NewRetriever(ms)
	results, err := r.Retrieve(ctx, "persona.a", "Alpha", 10)
	if err != nil {
		t.Fatalf("Retrieve: %v", err)
	}
	if len(results) == 0 {
		t.Fatal("expected a result")
	}
	if results[0].Project != nil {
		t.Fatalf("Result.Project = %+v, want nil without a project-index store", results[0].Project)
	}
}