250 lines
7.4 KiB
Go
250 lines
7.4 KiB
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 is a scriptable vector.Index for tests. With no fns it emits
|
|
// one-hot vectors and returns no search hits.
|
|
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 }
|
|
|
|
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)
|
|
|
|
// Both items are in English while the query is Chinese, so the keyword path
|
|
// shares no token and contributes nothing. The ranking is therefore driven
|
|
// entirely by the vector index, proving the persistent vector path is used.
|
|
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,
|
|
})
|
|
|
|
// The vector index only knows about item1; item2 is absent from it, so it
|
|
// carries no semantic score. With a Chinese query that the keyword path
|
|
// cannot match, only item1 is relevant: it must rank at the top with a
|
|
// positive score while item2 scores 0, proving the semantic contribution
|
|
// comes strictly from the vector index.
|
|
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 has no vector hit and no keyword overlap, so it scores 0 (matched
|
|
// mode) and sorts below.
|
|
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()
|
|
|
|
// Separate project-index DB, so it shares no connection with the memory store.
|
|
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,
|
|
})
|
|
|
|
// Without a project-index store wired, the hit still returns and Project is
|
|
// simply left unset (no error).
|
|
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)
|
|
}
|
|
}
|