diff --git a/internal/agent/core/agent.go b/internal/agent/core/agent.go index 6f7a253..5ff3559 100644 --- a/internal/agent/core/agent.go +++ b/internal/agent/core/agent.go @@ -147,8 +147,9 @@ type AgentConfig struct { PluginReg *plugin.Registry PluginDir string DistillInterval time.Duration - MaxContextSize int // 活跃上下文最大条数,超出按相关性裁剪 - ContextSavePath string // 上下文持久化路径,空则不持久化 + MaxContextSize int // 活跃上下文最大条数,超出按相关性裁剪 + ContextSavePath string // 上下文持久化路径,空则不持久化 + EmbeddingModelPath string // 预训练词嵌入模型路径(word2vec 文本格式),空则不使用 StageHost *StageHost EventBus *events.Bus ThinkingEnabled bool @@ -174,7 +175,7 @@ func New(cfg AgentConfig) *Agent { indexer: cfg.Indexer, skills: cfg.Skills, tracker: cfg.Tracker, - context: NewRelevanceContext(cfg.ContextSavePath), + context: NewRelevanceContext(cfg.ContextSavePath, memory.NewStaticEmbedder(strings.Split(cfg.EmbeddingModelPath, ",")...)), systemPrompt: cfg.SystemPrompt, ctx: ctx, cancel: cancel, diff --git a/internal/agent/core/context.go b/internal/agent/core/context.go index b7a7ec0..aa267dc 100644 --- a/internal/agent/core/context.go +++ b/internal/agent/core/context.go @@ -15,32 +15,29 @@ import ( "gitcode.com/JianFeeeee/HomeAgent/internal/memory/vector" ) -// ContextEvent — 单条上下文事件 type ContextEvent struct { - Timestamp time.Time `json:"timestamp"` - Source string `json:"source"` - Input string `json:"input"` - Response string `json:"response,omitempty"` - ToolsUsed []string `json:"tools_used,omitempty"` - Vector vector.Vector `json:"-"` // 缓存向量,避免重复计算 + Timestamp time.Time `json:"timestamp"` + Source string `json:"source"` + Input string `json:"input"` + Response string `json:"response,omitempty"` + ToolsUsed []string `json:"tools_used,omitempty"` + Vector vector.Vector `json:"-"` } const contextFlushInterval = 5 * time.Second -// RelevanceContext — 基于相关性的上下文管理,非固定阈值 type RelevanceContext struct { mu sync.Mutex events []*ContextEvent - embedder *memory.LocalWordEmbedder - trained bool + embedder *memory.StaticEmbedder savePath string saveTimer *time.Timer dirty bool } -func NewRelevanceContext(savePath string) *RelevanceContext { +func NewRelevanceContext(savePath string, embedder *memory.StaticEmbedder) *RelevanceContext { rc := &RelevanceContext{ - embedder: memory.NewLocalWordEmbedder(), + embedder: embedder, savePath: savePath, } if savePath != "" { @@ -49,7 +46,6 @@ func NewRelevanceContext(savePath string) *RelevanceContext { return rc } -// load 从文件恢复上下文事件 func (c *RelevanceContext) load() { data, err := os.ReadFile(c.savePath) if err != nil { @@ -60,12 +56,27 @@ func (c *RelevanceContext) load() { return } for _, evt := range events { - evt.Vector = c.embedder.Vectorize(evt.Input + " " + evt.Response) + evt.Input = memory.CleanTemplateText(evt.Input) + evt.Vector = c.computeVector(evt) } c.events = events } -// Save 持久化上下文事件到文件 +func textForVector(evt *ContextEvent) string { + switch { + case evt.Source == "agent" && evt.Response != "": + return memory.CleanTemplateText(evt.Response) + case evt.Source == "cold_storage": + return memory.CleanTemplateText(evt.Input + " " + evt.Response) + default: + return memory.CleanTemplateText(evt.Input) + } +} + +func (c *RelevanceContext) computeVector(evt *ContextEvent) vector.Vector { + return c.embedder.Vectorize(textForVector(evt)) +} + func (c *RelevanceContext) Save() error { if c.savePath == "" { return nil @@ -84,15 +95,13 @@ func (c *RelevanceContext) Append(evt ContextEvent) { c.mu.Lock() defer c.mu.Unlock() - evt.Vector = c.embedder.Vectorize(evt.Input + " " + evt.Response) + evt.Input = memory.CleanTemplateText(evt.Input) + evt.Vector = c.computeVector(&evt) c.events = append(c.events, &evt) - c.trained = false - c.save() } -// save 无锁版本,Append/Prune 内部持有锁时调用。带 debounce,每 5s 写一次盘。 func (c *RelevanceContext) save() error { if c.savePath == "" { return nil @@ -124,8 +133,6 @@ func (c *RelevanceContext) flush() { c.dirty = false } -// Prune — 基于当前输入计算每条上下文的相关性,归档最不相关的 -// 返回被归档的事件(转为文档),保留 topK 个最相关的在活跃上下文中 func (c *RelevanceContext) Prune(currentInput string, topK int, docStore *document.Store) int { c.mu.Lock() defer c.mu.Unlock() @@ -134,23 +141,19 @@ func (c *RelevanceContext) Prune(currentInput string, topK int, docStore *docume return 0 } - // 保护最近 10 条记录不被淘汰,从更早的记录中选择淘汰对象 - protectCount := 10 - if protectCount > len(c.events) { - protectCount = len(c.events) + pCount := 10 + if pCount > len(c.events) { + pCount = len(c.events) } - protected := c.events[len(c.events)-protectCount:] - candidates := c.events[:len(c.events)-protectCount] + protected := c.events[len(c.events)-pCount:] + candidates := c.events[:len(c.events)-pCount] if len(candidates) == 0 { return 0 } - c.ensureTrained() + queryVec := c.embedder.VectorizeClean(currentInput) - queryVec := c.embedder.Vectorize(currentInput) - - // 计算每条候选上下文与当前输入的相关性 type scored struct { event *ContextEvent score float64 @@ -162,12 +165,10 @@ func (c *RelevanceContext) Prune(currentInput string, topK int, docStore *docume scoredEvents[i] = scored{event: evt, score: score, idx: i} } - // 按相关性从高到低排序 sort.Slice(scoredEvents, func(i, j int) bool { return scoredEvents[i].score > scoredEvents[j].score }) - // 从候选中选 topK 最相关的保留,其余淘汰 keepCount := topK - len(protected) if keepCount < 0 { keepCount = 0 @@ -178,19 +179,16 @@ func (c *RelevanceContext) Prune(currentInput string, topK int, docStore *docume } archive := scoredEvents[keepCount:] - // 重建 events 为保留的候选 + 受保护的最新记录 c.events = make([]*ContextEvent, 0, len(keep)+len(protected)) for _, s := range keep { c.events = append(c.events, s.event) } c.events = append(c.events, protected...) - // 按时间重新排序 sort.Slice(c.events, func(i, j int) bool { return c.events[i].Timestamp.Before(c.events[j].Timestamp) }) - // 归档到文档记忆 archived := 0 if docStore != nil && len(archive) > 0 { entries := make([]document.ContextEntry, len(archive)) @@ -213,7 +211,6 @@ func (c *RelevanceContext) Prune(currentInput string, topK int, docStore *docume return archived } -// Format — 输出活跃上下文的文本,用于注入 prompt func (c *RelevanceContext) Format() string { c.mu.Lock() defer c.mu.Unlock() @@ -234,7 +231,6 @@ func (c *RelevanceContext) Format() string { return sb.String() } -// Recent — 返回最近 n 条 func (c *RelevanceContext) Recent(n int) []ContextEvent { c.mu.Lock() defer c.mu.Unlock() @@ -249,23 +245,8 @@ func (c *RelevanceContext) Recent(n int) []ContextEvent { return result } -// Len — 当前上下文事件数 func (c *RelevanceContext) Len() int { c.mu.Lock() defer c.mu.Unlock() return len(c.events) } - -func (c *RelevanceContext) ensureTrained() { - if !c.trained && len(c.events) > 0 { - texts := make([]string, len(c.events)) - for i, evt := range c.events { - texts[i] = evt.Input + " " + evt.Response - } - c.embedder.Train(texts) - for _, evt := range c.events { - evt.Vector = c.embedder.Vectorize(evt.Input + " " + evt.Response) - } - c.trained = true - } -} diff --git a/internal/agent/core/context_test.go b/internal/agent/core/context_test.go index 8ed489a..bbb6f97 100644 --- a/internal/agent/core/context_test.go +++ b/internal/agent/core/context_test.go @@ -1,12 +1,19 @@ package core import ( + "os" "testing" "time" + + "gitcode.com/JianFeeeee/HomeAgent/internal/memory" ) +func newTestCtx() *RelevanceContext { + return NewRelevanceContext("", memory.NewStaticEmbedder("")) +} + func TestContextAppendAndLen(t *testing.T) { - ctx := NewRelevanceContext("") + ctx := newTestCtx() if ctx.Len() != 0 { t.Errorf("new context should be empty, got %d", ctx.Len()) } @@ -18,7 +25,7 @@ func TestContextAppendAndLen(t *testing.T) { } func TestContextRecent(t *testing.T) { - ctx := NewRelevanceContext("") + ctx := newTestCtx() ctx.Append(ContextEvent{Timestamp: time.Now(), Source: "user", Input: "a"}) ctx.Append(ContextEvent{Timestamp: time.Now(), Source: "user", Input: "b"}) ctx.Append(ContextEvent{Timestamp: time.Now(), Source: "user", Input: "c"}) @@ -33,7 +40,7 @@ func TestContextRecent(t *testing.T) { } func TestContextFormat(t *testing.T) { - ctx := NewRelevanceContext("") + ctx := newTestCtx() f := ctx.Format() if f != "" { t.Errorf("empty context should format to empty string, got %q", f) @@ -54,7 +61,7 @@ func TestContextFormat(t *testing.T) { } func TestContextPruneKeepsTopK(t *testing.T) { - ctx := NewRelevanceContext("") + ctx := newTestCtx() for i := 0; i < 20; i++ { ctx.Append(ContextEvent{ Timestamp: time.Now(), @@ -63,7 +70,6 @@ func TestContextPruneKeepsTopK(t *testing.T) { Response: "是的天气不错", }) } - // 加一条不同主题的 ctx.Append(ContextEvent{ Timestamp: time.Now(), Source: "user", @@ -71,17 +77,16 @@ func TestContextPruneKeepsTopK(t *testing.T) { Response: "好的我来算", }) - archived := ctx.Prune("微积分", 5, nil) // nil docStore → 不归档,只裁剪 + archived := ctx.Prune("微积分", 5, nil) _ = archived - // protectCount=10 + topK=5 → 最多保留 15 if ctx.Len() > 15 { t.Errorf("after prune to 5, len should be ≤15, got %d", ctx.Len()) } } func TestContextPruneWithDocStore(t *testing.T) { - ctx := NewRelevanceContext("") + ctx := newTestCtx() for i := 0; i < 15; i++ { ctx.Append(ContextEvent{ Timestamp: time.Now(), @@ -98,8 +103,7 @@ func TestContextPruneWithDocStore(t *testing.T) { } func TestContextAppendAfterPrune(t *testing.T) { - ctx := NewRelevanceContext("") - // 需要超过 protectCount(10) + topK(3) 个事件才能产生修剪候选 + ctx := newTestCtx() for i := 0; i < 20; i++ { ctx.Append(ContextEvent{ Timestamp: time.Now(), @@ -109,7 +113,7 @@ func TestContextAppendAfterPrune(t *testing.T) { } ctx.Prune("hello", 3, nil) - if ctx.Len() > 13 { // 10 protected + 3 topK + if ctx.Len() > 13 { t.Errorf("expected ≤13 after prune, got %d", ctx.Len()) } @@ -119,6 +123,107 @@ func TestContextAppendAfterPrune(t *testing.T) { } } +func TestContextPruneWithStaticEmbedder(t *testing.T) { + tmpFile, err := os.CreateTemp("", "test_embeddings_*.txt") + if err != nil { + t.Fatal(err) + } + defer os.Remove(tmpFile.Name()) + + content := `8 4 +天气 0.1 0.2 0.3 0.4 +下雨 0.15 0.25 0.35 0.45 +台风 0.12 0.22 0.32 0.42 +股票 0.9 0.1 0.1 0.1 +基金 0.85 0.15 0.1 0.1 +微积分 0.1 0.1 0.9 0.1 +导数 0.15 0.1 0.85 0.15 +数学 0.1 0.1 0.8 0.2 +` + if _, err := tmpFile.WriteString(content); err != nil { + t.Fatal(err) + } + tmpFile.Close() + + embedder := memory.NewStaticEmbedder(tmpFile.Name()) + if !embedder.Loaded() { + t.Fatal("embedder should be loaded") + } + + ctx := NewRelevanceContext("", embedder) + + ctx.Append(ContextEvent{Timestamp: time.Now(), Source: "user", Input: "今天天气很好", Response: "是的"}) + ctx.Append(ContextEvent{Timestamp: time.Now(), Source: "user", Input: "会不会下雨", Response: "会"}) + ctx.Append(ContextEvent{Timestamp: time.Now(), Source: "user", Input: "台风来了", Response: "注意"}) + ctx.Append(ContextEvent{Timestamp: time.Now(), Source: "user", Input: "帮我算微积分", Response: "好的"}) + ctx.Append(ContextEvent{Timestamp: time.Now(), Source: "user", Input: "导数怎么求", Response: "公式"}) + ctx.Append(ContextEvent{Timestamp: time.Now(), Source: "user", Input: "数学题", Response: "解答"}) + ctx.Append(ContextEvent{Timestamp: time.Now(), Source: "user", Input: "股票涨了", Response: "恭喜"}) + ctx.Append(ContextEvent{Timestamp: time.Now(), Source: "user", Input: "基金定投", Response: "可以"}) + + if ctx.Len() != 8 { + t.Fatalf("expected 8 events, got %d", ctx.Len()) + } + + archived := ctx.Prune("最近基金怎么样", 3, nil) + + if ctx.Len() > 13 { + t.Errorf("prune should limit total events, got %d", ctx.Len()) + } + + remaining := ctx.Format() + t.Logf("query: 最近基金怎么样\nremaining events:\n%s", remaining) + t.Logf("archived: %d", archived) + + needsFund := contains(remaining, "基金定投") || contains(remaining, "股票涨了") + needsWeather := contains(remaining, "今天天气很好") || contains(remaining, "台风来了") + + t.Logf("has financial events: %v, has weather events: %v", needsFund, needsWeather) +} + +func TestContextPruneRecent10Protected(t *testing.T) { + ctx := newTestCtx() + + for i := 0; i < 15; i++ { + ctx.Append(ContextEvent{ + Timestamp: time.Now(), + Source: "user", + Input: "今天天气很好", + }) + } + + for i := 0; i < 5; i++ { + ctx.Append(ContextEvent{ + Timestamp: time.Now(), + Source: "user", + Input: "股票行情", + }) + } + + ctx.Prune("天气", 3, nil) + + // 最近 10 条全部是"股票行情"(第6-15条是天气,第16-20条是股票) + // protectCount=10 保护最近 10 条 → 5 条天气最多保留 5+3=8 条 + // 至少最近 10 条全部保留 → 至少包含 5 条股票 + remaining := ctx.Format() + t.Logf("after weather query:\n%s", remaining) + weatherCount := 0 + stockCount := 0 + for _, line := range splitLines(remaining) { + if contains(line, "天气") { + weatherCount++ + } + if contains(line, "股票") { + stockCount++ + } + } + t.Logf("weather events: %d, stock events: %d", weatherCount, stockCount) + + if stockCount < 5 { + t.Errorf("recent 10 should all be protected, expected at least 5 stock events, got %d", stockCount) + } +} + func contains(s, substr string) bool { return len(s) >= len(substr) && containsStr(s, substr) } @@ -131,3 +236,18 @@ func containsStr(s, substr string) bool { } return false } + +func splitLines(s string) []string { + var lines []string + start := 0 + for i := 0; i < len(s); i++ { + if s[i] == '\n' { + lines = append(lines, s[start:i]) + start = i + 1 + } + } + if start < len(s) { + lines = append(lines, s[start:]) + } + return lines +} diff --git a/internal/memory/bilingual_test.go b/internal/memory/bilingual_test.go new file mode 100644 index 0000000..010169d --- /dev/null +++ b/internal/memory/bilingual_test.go @@ -0,0 +1,243 @@ +package memory + +import ( + "fmt" + "sort" + "testing" +) + +type bilingualEvent struct { + idx int + source string + topic string + text string // cleaned text for vectorization + label string // short description +} + +func TestBilingualPruningAccuracy(t *testing.T) { + zhPath := "/tmp/cc.zh.top200k.vec" + enPath := "/tmp/cc.en.top200k.vec" + + // Test with Chinese-only vs Chinese+English + type modelConfig struct { + name string + paths []string + } + + configs := []modelConfig{ + {"中文仅", []string{zhPath}}, + {"中文+英文", []string{zhPath, enPath}}, + } + + events := genBilingualEvents() + + for _, cfg := range configs { + t.Run(cfg.name, func(t *testing.T) { + e := NewStaticEmbedder(cfg.paths...) + if !e.Loaded() { + t.Skipf("%s: embedder not loaded", cfg.name) + } + t.Logf("%s: %d words", cfg.name, len(e.words)) + + type scored struct { + idx int + topic string + label string + score float64 + } + + queries := []struct { + q string + qTopic string + desc string + }{ + {"老大说了关于 React 组件的事情", "老大私聊", "中英混合:老大+React"}, + {"帮我查一下 Nginx 反向代理配置", "服务器运维", "中英混合:Nginx+反向代理"}, + {"河南医药大学 Docker 部署", "大学招生", "中英混合:大学+Docker"}, + {"JavaScript 基金定投收益计算", "股票基金", "中英混合:JS+基金"}, + {"Server 前端组件封装 layout", "前端开发", "中英混合:Server+layout"}, + {"河南医药大学录取分数线", "大学招生", "纯中文:大学"}, + {"nginx reverse proxy config", "服务器运维", "纯英文:nginx"}, + } + + for _, q := range queries { + qVec := e.Vectorize(q.q) + t.Logf("\n query: %q (%s)", q.q, q.desc) + + all := make([]scored, len(events)) + for i, ev := range events { + text := textForBilingual(ev, cfg.paths) + vec := e.Vectorize(text) + all[i] = scored{idx: i, topic: ev.topic, label: ev.label, score: cosineSim(qVec, vec)} + } + sort.Slice(all, func(i, j int) bool { return all[i].score > all[j].score }) + + // Check top 5 for same-topic presence + var intraHits int + for _, s := range all[:5] { + if s.topic == q.qTopic { + intraHits++ + } + } + + topScore := all[0] + topIsCorrect := topScore.topic == q.qTopic + + t.Logf(" top5 intra=%d/5, top1=%q(%s) score=%.4f %s", + intraHits, topScore.topic, topScore.label, topScore.score, + map[bool]string{true: "✅", false: "❌"}[topIsCorrect]) + for _, s := range all[:5] { + mark := "" + if s.topic == q.qTopic { + mark = " ✓" + } + t.Logf(" [%.4f] [%-12s] %s%s", s.score, s.topic, trimLen(s.label, 50), mark) + } + + if !topIsCorrect { + t.Logf(" [WARN] top1 mismatch for %q", q.desc) + } + } + }) + } +} + +func TestBilingualCrossLingualSimilarity(t *testing.T) { + zhPath := "/tmp/cc.zh.top200k.vec" + enPath := "/tmp/cc.en.top200k.vec" + e := NewStaticEmbedder(zhPath, enPath) + if !e.Loaded() { + t.Skip("embedder not loaded") + } + + pairs := []struct { + a, b string + desc string + }{ + {"server", "服务器", "英中同义"}, + {"computer", "电脑", "英中同义"}, + {"老大", "boss", "中英同义"}, + {"大学", "university", "中英同义"}, + {"Nginx", "服务器", "专名+普通"}, + {"股票", "stock", "中英同义"}, + {"React", "前端", "专名+概念"}, + {"JavaScript", "编程", "专名+概念"}, + {"老大私聊", "boss private chat", "中英短语"}, + {"河南医药大学录取", "Henan Medical University admission", "中英专名"}, + {"nginx config", "Nginx 配置", "英中技术"}, + } + + t.Log("=== 跨语言相似度 ===") + for _, p := range pairs { + va := e.Vectorize(p.a) + vb := e.Vectorize(p.b) + sim := cosineSim(va, vb) + t.Logf(" %.4f %q ↔ %q [%s]", sim, trimLen(p.a, 30), trimLen(p.b, 30), p.desc) + } +} + +func TestBilingualEdgeCases(t *testing.T) { + zhPath := "/tmp/cc.zh.top200k.vec" + e := NewStaticEmbedder(zhPath) + if !e.Loaded() { + t.Skip("embedder not loaded") + } + + cases := []string{ + "纯英文文本 nginx react docker javascript", + "纯中文 服务器 配置 反向代理 部署", + "中英混合 nginx 反向代理 配置", + "代码片段 const foo = 'bar'; function test()", + "URL路径 /api/v1/users/123", + "中文含标点!@#¥%……&*()", + "空字符串", + } + + t.Log("=== 边缘情况向量化 ===") + for _, c := range cases { + v := e.Vectorize(c) + var dims int + for range v { + dims++ + } + t.Logf(" dims=%d %q", dims, trimLen(c, 60)) + } +} + +func TestBilingualVectorizeClean(t *testing.T) { + zhPath := "/tmp/cc.zh.top200k.vec" + e := NewStaticEmbedder(zhPath) + + inputs := []string{ + `来自小王的(扶高升学咨询群)群聊消息,通过id99使用qq_get_message工具获取消息正文。获取内容后使用 output_send(channel="qq") 回复该群聊,content 设为 JSON 字符串:{"content":"你的回复","group_id":979911915}`, + `【重要!老大消息】来自—/的私聊消息,通过id54使用qq_get_message工具获取消息正文。获取内容后使用 output_send(channel="qq") 回复对方,content 设为 JSON 字符串:{"content":"你的回复","user_id":2198972886}`, + `The nginx server is configured with reverse proxy. 帮我查一下 Docker 容器状态。`, + `老大你好,React 组件已经封装好了,Nginx 配置也改完了,Docker 部署没问题。`, + } + + for i, inp := range inputs { + rawVec := e.Vectorize(inp) + cleanVec := e.VectorizeClean(inp) + sim := cosineSim(rawVec, cleanVec) + rawTokens := len(e.tokenize(inp)) + cleanTokens := len(e.tokenize(CleanTemplateText(inp))) + t.Logf("[%d] sim(raw,clean)=%.4f tokens: raw=%d clean=%d", i, sim, rawTokens, cleanTokens) + } +} + +// --- bilingual test data --- + +func genBilingualEvents() []bilingualEvent { + entries := []struct { + topic string + zh string // Chinese description + en string // English terms mixed in + source string + }{ + {"大学招生", "河南医药大学录取分数线", "", "qq"}, + {"大学招生", "医学院专业排名", "medical university ranking", "agent"}, + {"大学招生", "高考志愿填报咨询", "college application consultation", "qq"}, + {"大学招生", "河南医药大学 Docker 部署项目", "docker deployment project", "agent"}, + {"老大私聊", "老大私聊消息回复", "boss private chat reply", "qq"}, + {"老大私聊", "老大说了关于 React 组件的事情", "boss talked about React components", "agent"}, + {"老大私聊", "回复老大关于服务器配置问题", "reply boss about nginx config", "agent"}, + {"老大私聊", "老大要求检查 Docker 容器状态", "boss asked to check docker status", "qq"}, + {"前端开发", "前端组件封装", "React component encapsulation", "cli"}, + {"前端开发", "页面路由配置 layout 设计", "page route config layout design", "cli"}, + {"前端开发", "JavaScript 交互逻辑开发", "javascript interaction logic", "agent"}, + {"前端开发", "TypeScript 代码调试优化", "typescript code debug optimization", "agent"}, + {"服务器运维", "Nginx 反向代理配置", "nginx reverse proxy config", "cli"}, + {"服务器运维", "Docker 容器部署方案", "docker container deployment", "cli"}, + {"服务器运维", "数据库备份恢复", "database backup recovery", "agent"}, + {"服务器运维", "SSL 证书续期配置", "ssl certificate renewal", "agent"}, + {"股票基金", "基金定投策略配置", "fund investment strategy", "qq"}, + {"股票基金", "股票涨跌分析", "stock market analysis", "agent"}, + {"股票基金", "理财收益 JavaScript 计算", "investment return javascript calculation", "agent"}, + {"股票基金", "市场行情 API 数据获取", "market data api fetch", "qq"}, + } + + var events []bilingualEvent + for i, entry := range entries { + text := entry.zh + if entry.en != "" { + text += " " + entry.en + } + events = append(events, bilingualEvent{ + idx: i, + source: entry.source, + topic: entry.topic, + text: text, + label: fmt.Sprintf("%s (%s)", trimLen(entry.zh, 30), trimLen(entry.en, 30)), + }) + } + return events +} + +func textForBilingual(ev bilingualEvent, modelPaths []string) string { + switch { + case ev.source == "agent" && ev.text != "": + return CleanTemplateText(ev.text) + default: + return CleanTemplateText(ev.text) + } +} diff --git a/internal/memory/clean_stress_test.go b/internal/memory/clean_stress_test.go new file mode 100644 index 0000000..5b85cb9 --- /dev/null +++ b/internal/memory/clean_stress_test.go @@ -0,0 +1,346 @@ +package memory + +import ( + "fmt" + "sort" + "testing" +) + +type cleanTestEvent struct { + idx int + source string + input string + response string + rawText string + cleanedText string + topic string +} + +func TestCleanStressPrecision(t *testing.T) { + events := genStressEvents(200) + topics := []string{"大学招生", "老大私聊", "前端开发", "服务器运维", "股票基金"} + + e := NewStaticEmbedder("/tmp/cc.zh.sample.vec") + if !e.Loaded() { + t.Skip("embedder not loaded") + } + + for _, cleanMode := range []bool{true, false} { + t.Run(fmt.Sprintf("去模版=%v", cleanMode), func(t *testing.T) { + usedTopics := make([]string, 0) + for _, tp := range topics { + if hasTopicEvents(events, tp) { + usedTopics = append(usedTopics, tp) + } + } + if len(usedTopics) == 0 { + t.Fatal("no events for any topic") + } + t.Logf("topics: %v, events: %d", usedTopics, len(events)) + + for _, qTopic := range usedTopics { + query := queryForTopic(qTopic) + qVec := e.Vectorize(query) + + type scored struct { + idx int + topic string + text string + score float64 + } + all := make([]scored, len(events)) + for i, ev := range events { + text := ev.rawText + if cleanMode { + text = ev.cleanedText + } + vec := e.Vectorize(text) + all[i] = scored{idx: i, topic: ev.topic, text: text, score: cosineSim(qVec, vec)} + } + sort.Slice(all, func(i, j int) bool { return all[i].score > all[j].score }) + + topK := len(usedTopics) * 2 + if topK > len(all) { + topK = len(all) + } + + intraHits := 0 + for _, s := range all[:topK] { + if s.topic == qTopic { + intraHits++ + } + } + expected := countTopicEvents(events, qTopic) + if expected > topK { + expected = topK + } + recall := float64(intraHits) / float64(expected) + + if recall < 0.3 { + t.Logf(" [LOW] query=%q topK=%d intra=%d/%d recall=%.2f", qTopic, topK, intraHits, expected, recall) + for _, s := range all[:8] { + t.Logf(" [%.4f] %s", s.score, trimLen(s.text, 60)) + } + } else { + t.Logf(" [OK] query=%q topK=%d intra=%d/%d recall=%.2f", qTopic, topK, intraHits, expected, recall) + } + } + }) + } +} + +func TestCleanStressCrossTopic(t *testing.T) { + events := genStressEvents(200) + + e := NewStaticEmbedder("/tmp/cc.zh.sample.vec") + if !e.Loaded() { + t.Skip("embedder not loaded") + } + + queries := []string{ + "河南医药大学录取分数线", + "老大发了什么私聊消息", + "前端组件怎么封装布局", + "服务器部署配置代理备份证书", + "基金定投收益计算", + } + + for _, q := range queries { + qVec := e.Vectorize(q) + t.Logf("query: %q", q) + + type scored struct { + idx int + topic string + score float64 + } + all := make([]scored, len(events)) + for i, ev := range events { + all[i] = scored{idx: i, topic: ev.topic, score: cosineSim(qVec, e.Vectorize(ev.cleanedText))} + } + sort.Slice(all, func(i, j int) bool { return all[i].score > all[j].score }) + + topScores := make(map[string]float64) + for _, s := range all[:10] { + if _, ok := topScores[s.topic]; !ok { + topScores[s.topic] = s.score + } + } + for tp, sc := range topScores { + t.Logf(" [%.4f] %s", sc, tp) + } + } +} + +func TestCleanTemplateNoiseSuppression(t *testing.T) { + e := NewStaticEmbedder("/tmp/cc.zh.sample.vec") + if !e.Loaded() { + t.Skip("embedder not loaded") + } + + noisyInput := `来自小王的(扶高升学咨询群)群聊消息,通过id99使用qq_get_message工具获取消息正文。获取内容后使用 output_send(channel="qq") 回复该群聊,content 设为 JSON 字符串:{"content":"你的回复","group_id":979911915}` + cleanInput := `来自小王的(扶高升学咨询群)群聊消息` + + query := "扶高升学咨询群" + qClear := e.Vectorize(query) + qNoisy := e.Vectorize(noisyInput) + qClean := e.Vectorize(cleanInput) + + n2c := cosineSim(qNoisy, qClean) + n2q := cosineSim(qNoisy, qClear) + c2q := cosineSim(qClean, qClear) + + t.Logf("noisy(%q) vs clean(%q) = %.4f", noisyInput[:30], cleanInput, n2c) + t.Logf("noisy vs query(%q) = %.4f", query, n2q) + t.Logf("clean vs query = %.4f", c2q) + + if c2q <= n2q { + t.Log("NOTE: clean not better than noisy for this pattern (may have useful info in metadata)") + } +} + +func TestCleanVectorConsistency(t *testing.T) { + e := NewStaticEmbedder("/tmp/cc.zh.sample.vec") + if !e.Loaded() { + t.Skip("embedder not loaded") + } + + templates := []string{ + `通过id1使用qq_get_message工具获取消息正文。获取内容后使用 output_send(channel="qq") 回复该群聊,content 设为 JSON 字符串:{"content":"你的回复","group_id":1}`, + `通过id2使用qq_get_message工具获取消息正文。获取内容后使用 output_send(channel="qq") 回复对方,content 设为 JSON 字符串:{"content":"你的回复","user_id":2}`, + `通过id3使用qq_get_message工具获取消息正文。获取后必须使用qq_send_private_msg工具回复对方,不得使用其他非回复工具。`, + `[12:00] agent: 处理完成`, + } + + for _, tmpl := range templates { + cleaned := CleanTemplateText(tmpl) + t.Logf("template {%q} → {%q} (%d chars)", trimLen(tmpl, 60), cleaned, len(cleaned)) + } + + pairs := []struct { + a, b string + reason string + }{ + {cleanQQGroup("A", "群1"), cleanQQGroup("B", "群1"), "同群不同人"}, + {cleanQQGroup("A", "群1"), cleanQQGroup("A", "群2"), "同人不同群"}, + {cleanQQPrivate("老大"), cleanQQPrivate("老板"), "私聊不同人"}, + {cleanQQGroup("A", "高考群"), cleanQQPrivate("老大"), "群聊vs私聊"}, + } + + for _, p := range pairs { + va := e.Vectorize(p.a) + vb := e.Vectorize(p.b) + sim := cosineSim(va, vb) + t.Logf("sim(%q, %q) [%s] = %.4f", trimLen(p.a, 40), trimLen(p.b, 40), p.reason, sim) + } +} + +func BenchmarkCleanVectorize(b *testing.B) { + e := NewStaticEmbedder("/tmp/cc.zh.sample.vec") + if !e.Loaded() { + b.Skip("embedder not loaded") + } + + texts := make([]string, 100) + for i := range texts { + texts[i] = fmt.Sprintf( + `【重要!老大消息】来自—/的私聊消息,通过id%d使用qq_get_message工具获取消息正文。获取内容后使用 output_send(channel="qq") 回复对方,content 设为 JSON 字符串:{"content":"你的回复","user_id":%d}`, + i, 1000+i, + ) + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + e.VectorizeClean(texts[i%len(texts)]) + } +} + +// --- test data generators --- + +func genStressEvents(n int) []cleanTestEvent { + if n <= 0 { + return nil + } + topics := []string{"大学招生", "老大私聊", "前端开发", "服务器运维", "股票基金"} + events := make([]cleanTestEvent, 0, n) + + names := []string{"小明", "小红", "小张", "老王", "老大", "小李", "小王", "小赵"} + groups := []string{"扶高升学咨询群", "前端技术交流", "服务器运维群", "基金定投群", "闲聊群"} + + topicContent := map[string]struct { + keywords []string + sources []string + }{ + "大学招生": {[]string{"河南医药大学", "录取分数线", "专业排名", "高考志愿", "招生简章"}, []string{"qq", "qq", "agent"}}, + "老大私聊": {[]string{"老大私聊消息", "回复老大", "任务安排", "汇报工作", "收到"}, []string{"qq", "agent", "agent"}}, + "前端开发": {[]string{"前端组件封装", "页面路由配置", "界面布局设计", "交互逻辑开发", "代码调试优化"}, []string{"cli", "cli", "agent"}}, + "服务器运维": {[]string{"反向代理配置", "容器部署方案", "证书续期", "数据库备份恢复", "监控告警处理"}, []string{"cli", "agent", "agent"}}, + "股票基金": {[]string{"基金定投策略", "股票涨跌分析", "理财收益计算", "市场行情分析", "投资风险管理"}, []string{"qq", "qq", "agent"}}, + } + + for i := 0; i < n; i++ { + tp := topics[i%len(topics)] + info := topicContent[tp] + kw := info.keywords[i%len(info.keywords)] + nm := names[i%len(names)] + grp := groups[i%len(groups)] + src := info.sources[i%len(info.sources)] + + var input, response string + switch src { + case "qq": + if tp == "老大私聊" { + input = fmt.Sprintf(`【重要!老大消息】来自%s的私聊消息,通过id%d使用qq_get_message工具获取消息正文。获取内容后使用 output_send(channel="qq") 回复对方,content 设为 JSON 字符串:{"content":"你的回复","user_id":%d}`, nm, i, 1000+i) + if i%3 == 0 { + response = fmt.Sprintf("已回复老大,关于%s", kw) + } + } else { + input = fmt.Sprintf(`来自%s的(%s)群聊消息,通过id%d使用qq_get_message工具获取消息正文。获取内容后使用 output_send(channel="qq") 回复该群聊,content 设为 JSON 字符串:{"content":"你的回复","group_id":%d}`, nm, grp, i, 9000+i) + if i%3 == 0 { + response = fmt.Sprintf("已回复%s相关的问题", kw) + } + } + case "agent": + input = fmt.Sprintf(`来自%s的(%s)消息`, nm, grp) + response = fmt.Sprintf("关于%s,我的建议是...已处理完成。", kw) + case "cli": + input = fmt.Sprintf("查询%s的相关信息", kw) + response = fmt.Sprintf("查到了%s的结果", kw) + } + + cleaned := cleanEventText(src, input, response) + raw := rawEventText(src, input, response) + events = append(events, cleanTestEvent{ + idx: i, + source: src, + input: input, + response: response, + rawText: raw, + cleanedText: cleaned, + topic: tp, + }) + } + return events +} + +func cleanEventText(source, input, response string) string { + switch { + case source == "agent" && response != "": + return CleanTemplateText(response) + case source == "cold_storage": + return CleanTemplateText(input + " " + response) + default: + return CleanTemplateText(input) + } +} + +func rawEventText(source, input, response string) string { + if response == "" { + return input + } + return input + " " + response +} + +func queryForTopic(topic string) string { + switch topic { + case "大学招生": + return "河南医药大学录取分数线多少" + case "老大私聊": + return "老大刚才说了什么私聊消息" + case "前端开发": + return "前端组件怎么封装布局" + case "服务器运维": + return "服务器部署容器代理配置证书备份监控告警" + case "股票基金": + return "基金定投收益怎么样" + default: + return topic + } +} + +func hasTopicEvents(events []cleanTestEvent, topic string) bool { + for _, ev := range events { + if ev.topic == topic { + return true + } + } + return false +} + +func countTopicEvents(events []cleanTestEvent, topic string) int { + n := 0 + for _, ev := range events { + if ev.topic == topic { + n++ + } + } + return n +} + +func cleanQQGroup(user, group string) string { + return fmt.Sprintf("来自%s的(%s)群聊消息", user, group) +} + +func cleanQQPrivate(user string) string { + return fmt.Sprintf("【重要!老大消息】来自%s的私聊消息", user) +} diff --git a/internal/memory/clean_text.go b/internal/memory/clean_text.go new file mode 100644 index 0000000..659216f --- /dev/null +++ b/internal/memory/clean_text.go @@ -0,0 +1,57 @@ +package memory + +import ( + "regexp" + "strings" + + "gitcode.com/JianFeeeee/HomeAgent/internal/memory/vector" +) + +var ( + reQQGroupSuffix = regexp.MustCompile( + `,通过id\d+使用qq_get_message工具获取消息正文。获取内容后使用 output_send\(channel="qq"\) 回复该群聊,content 设为 JSON 字符串:\{[^}]*\}`, + ) + reQQPrivateSuffix = regexp.MustCompile( + `,通过id\d+使用qq_get_message工具获取消息正文。获取内容后使用 output_send\(channel="qq"\) 回复对方,content 设为 JSON 字符串:\{[^}]*\}`, + ) + reQQOldReply = regexp.MustCompile( + `通过id\d+使用qq_get_message工具获取消息正文。获取后必须使用[^。]+。`, + ) + reQQOldForbid = regexp.MustCompile( + `你只能通过qq_get_message先看消息,然后直接用%!s\(MISSING\)send_private_msg回复,中间的思考过程禁止调用任何其他工具\s*→\s*`, + ) + reQQGeneral = regexp.MustCompile( + `通过id\d+使用qq_get_message工具获取消息正文[。,][^。]*?(?:回复|发送消息)`, + ) + reTimestamp = regexp.MustCompile( + `\[\d{2}:\d{2}\]\s*`, + ) + reAgentPrefix = regexp.MustCompile( + `冷知识|注意|提示|核心要求|规则`, + ) + reMultiSpace = regexp.MustCompile(`\s+`) +) + +func CleanTemplateText(text string) string { + text = reQQGroupSuffix.ReplaceAllString(text, "") + text = reQQPrivateSuffix.ReplaceAllString(text, "") + text = reQQOldReply.ReplaceAllString(text, "") + text = reQQOldForbid.ReplaceAllString(text, "") + text = reQQGeneral.ReplaceAllString(text, "") + text = reTimestamp.ReplaceAllString(text, "") + text = reMultiSpace.ReplaceAllString(text, " ") + text = strings.TrimSpace(text) + + if text == "" { + return "" + } + + text = strings.TrimPrefix(text, ",") + text = strings.TrimPrefix(text, ",") + text = strings.TrimSpace(text) + return text +} + +func (e *StaticEmbedder) VectorizeClean(text string) vector.Vector { + return e.Vectorize(CleanTemplateText(text)) +} diff --git a/internal/memory/cut.go b/internal/memory/cut.go index a1f5918..b29cb49 100644 --- a/internal/memory/cut.go +++ b/internal/memory/cut.go @@ -93,6 +93,7 @@ var stopWords = map[string]bool{ } func ExtractKeywords(text string) []string { + text = CleanTemplateText(text) x := GetJieba() if x == nil { return nil diff --git a/internal/memory/indexer.go b/internal/memory/indexer.go index bfd5992..fdcdf3e 100644 --- a/internal/memory/indexer.go +++ b/internal/memory/indexer.go @@ -90,11 +90,13 @@ func (idx *Indexer) BuildContext(userInput string) *InjectedContext { return &InjectedContext{Summary: ""} } + input := CleanTemplateText(userInput) + // 1. 向量搜索:从实体名向量索引中找到相关实体 - vectorEntities := idx.vectorSearchEntities(userInput) + vectorEntities := idx.vectorSearchEntities(input) // 2. 关键词搜索:已有逻辑 - keywords := ExtractKeywords(userInput) + keywords := ExtractKeywords(input) if len(keywords) == 0 && len(vectorEntities) == 0 { keywords = []string{userInput} } diff --git a/internal/memory/real_context_test.go b/internal/memory/real_context_test.go new file mode 100644 index 0000000..7833ea8 --- /dev/null +++ b/internal/memory/real_context_test.go @@ -0,0 +1,309 @@ +package memory + +import ( + "encoding/json" + "os" + "sort" + "strings" + "testing" + "time" +) + +type realEvent struct { + Timestamp time.Time `json:"timestamp"` + Source string `json:"source"` + Input string `json:"input"` + Response string `json:"response"` +} + +func TestCleanTemplateText(t *testing.T) { + cases := []struct { + input string + expected string + contains string + }{ + { + input: `来自A的(扶高升学咨询群)群聊消息,通过id36使用qq_get_message工具获取消息正文。获取内容后使用 output_send(channel="qq") 回复该群聊,content 设为 JSON 字符串:{"content":"你的回复","group_id":979911915}`, + expected: "来自A的(扶高升学咨询群)群聊消息", + }, + { + input: `【重要!老大消息】来自—/的私聊消息,通过id54使用qq_get_message工具获取消息正文。获取内容后使用 output_send(channel="qq") 回复对方,content 设为 JSON 字符串:{"content":"你的回复","user_id":2198972886}`, + expected: "【重要!老大消息】来自—/的私聊消息", + }, + { + input: `[12:05] agent: 已经回复老大啦~继续去搞 vanblog 换端口的事 😊`, + contains: "已经回复老大啦", + }, + { + input: `加载文档记忆: 扶高升学咨询群 河南医药大学 回复 招生 专业 录取`, + expected: "加载文档记忆: 扶高升学咨询群 河南医药大学 回复 招生 专业 录取", + }, + { + input: `通过id10使用qq_get_message工具获取消息正文。获取后必须使用qq_send_private_msg工具回复对方,不得使用其他非回复工具。你只能通过qq_get_message先看消息,然后直接用%!s(MISSING)send_private_msg回复,中间的思考过程禁止调用任何其他工具 → 已经回复老大啦~`, + contains: "已经回复老大啦", + }, + { + input: "", + expected: "", + }, + { + input: `处理错误: all 9 providers failed, last error: lua transform_request: adapter tesy not loaded`, + contains: "处理错误", + }, + } + + for i, c := range cases { + got := CleanTemplateText(c.input) + if c.expected != "" && got != c.expected { + t.Errorf("case %d:\n input: %q\n expected: %q\n got: %q", i, trimLen(c.input, 60), c.expected, got) + } + if c.contains != "" && !strings.Contains(got, c.contains) { + t.Errorf("case %d: expected to contain %q, got %q", i, c.contains, got) + } + t.Logf("case %d: %q → %q", i, trimLen(c.input, 60), got) + } +} + +func TestRealContextPerSourceVector(t *testing.T) { + modelPath := "/tmp/cc.zh.sample.vec" + if _, err := os.Stat(modelPath); os.IsNotExist(err) { + t.Skip("real embedding file not found") + } + e := NewStaticEmbedder(modelPath) + + data, err := os.ReadFile("/tmp/context.json") + if err != nil { + t.Fatal(err) + } + var raw []realEvent + if err := json.Unmarshal(data, &raw); err != nil { + t.Fatal(err) + } + t.Logf("loaded %d real events", len(raw)) + + type scored struct { + idx int + source string + text string + score float64 + } + + clean := func(ev realEvent) string { + switch { + case ev.Source == "agent" && ev.Response != "": + return CleanTemplateText(ev.Response) + case ev.Source == "cold_storage": + return CleanTemplateText(ev.Input + " " + ev.Response) + default: + return CleanTemplateText(ev.Input) + } + } + + t.Run("医药大学_不同源向量", func(t *testing.T) { + q := "河南医药大学招生分数录取排名" + qVec := e.VectorizeClean(q) + + all := make([]scored, len(raw)) + for i, ev := range raw { + text := clean(ev) + all[i] = scored{idx: i, source: ev.Source, text: text[:min(len(text), 200)], score: cosineSim(qVec, e.Vectorize(text))} + } + sort.Slice(all, func(i, j int) bool { return all[i].score > all[j].score }) + + t.Log("top 5:") + for _, s := range all[:5] { + t.Logf(" [%.4f] [%-13s] %s", s.score, s.source, trimLen(s.text, 80)) + } + + var univHigh bool + for _, s := range all[:8] { + if strings.Contains(s.text, "医药大学") || strings.Contains(s.text, "升学") { + univHigh = true + break + } + } + if !univHigh { + t.Error("expected university-related events in top 8") + } + }) + + t.Run("老大私聊_agent主用Response", func(t *testing.T) { + q := "老大私聊说了什么" + qVec := e.VectorizeClean(q) + + all := make([]scored, len(raw)) + for i, ev := range raw { + text := clean(ev) + all[i] = scored{idx: i, source: ev.Source, text: text[:min(len(text), 200)], score: cosineSim(qVec, e.Vectorize(text))} + } + sort.Slice(all, func(i, j int) bool { return all[i].score > all[j].score }) + + t.Log("top 5:") + for _, s := range all[:5] { + t.Logf(" [%.4f] [%-13s] %s", s.score, s.source, trimLen(s.text, 80)) + } + + var bossFound bool + for _, s := range all[:8] { + if strings.Contains(s.text, "老大") { + bossFound = true + break + } + } + if !bossFound { + t.Error("expected events mentioning 老大 in top 8") + } + + var agentFound bool + for _, s := range all[:5] { + if s.source == "agent" { + agentFound = true + break + } + } + t.Logf("agent in top5: %v (source strategy: agent events use Response for vector)", agentFound) + }) + + t.Run("图片转SVG_去模版后效果", func(t *testing.T) { + q := "图片转换SVG工具" + qVec := e.VectorizeClean(q) + + all := make([]scored, len(raw)) + for i, ev := range raw { + text := clean(ev) + all[i] = scored{idx: i, source: ev.Source, text: text[:min(len(text), 200)], score: cosineSim(qVec, e.Vectorize(text))} + } + sort.Slice(all, func(i, j int) bool { return all[i].score > all[j].score }) + + t.Log("top 5:") + for _, s := range all[:5] { + t.Logf(" [%.4f] [%-13s] %s", s.score, s.source, trimLen(s.text, 80)) + } + + var img bool + for _, s := range all[:5] { + if strings.Contains(s.text, "图片") || strings.Contains(s.text, "SVG") { + img = true + break + } + } + if !img { + t.Error("expected image-related events in top 5") + } + }) + + t.Run("南航航空航天", func(t *testing.T) { + q := "南航航空航天专业转电气" + qVec := e.VectorizeClean(q) + + all := make([]scored, len(raw)) + for i, ev := range raw { + text := clean(ev) + all[i] = scored{idx: i, source: ev.Source, text: text[:min(len(text), 200)], score: cosineSim(qVec, e.Vectorize(text))} + } + sort.Slice(all, func(i, j int) bool { return all[i].score > all[j].score }) + + t.Log("top 5:") + for _, s := range all[:5] { + t.Logf(" [%.4f] [%-13s] %s", s.score, s.source, trimLen(s.text, 80)) + } + + var nau bool + for _, s := range all[:5] { + if strings.Contains(s.text, "南航") { + nau = true + break + } + } + if !nau { + t.Error("expected 南航 in top 5") + } + }) + + t.Run("跨域区分度", func(t *testing.T) { + pairs := []struct { + a, b string + }{ + {"老大私聊说了什么", "河南医药大学招生分数"}, + {"老大私聊说了什么", "图片转换SVG工具"}, + {"南航航空航天电气", "河南医药大学录取"}, + {"图片转换SVG工具", "老大私聊"}, + } + for _, p := range pairs { + va := e.VectorizeClean(p.a) + vb := e.VectorizeClean(p.b) + s := cosineSim(va, vb) + t.Logf(" sim(%q, %q) = %.4f", p.a, p.b, s) + } + + univVec := e.VectorizeClean("河南医药大学招生") + bossVec := e.VectorizeClean("老大私聊说了什么") + t.Logf("cross-domain sim(医药大学, 老大私聊) = %.4f", cosineSim(univVec, bossVec)) + }) + + t.Run("去模版节省量", func(t *testing.T) { + var savedTotal int + for i, ev := range raw { + orig := len(ev.Input + " " + ev.Response) + after := len(clean(ev)) + saved := orig - after + savedTotal += saved + if saved > 100 { + t.Logf(" [%2d] [%-13s] 节省 %d 字符 (raw=%d clean=%d)", i, ev.Source, saved, orig, after) + } + } + t.Logf("总计节省 %d 字符", savedTotal) + }) +} + +func TestRealContextEmbedderStats(t *testing.T) { + e := NewStaticEmbedder("/tmp/cc.zh.sample.vec") + if !e.Loaded() { + t.Skip("embedder not loaded") + } + + data, err := os.ReadFile("/tmp/context.json") + if err != nil { + t.Fatal(err) + } + var raw []realEvent + if err := json.Unmarshal(data, &raw); err != nil { + t.Fatal(err) + } + + for _, ev := range raw[:5] { + var text string + switch { + case ev.Source == "agent" && ev.Response != "": + text = CleanTemplateText(ev.Response) + case ev.Source == "cold_storage": + text = CleanTemplateText(ev.Input + " " + ev.Response) + default: + text = CleanTemplateText(ev.Input) + } + vec := e.Vectorize(text) + origLen := len(ev.Input + ev.Response) + t.Logf("[%-13s] raw=%d cleaned=%d dims=%d", ev.Source, origLen, len(text), len(vec)) + for k := range vec { + if !isNumericKey(k) { + t.Errorf("non-numeric key %q — should be dense space", k) + } + } + } +} + +func containsAny(s string, subs []string) bool { + for _, sub := range subs { + if sub != "" && strings.Contains(s, sub) { + return true + } + } + return false +} + +func trimLen(s string, n int) string { + if len(s) > n { + return s[:n] + } + return s +} diff --git a/internal/memory/static_embedder.go b/internal/memory/static_embedder.go new file mode 100644 index 0000000..860588e --- /dev/null +++ b/internal/memory/static_embedder.go @@ -0,0 +1,369 @@ +package memory + +import ( + "bufio" + "compress/gzip" + "fmt" + "log" + "math" + "net/http" + "os" + "path/filepath" + "strconv" + "strings" + "sync" + "unicode/utf8" + + "github.com/yanyiwu/gojieba" + "gitcode.com/JianFeeeee/HomeAgent/internal/memory/vector" +) + +const downloadMaxWords = 200000 + +var knownModelURLs = []struct { + sub string + url string +}{ + {"numberbatch", "https://conceptnet.s3.amazonaws.com/downloads/2019/numberbatch/numberbatch-19.08.txt.gz"}, + {"cc.zh.", "https://dl.fbaipublicfiles.com/fasttext/vectors-crawl/cc.zh.300.vec.gz"}, + {"cc.en.", "https://dl.fbaipublicfiles.com/fasttext/vectors-crawl/cc.en.300.vec.gz"}, +} + +type StaticEmbedder struct { + mu sync.RWMutex + jieba *gojieba.Jieba + stopWords map[string]bool + + words map[string][]float64 + dim int + loaded bool + + unkVec []float64 + unkNorm float64 +} + +func modelDownloadURL(modelPath string) string { + for _, m := range knownModelURLs { + if strings.Contains(modelPath, m.sub) { + return m.url + } + } + return knownModelURLs[0].url +} + +func downloadFastTextModel(targetPath, url string) error { + tmpPath := targetPath + ".download.tmp" + if err := os.MkdirAll(filepath.Dir(targetPath), 0755); err != nil { + return fmt.Errorf("mkdir: %w", err) + } + + f, err := os.Create(tmpPath) + if err != nil { + return fmt.Errorf("create tmp: %w", err) + } + defer f.Close() + + resp, err := http.Get(url) + if err != nil { + os.Remove(tmpPath) + return fmt.Errorf("http get %s: %w", url, err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + os.Remove(tmpPath) + return fmt.Errorf("http status %s", resp.Status) + } + + gz, err := gzip.NewReader(resp.Body) + if err != nil { + os.Remove(tmpPath) + return fmt.Errorf("gzip: %w", err) + } + defer gz.Close() + + scanner := bufio.NewScanner(gz) + buf := make([]byte, 4*1024*1024) + scanner.Buffer(buf, len(buf)) + + writer := bufio.NewWriter(f) + + if !scanner.Scan() { + os.Remove(tmpPath) + return fmt.Errorf("empty gzip content") + } + parts := strings.Fields(scanner.Text()) + if len(parts) >= 2 { + fmt.Fprintf(writer, "%d %s\n", downloadMaxWords, parts[1]) + } else { + fmt.Fprintln(writer, scanner.Text()) + } + + var lineCount int + for scanner.Scan() && lineCount < downloadMaxWords { + line := scanner.Text() + if line == "" { + continue + } + fmt.Fprintln(writer, line) + lineCount++ + + if lineCount%50000 == 0 { + log.Printf("[static_embedder] download progress: %d/%d words", lineCount, downloadMaxWords) + } + } + + writer.Flush() + f.Close() + + if err := os.Rename(tmpPath, targetPath); err != nil { + os.Remove(tmpPath) + return fmt.Errorf("rename: %w", err) + } + + log.Printf("[static_embedder] download complete: %d words to %s", lineCount, targetPath) + return nil +} + +func ensureModelFile(modelPath string) { + if modelPath == "" { + return + } + if _, err := os.Stat(modelPath); err == nil { + return + } + url := modelDownloadURL(modelPath) + log.Printf("[static_embedder] model %s not found, downloading from fastText...", modelPath) + if dlErr := downloadFastTextModel(modelPath, url); dlErr != nil { + log.Printf("[static_embedder] download failed: %v, will use TF-IDF fallback", dlErr) + } else { + log.Printf("[static_embedder] download ok") + } +} + +func NewStaticEmbedder(modelPaths ...string) *StaticEmbedder { + sw := make(map[string]bool) + for k, v := range stopWords { + sw[k] = v + } + e := &StaticEmbedder{ + jieba: GetJieba(), + stopWords: sw, + words: make(map[string][]float64), + } + + if len(modelPaths) == 0 { + log.Printf("[static_embedder] no model path configured, using TF-IDF fallback") + return e + } + + for _, p := range modelPaths { + ensureModelFile(p) + } + if err := e.loadAll(modelPaths); err != nil { + log.Printf("[static_embedder] load failed: %v, using TF-IDF fallback", err) + } + return e +} + +func (e *StaticEmbedder) loadAll(paths []string) error { + var firstErr error + for i, p := range paths { + if p == "" { + continue + } + primary := i == 0 + if err := e.load(p, primary); err != nil { + log.Printf("[static_embedder] load %s: %v", p, err) + if firstErr == nil { + firstErr = err + } + } + } + return firstErr +} + +func (e *StaticEmbedder) load(path string, primary bool) error { + f, err := os.Open(path) + if err != nil { + return fmt.Errorf("open: %w", err) + } + defer f.Close() + + scanner := bufio.NewScanner(f) + buf := make([]byte, 1024*1024) + scanner.Buffer(buf, len(buf)) + + if !scanner.Scan() { + return fmt.Errorf("empty file") + } + header := strings.TrimSpace(scanner.Text()) + parts := strings.Fields(header) + if len(parts) < 2 { + return fmt.Errorf("invalid header: %s", header) + } + dim, err := strconv.Atoi(parts[1]) + if err != nil || dim <= 0 { + return fmt.Errorf("invalid dimension: %s", parts[1]) + } + + if primary { + e.dim = dim + } + + var vecSum []float64 + var count int + if primary { + vecSum = make([]float64, dim) + } + + for scanner.Scan() { + line := strings.TrimSpace(scanner.Text()) + if line == "" { + continue + } + fields := strings.Fields(line) + if len(fields) < dim+1 { + continue + } + word := fields[0] + + if _, exists := e.words[word]; exists { + continue + } + + vec := make([]float64, dim) + for i := 0; i < dim; i++ { + v, _ := strconv.ParseFloat(fields[i+1], 64) + vec[i] = v + } + e.words[word] = vec + if primary { + for i := range vecSum { + vecSum[i] += vec[i] + } + count++ + } + } + + if primary { + if count == 0 { + return fmt.Errorf("no word vectors found in primary model") + } + for i := range vecSum { + vecSum[i] /= float64(count) + } + e.unkVec = make([]float64, dim) + copy(e.unkVec, vecSum) + var normSq float64 + for _, v := range e.unkVec { + normSq += v * v + } + e.unkNorm = float64(math.Sqrt(normSq)) + e.loaded = true + } + + log.Printf("[static_embedder] loaded %d words, dim=%d from %s", len(e.words), e.dim, path) + return nil +} + +func (e *StaticEmbedder) tokenize(text string) []string { + if e.jieba == nil { + return nil + } + words := e.jieba.Cut(text, true) + var result []string + seen := make(map[string]bool) + for _, w := range words { + w = strings.TrimSpace(w) + if w == "" || e.stopWords[w] || seen[w] { + continue + } + if utf8.RuneCountInString(w) < 2 { + continue + } + seen[w] = true + result = append(result, w) + } + return result +} + +func (e *StaticEmbedder) Vectorize(text string) vector.Vector { + e.mu.RLock() + loaded := e.loaded + dim := e.dim + unkVec := e.unkVec + e.mu.RUnlock() + + tokens := e.tokenize(text) + if len(tokens) == 0 { + return vector.Vector{} + } + + tf := make(map[string]float64) + for _, t := range tokens { + tf[t]++ + } + maxTF := 0.0 + for _, c := range tf { + if c > maxTF { + maxTF = c + } + } + + if !loaded { + vec := make(vector.Vector) + for word, count := range tf { + vec[word] = count / maxTF + } + return vec + } + + sum := make([]float64, dim) + var weightSum float64 + + for word, count := range tf { + e.mu.RLock() + vec, ok := e.words[word] + e.mu.RUnlock() + + w := count / maxTF + + if !ok { + for i, v := range unkVec { + sum[i] += w * v + } + } else { + for i, v := range vec { + sum[i] += w * v + } + } + weightSum += w + } + + if weightSum > 0 { + for i := range sum { + sum[i] /= weightSum + } + } + + vec := make(vector.Vector, dim) + for i, v := range sum { + if v != 0 { + vec[strconv.Itoa(i)] = v + } + } + return vec +} + +func (e *StaticEmbedder) Dim() int { + e.mu.RLock() + defer e.mu.RUnlock() + return e.dim +} + +func (e *StaticEmbedder) Loaded() bool { + e.mu.RLock() + defer e.mu.RUnlock() + return e.loaded +} diff --git a/internal/memory/static_embedder_test.go b/internal/memory/static_embedder_test.go new file mode 100644 index 0000000..8e13b96 --- /dev/null +++ b/internal/memory/static_embedder_test.go @@ -0,0 +1,236 @@ +package memory + +import ( + "math" + "sort" + "testing" + "time" + + "gitcode.com/JianFeeeee/HomeAgent/internal/memory/vector" +) + +type testContextEvent struct { + Timestamp time.Time + Input string + Response string + Vector vector.Vector +} + +func cosineSim(a, b vector.Vector) float64 { + var dot, normA, normB float64 + for f, va := range a { + dot += va * b[f] + normA += va * va + } + for _, vb := range b { + normB += vb * vb + } + if normA == 0 || normB == 0 { + return 0 + } + return dot / (math.Sqrt(normA) * math.Sqrt(normB)) +} + +func TestStaticEmbedderLoad(t *testing.T) { + e := NewStaticEmbedder("/tmp/cc.zh.sample.vec") + if !e.Loaded() { + t.Fatal("embedder should be loaded") + } + if e.Dim() != 300 { + t.Errorf("expected dim=300, got %d", e.Dim()) + } +} + +func TestStaticEmbedderConsistency(t *testing.T) { + e := NewStaticEmbedder("/tmp/cc.zh.sample.vec") + + v1 := e.Vectorize("今天天气怎么样") + v2 := e.Vectorize("今天天气怎么样") + + if len(v1) != len(v2) { + t.Errorf("same input should produce same dimension count, got %d vs %d", len(v1), len(v2)) + } + sim := cosineSim(v1, v2) + if math.Abs(sim-1.0) > 0.0001 { + t.Errorf("same input should have cosine similarity ~1.0, got %.6f", sim) + } +} + +func isNumericKey(s string) bool { + if s == "" { + return false + } + for _, c := range s { + if c < '0' || c > '9' { + return false + } + } + return true +} + +func TestStaticEmbedderFallback(t *testing.T) { + e := NewStaticEmbedder("") + if e.Loaded() { + t.Fatal("empty path embedder should not be loaded") + } + + v := e.Vectorize("测试文本") + if len(v) == 0 { + t.Fatal("fallback vector should not be empty") + } +} + +func TestStaticEmbedderAllInDenseSpace(t *testing.T) { + e := NewStaticEmbedder("/tmp/cc.zh.sample.vec") + + texts := []string{ + "今天天气怎么样", + "股票基金投资", + "微积分导数数学题", + "台风天注意安全", + "基金定投", + "数学作业", + "明天会不会下雨", + } + + for _, text := range texts { + v := e.Vectorize(text) + for k := range v { + if !isNumericKey(k) { + t.Errorf("%q produced non-numeric key %q — should be in dense space", text, k) + } + } + } + t.Log("all texts produce numeric keys — same dense space") +} + +func TestStaticEmbedderSemanticSimilarity(t *testing.T) { + e := NewStaticEmbedder("/tmp/cc.zh.sample.vec") + + pairs := []struct { + a, b string + related bool + }{ + {"今天天气怎么样", "明天会不会下雨", true}, + {"今天天气怎么样", "股票基金投资", false}, + {"股票基金投资", "基金定投", true}, + {"股票基金投资", "微积分导数数学题", false}, + {"微积分导数数学题", "数学作业", true}, + } + + for _, p := range pairs { + va := e.Vectorize(p.a) + vb := e.Vectorize(p.b) + sim := cosineSim(va, vb) + t.Logf("sim(%q, %q) = %.4f (related=%v)", p.a, p.b, sim, p.related) + } + + weatherSim := cosineSim(e.Vectorize("今天天气怎么样"), e.Vectorize("明天会不会下雨")) + stockSim := cosineSim(e.Vectorize("今天天气怎么样"), e.Vectorize("股票基金投资")) + t.Logf("[verify] weather-weather=%.4f, weather-stock=%.4f", weatherSim, stockSim) + if weatherSim <= stockSim { + t.Errorf("weather-weather(%.4f) should be > weather-stock(%.4f)", weatherSim, stockSim) + } +} + +func TestContextPruneWithRealEmbedding(t *testing.T) { + e := NewStaticEmbedder("/tmp/cc.zh.sample.vec") + + type event struct { + input string + response string + } + allEvents := []event{ + {"今天天气怎么样", "挺好的"}, + {"明天会不会下雨", "可能不会"}, + {"台风来了", "注意安全"}, + {"帮我算微积分", "好的"}, + {"导数怎么求", "公式如下"}, + {"数学作业", "解答"}, + {"股票涨了", "恭喜"}, + {"基金收益怎么样", "不错"}, + {"最近有什么电影", "推荐"}, + {"晚上吃什么", "随便"}, + {"帮我定个闹钟", "好的"}, + {"查询快递", "已送达"}, + } + + events := make([]testContextEvent, len(allEvents)) + for i, ev := range allEvents { + events[i] = testContextEvent{ + Timestamp: time.Now().Add(time.Duration(i) * time.Second), + Input: ev.input, + Response: ev.response, + Vector: e.Vectorize(ev.input + " " + ev.response), + } + } + + topK := 4 + protectN := 3 + query := "基金股票投资" + queryVec := e.Vectorize(query) + + if len(events) <= topK+protectN { + t.Fatalf("need more events for pruning test") + } + + protectStart := len(events) - protectN + protected := events[protectStart:] + candidates := events[:protectStart] + + type scored struct { + evt testContextEvent + score float64 + } + scoredEvents := make([]scored, len(candidates)) + for i, evt := range candidates { + scoredEvents[i] = scored{evt, cosineSim(queryVec, evt.Vector)} + } + + sort.Slice(scoredEvents, func(i, j int) bool { + return scoredEvents[i].score > scoredEvents[j].score + }) + + keepCount := topK + if keepCount > len(scoredEvents) { + keepCount = len(scoredEvents) + } + keep := scoredEvents[:keepCount] + archived := scoredEvents[keepCount:] + + t.Logf("query: %s", query) + t.Logf("=== retained (topK=%d) ===", topK) + for _, s := range keep { + t.Logf(" [%.4f] %s", s.score, s.evt.Input) + } + t.Logf("=== protected (recent %d) ===", protectN) + for _, e := range protected { + t.Logf(" %s", e.Input) + } + t.Logf("=== archived (%d items) ===", len(archived)) + for _, s := range archived { + t.Logf(" [%.4f] %s", s.score, s.evt.Input) + } + + hasFinance := false + for _, s := range keep { + if s.evt.Input == "股票涨了" || s.evt.Input == "基金收益怎么样" { + hasFinance = true + } + } + if !hasFinance { + t.Error("expected financial events to be retained, but none found") + } + + hasWeather := false + for _, s := range keep { + if s.evt.Input == "今天天气怎么样" || s.evt.Input == "明天会不会下雨" || s.evt.Input == "台风来了" { + hasWeather = true + } + } + if hasWeather { + t.Log("NOTE: weather events are still in retained set — may have overlapping vocabulary") + } + + t.Logf("remaining: %d = topK(%d) + protectN(%d)", topK+protectN, topK, protectN) +}