package pipeline import ( "fmt" "os" "path/filepath" "testing" "time" "gitcode.com/JianFeeeee/HomeAgent/internal/memory" "gitcode.com/JianFeeeee/HomeAgent/internal/memory/vector" ) // constVectorizer 恒等向量器:所有候选统一过融合阈值,验证嵌入接线是否生效。 type constVectorizer struct{} func (constVectorizer) Vectorize(text string) vector.Vector { return vector.Vector{"0": 1.0, "1": 0.5} } func TestNewDistiller(t *testing.T) { d := NewDistiller(nil, t.TempDir(), DistillerConfig{ Interval: 10 * time.Minute, RetentionDays: 7, BatchSize: 50, }) if d == nil { t.Fatal("expected non-nil distiller") } if d.cfg.Interval != 10*time.Minute { t.Errorf("expected interval 10m, got %v", d.cfg.Interval) } } func TestAppendAndFlush(t *testing.T) { dir := t.TempDir() d := NewDistiller(nil, dir, DistillerConfig{ Interval: 10 * time.Minute, RetentionDays: 7, BatchSize: 50, }) // flush 需要 rawPath 目录已存在(Start() 负责创建) os.MkdirAll(d.rawPath, 0755) d.Append("sess1", "user", "今天天气怎么样?") d.Append("sess1", "assistant", "今天天气很好,适合出行。") if len(d.records) != 2 { t.Fatalf("expected 2 records, got %d", len(d.records)) } d.flush() // 验证 JSONL 已写入 entries, err := os.ReadDir(d.rawPath) if err != nil { t.Fatal(err) } if len(entries) == 0 { t.Fatal("expected flushed files") } } func TestLoadExisting(t *testing.T) { dir := t.TempDir() // 先写一个文件再创建 Distiller 验证 load d1 := NewDistiller(nil, dir, DistillerConfig{ Interval: 10 * time.Minute, RetentionDays: 7, BatchSize: 50, }) os.MkdirAll(d1.rawPath, 0755) d1.Append("sess1", "user", "test content") d1.flush() d2 := NewDistiller(nil, dir, DistillerConfig{ Interval: 10 * time.Minute, RetentionDays: 7, BatchSize: 50, }) d2.loadExisting() if len(d2.records) != 1 { t.Fatalf("expected 1 record after load, got %d", len(d2.records)) } } func TestDistillOnce(t *testing.T) { db, err := memory.NewGraphDB(filepath.Join(t.TempDir(), "test.db")) if err != nil { t.Fatal(err) } defer db.Close() dir := t.TempDir() d := NewDistiller(db, dir, DistillerConfig{ Interval: 10 * time.Minute, RetentionDays: 7, BatchSize: 50, }) // 写入一条旧记录(超过保留期) past := time.Now().Add(-8 * 24 * time.Hour) d.records = append(d.records, RawRecord{ ID: 1, SessionID: "sess1", Role: "user", Content: "我的名字是张三", CreatedAt: past, Distilled: false, }) d.records = append(d.records, RawRecord{ ID: 2, SessionID: "sess1", Role: "assistant", Content: "你好张三!", CreatedAt: past, Distilled: false, }) d.distillOnce() // 验证蒸馏后已有记录被标记(distillOnce 从 records 中移除已蒸馏记录,检查剩余数量) if len(d.records) != 0 { t.Logf("records after distill: %d (all should have been removed)", len(d.records)) } } // Phase 4: 新记录无需等待 RetentionDays,下一 tick 立即蒸馏(文档所述 10min 频率) func TestDistillOnceFreshRecords(t *testing.T) { db, err := memory.NewGraphDB(filepath.Join(t.TempDir(), "test.db")) if err != nil { t.Fatal(err) } defer db.Close() dir := t.TempDir() d := NewDistiller(db, dir, DistillerConfig{ Interval: 10 * time.Minute, RetentionDays: 7, BatchSize: 50, }) d.Append("sess1", "user", "我的名字是李四") d.Append("sess1", "assistant", "你好李四!") if len(d.records) != 2 { t.Fatalf("expected 2 fresh records, got %d", len(d.records)) } d.distillOnce() if len(d.records) != 0 { t.Errorf("fresh records should be distilled on next tick (no retention gate), got %d remaining", len(d.records)) } // 二次蒸馏不重复(已蒸馏记录已被移除) d.distillOnce() if len(d.records) != 0 { t.Errorf("second distill should be no-op, got %d records", len(d.records)) } } // Phase 4: BatchSize 限制每 tick 处理前 N 条,未蒸馏记录留待下个 tick func TestDistillOnceBatchLimit(t *testing.T) { db, err := memory.NewGraphDB(filepath.Join(t.TempDir(), "test.db")) if err != nil { t.Fatal(err) } defer db.Close() dir := t.TempDir() d := NewDistiller(db, dir, DistillerConfig{ Interval: 10 * time.Minute, RetentionDays: 7, BatchSize: 3, }) for i := 0; i < 10; i++ { d.Append("sess1", "user", fmt.Sprintf("第 %d 条消息内容", i)) } d.distillOnce() if len(d.records) != 7 { t.Fatalf("expected 7 records remaining after batch 3, got %d", len(d.records)) } // 后续 tick 继续消化,最终全部蒸馏 for i := 0; i < 5 && len(d.records) > 0; i++ { d.distillOnce() } if len(d.records) != 0 { t.Errorf("all records should be distilled after several ticks, got %d remaining", len(d.records)) } } func TestExtractKeyTriples(t *testing.T) { tests := []struct { user string assistant string check func([]memory.Triple) bool }{ { user: "我住在北京", check: func(triples []memory.Triple) bool { for _, tr := range triples { if tr.Subject == "我" && tr.Relation == "住" && tr.Object == "北京" { return true } } return false }, }, { user: "我在杭州读书", assistant: "好的", check: func(triples []memory.Triple) bool { for _, tr := range triples { if tr.Subject == "我" && tr.Relation == "读书" && tr.Object == "杭州" { return true } } return false }, }, { user: "今天天气真好", check: func(triples []memory.Triple) bool { return true // NLP 提取器可能不提取形容词谓语句,0 个也没关系 }, }, } for _, tt := range tests { triples := extractKeyTriples(tt.user, tt.assistant, constVectorizer{}) if tt.check != nil && !tt.check(triples) { t.Errorf("extractKeyTriples(%q) = %v, check failed", tt.user, triples) } } } func TestDistillerGetRecentRecords(t *testing.T) { d := NewDistiller(nil, t.TempDir(), DistillerConfig{}) d.Append("s1", "user", "a") d.Append("s1", "user", "b") d.Append("s1", "user", "c") recent := d.GetRecentRecords(2) if len(recent) != 2 { t.Fatalf("expected 2 recent, got %d", len(recent)) } if recent[0].Content != "b" || recent[1].Content != "c" { t.Errorf("expected [b c], got %v", recent) } } func TestDistillerStats(t *testing.T) { d := NewDistiller(nil, t.TempDir(), DistillerConfig{ Interval: 5 * time.Minute, RetentionDays: 3, BatchSize: 10, }) d.Append("s1", "user", "hello") stats := d.Stats() if stats["raw_records"] != 1 { t.Errorf("expected 1 raw record, got %v", stats["raw_records"]) } } func TestTruncate(t *testing.T) { if truncate("hello world", 5) != "hello..." { t.Errorf("expected 'hello...', got %q", truncate("hello world", 5)) } if truncate("hi", 10) != "hi" { t.Errorf("expected 'hi', got %q", truncate("hi", 10)) } }