Cyrene_v2/backend/tools/memory-bench/main.go

package main

import (
	"context"
	"fmt"
	"log"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/db"
	"git.yeij.top/AskaEth/Cyrene_v2/services/memory"
)

func main() {
	ctx := context.Background()
	cfg := db.DefaultConfig()
	database, err := db.Connect(cfg)
	if err != nil { log.Fatalf("DB: %v", err) }
	defer database.Close()
	database.Exec(memory.CreateTableSQL())
	repo := memory.NewRepository(database)
	svc := memory.NewService(repo)
	database.Exec("DELETE FROM memories WHERE user_id LIKE 'bench-%'")

	fmt.Println("═══════════════════════════════════════════")
	fmt.Println("  Cyrene v2 Memory Benchmark")
	fmt.Println("═══════════════════════════════════════════")
	fmt.Println()

	// 1. Write Latency
	fmt.Println("📝 写入延迟 (20条)")
	t0 := time.Now()
	for i := range 20 {
		svc.Add(ctx, "bench-w", "", fmt.Sprintf("测试%d", i), "fact", 0.5)
	}
	d := time.Since(t0)
	fmt.Printf("   %.0f 条/秒  |  平均 %v\n", 20/d.Seconds(), (d/20).Round(time.Microsecond))

	// 2. Search Latency
	fmt.Println("\n🔍 搜索延迟")
	// Insert known content
	svc.Add(ctx, "bench-s", "", "开拓者喜欢喝咖啡", "preference", 0.8)
	svc.Add(ctx, "bench-s", "", "开拓者喜欢听音乐", "preference", 0.8)
	qStart := time.Now()
	r, _ := svc.Search(ctx, memory.SearchQuery{UserID: "bench-s", Query: "咖啡", Limit: 10})
	fmt.Printf("   '咖啡' → %d 条  %v\n", len(r), time.Since(qStart))

	// 3. Dedup Accuracy
	fmt.Println("\n🔁 去重准确率 (10组)")
	tests := []struct{ a, b string; expect bool }{
		{"喜欢喝咖啡", "喜欢喝咖啡", true},
		{"喜欢喝咖啡", "喜欢喝茶", false},
		{"喜欢喝咖啡", "每天喝咖啡", true},
		{"喜欢喝咖啡", "喜欢喝咖啡每天两杯", true},
		{"喜欢喝咖啡", "今天天气真好", false},
		{"开拓者喜欢听流行音乐", "开拓者喜欢听音乐", true},
		{"用户是程序员", "用户做软件开发", true},
		{"用户住在北京", "用户在东京旅游", false},
		{"喜欢吃火锅", "喜欢吃辣的", true},
		{"今天下雨了", "明天可能会下雨", false},
	}
	correct, fp, fn := 0, 0, 0
	for _, tt := range tests {
		uid := fmt.Sprintf("bench-dd-%d", time.Now().UnixNano())
		m1, _ := svc.Add(ctx, uid, "", tt.a, "fact", 0.5)
		m2, _ := svc.Add(ctx, uid, "", tt.b, "fact", 0.5)
		merged := m1 != nil && m2 != nil && m1.ID == m2.ID
		if tt.expect && merged { correct++
		} else if tt.expect && !merged { fn++
		} else if !tt.expect && merged { fp++
		} else { correct++ }
	}
	t := len(tests)
	fmt.Printf("   正确: %d/%d (%.0f%%)  漏合并: %d  误合并: %d\n", correct, t, float64(correct)/float64(t)*100, fn, fp)

	// 4. Keyword Quality
	fmt.Println("\n🏷 关键词质量")
	kwTests := []struct{ c string; want []string }{
		{"开拓者喜欢喝咖啡", []string{"开拓者", "咖啡"}},
		{"用户是程序员在北京工作", []string{"程序员", "北京"}},
	}
	for _, tt := range kwTests {
		kws := memory.ExtractKeywordsPublic(tt.c)
		hits := 0
		for _, w := range tt.want {
			for _, k := range kws[:min(8, len(kws))] {
				if k == w { hits++; break }
			}
		}
		fmt.Printf("   \"%s\" → %d/%d 命中  %v\n", tt.c, hits, len(tt.want), kws[:min(6, len(kws))])
	}

	// 5. Bulk Throughput
	fmt.Println("\n⚡ 批量吞吐 (50条)")
	database.Exec("DELETE FROM memories WHERE user_id='bench-tp'")
	tb := time.Now()
	for i := range 50 {
		svc.Add(ctx, "bench-tp", "", fmt.Sprintf("批量%d", i), "fact", 0.5)
	}
	db := time.Since(tb)
	fmt.Printf("   %.0f 条/秒\n", 50/db.Seconds())

	// 6. Real dedup rate
	fmt.Println("\n📊 实际去重率")
	database.Exec("DELETE FROM memories WHERE user_id='bench-dr'")
	variants := []string{"开拓者喜欢喝咖啡", "开拓者喜欢喝茶", "开拓者喜欢喝咖啡每天两杯", "开拓者喜欢听音乐", "开拓者在咖啡店喝拿铁", "开拓者喜欢喝手冲咖啡", "开拓者喜欢吃甜食"}
	seen := make(map[string]bool)
	for _, v := range variants {
		m, _ := svc.Add(ctx, "bench-dr", "", v, "preference", 0.5)
		seen[m.ID] = true
	}
	fmt.Printf("   %d 次写入 → %d 唯一记忆  →  去重率 %.0f%%\n", len(variants), len(seen), float64(len(variants)-len(seen))/float64(len(variants))*100)

	// Cleanup
	database.Exec("DELETE FROM memories WHERE user_id LIKE 'bench-%'")
	fmt.Println("\n═══════════════════════════════════════════")
}

func min(a, b int) int { if a < b { return a }; return b }