mirror of
https://gitcode.com/JianFeeeee/HomeAgent.git
synced 2026-09-21 09:28:14 +00:00
feat: multi-language embedding with CleanTemplateText + three-branch vector strategy
- StaticEmbedder: pre-trained ConceptNet Numberbatch/fastText word embeddings, auto-download with TF-IDF fallback, comma-separated multi-model paths - CleanTemplateText: regex stripping of QQ tool call templates and noise - textForVector: per-source vector strategy (agent→Response, user→Input, cold_storage→both) - Indexer.BuildContext and ExtractKeywords now clean input before vectorization - Protect recent 10 events in Prune (regression fix: use local var not const)
This commit is contained in:
243
internal/memory/bilingual_test.go
Normal file
243
internal/memory/bilingual_test.go
Normal file
@ -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)
|
||||
}
|
||||
}
|
||||
346
internal/memory/clean_stress_test.go
Normal file
346
internal/memory/clean_stress_test.go
Normal file
@ -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)
|
||||
}
|
||||
57
internal/memory/clean_text.go
Normal file
57
internal/memory/clean_text.go
Normal file
@ -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))
|
||||
}
|
||||
@ -93,6 +93,7 @@ var stopWords = map[string]bool{
|
||||
}
|
||||
|
||||
func ExtractKeywords(text string) []string {
|
||||
text = CleanTemplateText(text)
|
||||
x := GetJieba()
|
||||
if x == nil {
|
||||
return nil
|
||||
|
||||
@ -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}
|
||||
}
|
||||
|
||||
309
internal/memory/real_context_test.go
Normal file
309
internal/memory/real_context_test.go
Normal file
@ -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
|
||||
}
|
||||
369
internal/memory/static_embedder.go
Normal file
369
internal/memory/static_embedder.go
Normal file
@ -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
|
||||
}
|
||||
236
internal/memory/static_embedder_test.go
Normal file
236
internal/memory/static_embedder_test.go
Normal file
@ -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)
|
||||
}
|
||||
Reference in New Issue
Block a user