From 8ae1d1986925e0cb43957ea1df71feacbf8e13bb Mon Sep 17 00:00:00 2001 From: root Date: Sun, 19 Jul 2026 10:29:21 +0800 Subject: [PATCH] refactor: output channel interface (payload/meta/type) + memory fixes - Redesign output_send__ tools: content JSON string -> structured payload/meta/type params for LLM reliability - executeOutputSendTool: route by type with capability check - executeOutputSendHelp: show meta format + type enum - Updated system prompt rules for new interface - docToTriples: use jieba exact mode adjacent co-occurrence - Unify vector space: Doc.Vector field, ContextToDoc vectorizer, ReindexWithVectorizer on startup --- go.mod | 2 +- go.sum | 2 + internal/agent/core/agent.go | 192 +++++++++++++------- internal/agent/core/agent_functions_test.go | 69 +++---- internal/agent/core/agent_helpers_test.go | 98 ++++++++++ internal/agent/core/context.go | 2 +- internal/memory/cut.go | 23 +++ internal/memory/cut_test.go | 107 +++++++++++ internal/memory/document/document.go | 65 +++++-- internal/memory/document/document_test.go | 2 +- 10 files changed, 453 insertions(+), 109 deletions(-) create mode 100644 internal/memory/cut_test.go diff --git a/go.mod b/go.mod index 7082b75..d4662da 100644 --- a/go.mod +++ b/go.mod @@ -3,7 +3,7 @@ module gitcode.com/JianFeeeee/HomeAgent go 1.25.0 require ( - github.com/mattn/go-sqlite3 v1.14.22 + github.com/mattn/go-sqlite3 v1.14.48 github.com/yuin/gopher-lua v1.1.2 gopkg.in/yaml.v3 v3.0.1 ) diff --git a/go.sum b/go.sum index faf3a7f..4057385 100644 --- a/go.sum +++ b/go.sum @@ -1,5 +1,7 @@ github.com/mattn/go-sqlite3 v1.14.22 h1:2gZY6PC6kBnID23Tichd1K+Z0oS6nE/XwU+Vz/5o4kU= github.com/mattn/go-sqlite3 v1.14.22/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y= +github.com/mattn/go-sqlite3 v1.14.48 h1:7XHIgl0a8HwOaiK4E47ozLkST78rR9+OtNGx27D/TFs= +github.com/mattn/go-sqlite3 v1.14.48/go.mod h1:6JTjA44L93a0QCyJef5YvlPoKXntQPjzWv5gtm9sB6w= github.com/yanyiwu/gojieba v1.4.7 h1:2YkXELcYLTE0SJetq6xv4MjpEikWga6VpFn4jIFFQ/k= github.com/yanyiwu/gojieba v1.4.7/go.mod h1:JUq4DddFVGdHXJHxxepxRmhrKlDpaBxR8O28v6fKYLY= github.com/yuin/gopher-lua v1.1.2 h1:yF/FjE3hD65tBbt0VXLE13HWS9h34fdzJmrWRXwobGA= diff --git a/internal/agent/core/agent.go b/internal/agent/core/agent.go index 4145fc5..8b72946 100644 --- a/internal/agent/core/agent.go +++ b/internal/agent/core/agent.go @@ -165,9 +165,16 @@ func New(cfg AgentConfig) *Agent { if cfg.MaxContextSize <= 0 { cfg.MaxContextSize = 30 } + + embedder := memory.NewStaticEmbedder(strings.Split(cfg.EmbeddingModelPath, ",")...) + if cfg.DocStore != nil { + cfg.DocStore.SetVectorizer(embedder) + cfg.DocStore.ReindexWithVectorizer(embedder) + } + return &Agent{ id: cfg.ID, - startTime: time.Now(), + startTime: time.Now(), provider: cfg.Provider, providerManager: cfg.ProviderManager, io: cfg.IO, @@ -175,7 +182,7 @@ func New(cfg AgentConfig) *Agent { indexer: cfg.Indexer, skills: cfg.Skills, tracker: cfg.Tracker, - context: NewRelevanceContext(cfg.ContextSavePath, memory.NewStaticEmbedder(strings.Split(cfg.EmbeddingModelPath, ",")...)), + context: NewRelevanceContext(cfg.ContextSavePath, embedder), systemPrompt: cfg.SystemPrompt, ctx: ctx, cancel: cancel, @@ -734,6 +741,10 @@ func (a *Agent) process(input string, stageCtx *sdk.StageContext) (response stri break } + if errors.Is(llmErr, context.Canceled) { + // 打断不是 provider 故障，不标记不可用，直接让外层循环重试 + break + } var pe *agentAPI.ProviderError if errors.As(llmErr, &pe) && (pe.StatusCode == 401 || pe.StatusCode == 403) { a.providerManager.ReportStatus(fbProvider.Name(), pe.StatusCode) @@ -1612,8 +1623,8 @@ func (a *Agent) buildSystemPrompt(memContext string, userInput string) string { // 输出指令：使用 output_send__{channel} 作为回复手段 prompt += "\n\n【输出规则】你有多组输出门工具（type=output），每个对应一个输出通道。回复用户时必须调用对应的 output_send__{通道名} 工具。\n" - prompt += "- content 参数是 JSON 字符串，包含要发送的内容。具体格式因通道而异，用 output_send__{通道名}_help 查看每个通道的 JSON 格式说明。\n" - prompt += "- output_send__{通道名}_help 是普通 function 类型工具，调用后返回该通道的 JSON 格式详情和示例。\n" + prompt += "- payload 参数是消息载荷（文本直接填文字），type 指定载荷类型（text/voice/image/file），meta 是 JSON 发送元数据（群号/用户号等）。\n" + prompt += "- 用 output_send__{通道名}_help 查看该通道的 meta 格式和 type 枚举。\n" prompt += "- 同一轮对话中可多次调用输出门工具。长消息应当分多次发出，而不是一口气发完。\n" prompt += "- 直接返回纯文本不会到达任何用户端。" @@ -2081,16 +2092,24 @@ func (a *Agent) buildToolDefs() []interface{} { "type": "function", "function": map[string]interface{}{ "name": "output_send__" + ch.Name, - "description": desc + "。能力: " + capStr + "。content 参数为 JSON 字符串，具体格式请调用 output_send__" + ch.Name + "_help 查看。", + "description": desc + "。能力: " + capStr + "。payload 为消息载荷，meta 为 JSON 发送元数据，type 为载荷类型。用 _help 查看 meta 格式和 type 枚举。", "parameters": map[string]interface{}{ "type": "object", "properties": map[string]interface{}{ - "content": map[string]interface{}{ + "payload": map[string]interface{}{ "type": "string", - "description": "JSON 字符串，包含要发送的内容和路由信息。格式因通道而异，用 _help 工具查看详情。", + "description": "消息载荷。type=text 时填文字，type=file/image 时填 URL 或路径", + }, + "meta": map[string]interface{}{ + "type": "string", + "description": "JSON 对象，包含发送所需的元数据。用 output_send__" + ch.Name + "_help 查看 meta 格式", + }, + "type": map[string]interface{}{ + "type": "string", + "description": "载荷类型，用 channel._help 查看支持的枚举值", }, }, - "required": []string{"content"}, + "required": []string{"payload", "type"}, }, }, }) @@ -2100,7 +2119,7 @@ func (a *Agent) buildToolDefs() []interface{} { "type": "function", "function": map[string]interface{}{ "name": "output_send__" + ch.Name + "_help", - "description": "查看 " + ch.Name + " 输出通道的 JSON 格式说明和示例", + "description": "查看 " + ch.Name + " 输出通道的 meta 格式说明和 type 枚举", "parameters": map[string]interface{}{ "type": "object", "properties": map[string]interface{}{}, @@ -2425,12 +2444,16 @@ func (a *Agent) evaluateGraphQuality() { // 选出低质量候选：generic 关系（如 distiller 自动生成的泛化关系） var lowQuality []string for _, r := range result.Relations { - // 自动蒸馏生成的 (用户, 提及, ...) 和 (AI, 回应, ...) 通常是噪音 + // 自动蒸馏生成的 (用户, 提及, ...), (AI, 回应, ...) 和 jieba 共现 (..., 关联, ...) 通常是噪音 if (r.SourceName == "用户" || r.SourceName == "AI") && (r.RelationType == "提及" || r.RelationType == "回应") { lowQuality = append(lowQuality, fmt.Sprintf("「%s」-「%s」→「%s」", r.SourceName, r.RelationType, r.TargetName)) continue } + if r.RelationType == "关联" { + lowQuality = append(lowQuality, fmt.Sprintf("「%s」-「%s」→「%s」(jieba 共现)", r.SourceName, r.RelationType, r.TargetName)) + continue + } // 极低 mention 的实体+generic 关系 if r.Confidence < 0.3 && r.RelationType != "" { lowQuality = append(lowQuality, fmt.Sprintf("「%s」-「%s」→「%s」(confidence=%.1f)", r.SourceName, r.RelationType, r.TargetName, r.Confidence)) @@ -2518,32 +2541,42 @@ func docToTriples(doc *document.Doc) []memory.Triple { } triples = append(triples, memory.Triple{ - Subject: "文档", - Relation: "包含内容", - Object: doc.Summary, + Subject: "文档", + SubjectType: "Concept", + Relation: "主题", + Object: doc.Summary, + ObjectType: "Topic", + Confidence: 1.0, }) - for _, entity := range doc.Entities { - triples = append(triples, memory.Triple{ - Subject: "文档", - Relation: "提及实体", - Object: entity, - }) - } - - for _, tag := range doc.Tags { - triples = append(triples, memory.Triple{ - Subject: "文档", - Relation: "标签", - Object: tag, - }) + // 用 jieba 精确模式逐句分词 → 相邻词共现三元组 + lines := strings.Split(doc.Content, "\n") + for _, line := range lines { + line = strings.TrimSpace(line) + if line == "" { + continue + } + terms := memory.CutExact(line) + for i := 0; i < len(terms)-1; i++ { + triples = append(triples, memory.Triple{ + Subject: terms[i], + SubjectType: "Concept", + Relation: "关联", + Object: terms[i+1], + ObjectType: "Concept", + Confidence: 0.8, + }) + } } if doc.Source != "" { triples = append(triples, memory.Triple{ - Subject: "文档", - Relation: "来源", - Object: doc.Source, + Subject: "文档", + SubjectType: "Concept", + Relation: "来源", + Object: doc.Source, + ObjectType: "Source", + Confidence: 1.0, }) } @@ -2601,55 +2634,77 @@ func (a *Agent) processConsolidation(evt *agentIO.InputEvent, input string) { // executeOutputSendTool — AI 通过指定通道发送消息（校验通道能力） func (a *Agent) executeOutputSendTool(tc agentAPI.ToolCall) string { - // tool name is "output_send__{channel}" channel := strings.TrimPrefix(tc.Name, "output_send__") - content, _ := tc.Arguments["content"].(string) - if channel == "" || content == "" { - return "工具名称格式: output_send__{channel}，content 不能为空" + payload, _ := tc.Arguments["payload"].(string) + rawType, _ := tc.Arguments["type"].(string) + if channel == "" || payload == "" || rawType == "" { + return "工具名称格式: output_send__{channel}，payload 和 type 不能为空" } - - // content 是一个 JSON 字符串，插件通过解析它确定如何发送消息 - // === Stage: before_output — 输出前插件可审查/改写/拦截 === - stageCtx := &sdk.StageContext{ - FinalText: content, - Phase: sdk.StageBeforeOutput, - } - a.runStage(sdk.StageBeforeOutput, stageCtx) - if stageCtx.Response != nil { - return fmt.Sprintf("输出被插件拦截: %s", *stageCtx.Response) - } - content = stageCtx.FinalText - if content == "" { - return "输出被插件清空" - } - tc.Arguments["content"] = content + meta, _ := tc.Arguments["meta"].(string) // 通道能力检查 caps := a.io.GetChannelCapabilities(channel) if caps == 0 { return fmt.Sprintf("通道 [%s] 不存在或不可用。可用输出工具列表见 output_list_channels", channel) } - if !caps.Supports(agentIO.CapText) { - return fmt.Sprintf("通道 [%s] 不支持文本输出（能力: %s）", channel, caps.String()) + switch rawType { + case "text": + if !caps.Supports(agentIO.CapText) { + return fmt.Sprintf("通道 [%s] 不支持文本输出（能力: %s）", channel, caps.String()) + } + case "voice", "audio": + if !caps.Supports(agentIO.CapAudio) { + return fmt.Sprintf("通道 [%s] 不支持语音输出（能力: %s）", channel, caps.String()) + } + case "image": + if !caps.Supports(agentIO.CapImage) { + return fmt.Sprintf("通道 [%s] 不支持图片输出（能力: %s）", channel, caps.String()) + } + case "file": + if !caps.Supports(agentIO.CapFile) { + return fmt.Sprintf("通道 [%s] 不支持文件输出（能力: %s）", channel, caps.String()) + } } + // 组装参数传给设备处理器 + args := map[string]interface{}{ + "payload": payload, + "type": rawType, + } + if meta != "" { + args["meta"] = meta + } + + // === Stage: before_output — 输出前插件可审查/改写/拦截 === + stageCtx := &sdk.StageContext{ + FinalText: payload, + Phase: sdk.StageBeforeOutput, + } + a.runStage(sdk.StageBeforeOutput, stageCtx) + if stageCtx.Response != nil { + return fmt.Sprintf("输出被插件拦截: %s", *stageCtx.Response) + } + if stageCtx.FinalText == "" { + return "输出被插件清空" + } + args["payload"] = stageCtx.FinalText + // 通过设备处理器投递 if dev := a.io.GetDevice(channel); dev != nil { - result, err := dev.Execute("output", tc.Arguments) + result, err := dev.Execute("output", args) if err != nil { return fmt.Sprintf("通过 [%s] 通道发送失败: %v", channel, err) } return fmt.Sprintf("已通过 [%s] 通道发送: %v", channel, result) } - // 降级：发送到 outputCh（供 OutputChan 消费者） - a.io.EmitTextTo("agent_io", channel, content) + // 降级 + a.io.EmitTextTo("agent_io", channel, payload) return fmt.Sprintf("已通过 [%s] 通道发送", channel) } -// executeOutputSendHelp — 返回指定通道的 JSON 格式说明 +// executeOutputSendHelp — 返回指定通道的 meta 格式和 type 枚举 func (a *Agent) executeOutputSendHelp(tc agentAPI.ToolCall) string { - // tool name is "output_send__{channel}_help" suffix := strings.TrimPrefix(tc.Name, "output_send__") channel := strings.TrimSuffix(suffix, "_help") if channel == "" { @@ -2676,13 +2731,24 @@ func (a *Agent) executeOutputSendHelp(tc agentAPI.ToolCall) string { 描述: %s 能力: %s -【content JSON 格式说明】 -发送到此通道时 content 必须是 JSON 字符串，包含以下字段： -- "content": 消息正文（必填） -- 根据通道不同可能还需要路由字段（如 "group_id", "user_id" 等） +【参数说明】 +payload — 消息载荷（必填）。type=text 时直接填文字，type=file/image 时填 URL 或路径 +meta — JSON 对象，发送所需的元数据（可选，取决于通道是否需要路由信息） +type — 载荷类型（必填），枚举值见下方 -请在通道描述中查看具体字段要求。 -示例: {"content":"你好"}`, channel, desc, capStr) +【type 枚举】 +- text — 文本消息 +- voice — 语音消息 +- image — 图片 +- file — 文件 + +【meta JSON 格式】 +由通道描述定义，通常包含： +- "group_id" 群号（群聊时必填） +- "user_id" 目标用户 QQ 号（私聊时必填） +- "reply_to" 回复某条消息 ID（可选） + +示例: output_send__%s(payload="你好", meta="{\"group_id\": 123456789}", type="text")`, channel, desc, capStr, channel) } // executeOutputListChannels — 列出所有可用通道及其能力 diff --git a/internal/agent/core/agent_functions_test.go b/internal/agent/core/agent_functions_test.go index 0765020..55d6601 100644 --- a/internal/agent/core/agent_functions_test.go +++ b/internal/agent/core/agent_functions_test.go @@ -27,50 +27,36 @@ func TestEntitySimilarity(t *testing.T) { func TestDocToTriples(t *testing.T) { doc := &document.Doc{ - Summary: "用户喜欢编程", - Content: "用户提到喜欢Go和Python", - Tags: []string{"编程", "Go"}, - Entities: []string{"Go", "Python"}, - Source: "context", + Summary: "用户喜欢编程", + Content: "用户提到喜欢Go和Python", + Source: "context", } triples := docToTriples(doc) - if len(triples) == 0 { - t.Fatal("expected non-empty triples") - } foundSummary := false - foundEntity := false - foundTag := false + foundRel := false foundSource := false - for _, tr := range triples { - if tr.Subject == "文档" && tr.Relation == "包含内容" { + switch { + case tr.Subject == "文档" && tr.Relation == "主题": foundSummary = true - } - if tr.Subject == "文档" && tr.Relation == "提及实体" { - foundEntity = true - } - if tr.Subject == "文档" && tr.Relation == "标签" { - foundTag = true - } - if tr.Subject == "文档" && tr.Relation == "来源" { + case tr.Relation == "关联": + foundRel = true + case tr.Subject == "文档" && tr.Relation == "来源": foundSource = true } } if !foundSummary { - t.Error("missing '包含内容' triple") - } - if !foundEntity { - t.Error("missing '提及实体' triple") - } - if !foundTag { - t.Error("missing '标签' triple") + t.Error("missing '主题' triple") } if !foundSource { t.Error("missing '来源' triple") } + if needJieba() && !foundRel { + t.Error("missing '关联' triple with jieba available") + } } func TestDocToTriplesNil(t *testing.T) { @@ -95,15 +81,34 @@ func TestDocToTriplesNoSource(t *testing.T) { func TestDocToTriplesTypes(t *testing.T) { doc := &document.Doc{ - Summary: "测试三元组类型", - Content: "用于验证 SubjectType 和 ObjectType", - Entities: []string{"Go"}, + Summary: "测试三元组类型", + Content: "用于验证 SubjectType 和 ObjectType", + Source: "test", } triples := docToTriples(doc) + + // 主题 and 来源 triples have Subject=文档 for _, tr := range triples { - if tr.Subject != "文档" { - t.Errorf("expected subject '文档', got %q", tr.Subject) + if tr.Subject == "文档" { + if tr.SubjectType != "Concept" { + t.Errorf("文档 subject_type should be Concept, got %q", tr.SubjectType) + } + if tr.Confidence != 1.0 { + t.Errorf("文档 triple confidence should be 1.0, got %f", tr.Confidence) + } + } else { + // 关联 triples use extracted terms as subject/object + if tr.Relation != "关联" { + t.Errorf("non-文档 triple should have 关联 relation, got %q", tr.Relation) + } + if tr.Confidence != 0.8 { + t.Errorf("关联 triple confidence should be 0.8, got %f", tr.Confidence) + } + } + // all should have SubjectType/ObjectType set + if tr.SubjectType == "" || tr.ObjectType == "" { + t.Errorf("triple %+v missing SubjectType or ObjectType", tr) } } } diff --git a/internal/agent/core/agent_helpers_test.go b/internal/agent/core/agent_helpers_test.go index 4629612..005e4e2 100644 --- a/internal/agent/core/agent_helpers_test.go +++ b/internal/agent/core/agent_helpers_test.go @@ -2,8 +2,106 @@ package core import ( "testing" + + "gitcode.com/JianFeeeee/HomeAgent/internal/memory" + "gitcode.com/JianFeeeee/HomeAgent/internal/memory/document" ) +func needJieba() bool { + return memory.GetJieba() != nil +} + +func TestDocToTriplesEmpty(t *testing.T) { + doc := &document.Doc{ + Summary: "empty doc", + Content: "", + Source: "test", + } + triples := docToTriples(doc) + if len(triples) < 2 { + t.Fatalf("expected at least 2 triples (主题+来源), got %d", len(triples)) + } + + if triples[0].Subject != "文档" || triples[0].Relation != "主题" || triples[0].Object != "empty doc" { + t.Errorf("first triple mismatch: %+v", triples[0]) + } + + last := triples[len(triples)-1] + if last.Subject != "文档" || last.Relation != "来源" || last.Object != "test" { + t.Errorf("last triple mismatch: %+v", last) + } +} + +func TestDocToTriplesConversation(t *testing.T) { + doc := &document.Doc{ + Summary: "测试对话 (qq) 涉及: 天气", + Content: "[15:04] qq: 今天天气怎么样\n[15:05] agent: 今天天气很好", + Source: "qq", + } + triples := docToTriples(doc) + + minLen := 2 + hasJieba := needJieba() + + if hasJieba && len(triples) <= minLen { + t.Errorf("expected more than %d triples with jieba, got %d", minLen, len(triples)) + } + + for i, tr := range triples { + if tr.Subject == "" || tr.Relation == "" || tr.Object == "" { + t.Errorf("triple[%d] has empty field: %+v", i, tr) + } + if tr.Confidence <= 0 { + t.Errorf("triple[%d] has non-positive confidence: %+v", i, tr) + } + } + + relCount := 0 + for _, tr := range triples { + if tr.Relation == "关联" { + relCount++ + if tr.Subject == tr.Object { + t.Errorf("关联 triple has same subject and object: %+v", tr) + } + } + } + if hasJieba && relCount == 0 { + t.Errorf("expected 关联 triples with jieba enabled, got 0 in %+v", triples) + } +} + +func TestDocToTriplesMultiLine(t *testing.T) { + doc := &document.Doc{ + Summary: "多轮对话", + Content: "[10:00] user: 你好\n[10:01] agent: 你好，有什么可以帮助你的\n[10:02] user: 今天天气如何\n[10:03] agent: 今天天气很好", + Source: "qq", + } + triples := docToTriples(doc) + if len(triples) < 2 { + t.Fatalf("expected at least 2 triples, got %d", len(triples)) + } + + if triples[0].Subject != "文档" || triples[0].Relation != "主题" { + t.Errorf("first triple should be 主题, got %+v", triples[0]) + } + last := triples[len(triples)-1] + if last.Subject != "文档" || last.Relation != "来源" { + t.Errorf("last triple should be 来源, got %+v", last) + } +} + +func TestDocToTriplesEmptyContent(t *testing.T) { + doc := &document.Doc{ + Summary: "空内容", + Content: "", + Source: "test", + } + triples := docToTriples(doc) + if len(triples) != 2 { + t.Fatalf("expected exactly 2 triples (主题+来源) for empty content, got %d", len(triples)) + } +} + func TestTruncateStr(t *testing.T) { tests := []struct { input string diff --git a/internal/agent/core/context.go b/internal/agent/core/context.go index aa267dc..53ac36c 100644 --- a/internal/agent/core/context.go +++ b/internal/agent/core/context.go @@ -200,7 +200,7 @@ func (c *RelevanceContext) Prune(currentInput string, topK int, docStore *docume Response: s.event.Response, } } - doc, err := docStore.ContextToDoc("context_archived", entries) + doc, err := docStore.ContextToDoc("context_archived", entries, c.embedder) if err == nil && doc != nil { archived = len(archive) } diff --git a/internal/memory/cut.go b/internal/memory/cut.go index b29cb49..092aec5 100644 --- a/internal/memory/cut.go +++ b/internal/memory/cut.go @@ -117,3 +117,26 @@ func ExtractKeywords(text string) []string { } return keywords } + +// CutExact 精确模式分词：返回去停用词后的所有有义项（不限数量），用于 doc→graph 蒸馏 +func CutExact(text string) []string { + text = CleanTemplateText(text) + x := GetJieba() + if x == nil { + return nil + } + words := x.Cut(text, false) + var result []string + seen := make(map[string]bool) + for _, w := range words { + if stopWords[w] || seen[w] { + continue + } + if !validEntityName(w) { + continue + } + seen[w] = true + result = append(result, w) + } + return result +} diff --git a/internal/memory/cut_test.go b/internal/memory/cut_test.go new file mode 100644 index 0000000..5e70387 --- /dev/null +++ b/internal/memory/cut_test.go @@ -0,0 +1,107 @@ +package memory + +import ( + "testing" +) + +func needJieba(t *testing.T) { + t.Helper() + if GetJieba() == nil { + t.Skip("jieba dictionaries not found") + } +} + +func TestCutExact(t *testing.T) { + needJieba(t) + + tests := []struct { + name string + text string + min int + not []string + }{ + { + name: "chinese_sentence", + text: "今天天气怎么样", + min: 2, + not: nil, + }, + { + name: "stop_words_removed", + text: "和天气地", + min: 1, + not: []string{"的", "和"}, + }, + { + name: "short_words_filtered", + text: "今天好天气", + min: 1, + not: []string{"好"}, + }, + { + name: "empty_text", + text: "", + min: 0, + not: nil, + }, + { + name: "qq_conversation", + text: "今天天气怎么样 → 今天天气很好", + min: 2, + not: nil, + }, + { + name: "all_stop_words", + text: "的了呢吗", + min: 0, + not: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := CutExact(tt.text) + if len(got) < tt.min { + t.Errorf("CutExact(%q) = %v (len=%d), want at least %d terms", tt.text, got, len(got), tt.min) + } + for _, forbid := range tt.not { + for _, g := range got { + if g == forbid { + t.Errorf("CutExact(%q) = %v, should not contain %q", tt.text, got, forbid) + } + } + } + }) + } +} + +func TestCutExactNoDuplicates(t *testing.T) { + needJieba(t) + + got := CutExact("天气天气天气") + if len(got) > 1 { + t.Errorf("expected deduplicated result, got %v (len=%d)", got, len(got)) + } +} + +func TestCutExactValidEntityName(t *testing.T) { + needJieba(t) + + got := CutExact("a b c") + for _, g := range got { + if !validEntityName(g) { + t.Errorf("CutExact returned invalid entity name %q", g) + } + } +} + +func TestCutExactRemoveTimestamp(t *testing.T) { + needJieba(t) + + got := CutExact("[15:04] 今天天气不错") + for _, g := range got { + if g == "15" || g == "04" || g == "15:04" { + t.Errorf("timestamp should be removed by CleanTemplateText, got %q in %v", g, got) + } + } +} diff --git a/internal/memory/document/document.go b/internal/memory/document/document.go index 3d494e6..32005a1 100644 --- a/internal/memory/document/document.go +++ b/internal/memory/document/document.go @@ -28,6 +28,7 @@ type Doc struct { Meta map[string]string `json:"meta,omitempty"` AccessCount int `json:"access_count"` // 访问次数 LastAccess time.Time `json:"last_access"` // 最后访问时间 + Vector vector.Vector `json:"vector,omitempty"` // 预计算向量（与 context 同空间），nil 则用 TF-IDF 兜底 } // Store — 文档记忆存储，包含向量索引 @@ -37,12 +38,31 @@ type Store struct { veczer *vector.TFIDFVectorizer mu sync.RWMutex - docs map[string]*Doc - summaries []string // 用于训练向量化器，最大 10000 条 + docs map[string]*Doc + summaries []string // 用于训练向量化器，最大 10000 条 + vectorizer vector.Vectorizer // 可选：与 context 同空间的向量化器 dirty bool } +func (s *Store) SetVectorizer(v vector.Vectorizer) { + s.vectorizer = v +} + +// ReindexWithVectorizer 用给定的向量化器重建所有文档的向量索引 +func (s *Store) ReindexWithVectorizer(v vector.Vectorizer) { + s.mu.Lock() + defer s.mu.Unlock() + + log.Printf("[document memory] reindex with vectorizer (%d docs)", len(s.docs)) + s.vec = vector.NewStore() + for _, doc := range s.docs { + doc.Vector = v.Vectorize(doc.Summary + " " + doc.Content) + s.vec.Insert(doc.ID, doc.Summary, doc.Vector, doc.Meta) + } + log.Printf("[document memory] reindex with vectorizer complete (%d vectors)", s.vec.Size()) +} + const maxSummaries = 10000 func NewStore(dir string) *Store { @@ -88,7 +108,10 @@ func (s *Store) Insert(doc *Doc) error { // 增量训练向量化器并加入向量索引 s.addSummary(doc.Summary) - vec := s.veczer.Vectorize(doc.Summary + " " + doc.Content) + vec := doc.Vector + if vec == nil { + vec = s.veczer.Vectorize(doc.Summary + " " + doc.Content) + } s.vec.Insert(doc.ID, doc.Summary, vec, doc.Meta) // 立即写盘 @@ -101,7 +124,7 @@ func (s *Store) Insert(doc *Doc) error { } // ContextToDoc — 将一段上下文对话历史提炼为文档（带内容去重） -func (s *Store) ContextToDoc(source string, entries []ContextEntry) (*Doc, error) { +func (s *Store) ContextToDoc(source string, entries []ContextEntry, vec vector.Vectorizer) (*Doc, error) { if len(entries) == 0 { return nil, nil } @@ -140,6 +163,12 @@ func (s *Store) ContextToDoc(source string, entries []ContextEntry) (*Doc, error } id := fmt.Sprintf("doc_%d", time.Now().UnixNano()) + var docVec vector.Vector + if vec != nil { + docVec = vec.Vectorize(summary + " " + content) + } else { + docVec = s.veczer.Vectorize(summary + " " + content) + } doc := &Doc{ ID: id, Summary: summary, @@ -152,13 +181,13 @@ func (s *Store) ContextToDoc(source string, entries []ContextEntry) (*Doc, error AccessCount: 1, Source: source, Meta: map[string]string{"content_hash": contentHash}, + Vector: docVec, } s.docs[id] = doc - // 增量训练向量化器并加入向量索引 + // 加入向量索引 s.addSummary(summary) - vec := s.veczer.Vectorize(summary + " " + content) - s.vec.Insert(id, summary, vec, nil) + s.vec.Insert(id, summary, doc.Vector, nil) s.dirty = true s.mu.Unlock() @@ -180,7 +209,7 @@ func (s *Store) Consume(text string, topK int) []*Doc { topK = 5 } - vec := s.veczer.Vectorize(text) + vec := s.vectorizeQuery(text) results := s.vec.Search(vec, topK) var docs []*Doc @@ -194,6 +223,14 @@ func (s *Store) Consume(text string, topK int) []*Doc { return docs } +// vectorizeQuery 用语义向量化器（首选）或 TF-IDF（兜底）处理查询文本 +func (s *Store) vectorizeQuery(text string) vector.Vector { + if s.vectorizer != nil { + return s.vectorizer.Vectorize(text) + } + return s.veczer.Vectorize(text) +} + // Query — 向量相似度查询文档 func (s *Store) Query(text string, topK int) []*Doc { s.mu.RLock() @@ -203,7 +240,7 @@ func (s *Store) Query(text string, topK int) []*Doc { topK = 5 } - vec := s.veczer.Vectorize(text) + vec := s.vectorizeQuery(text) results := s.vec.Search(vec, topK) var docs []*Doc @@ -228,7 +265,10 @@ func (s *Store) Reindex() { s.vec = vector.NewStore() for _, doc := range s.docs { - vec := s.veczer.Vectorize(doc.Summary + " " + doc.Content) + vec := doc.Vector + if vec == nil { + vec = s.veczer.Vectorize(doc.Summary + " " + doc.Content) + } s.vec.Insert(doc.ID, doc.Summary, vec, doc.Meta) } @@ -348,7 +388,10 @@ func (s *Store) loadAll() error { // 重建向量索引 for _, doc := range s.docs { - vec := s.veczer.Vectorize(doc.Summary + " " + doc.Content) + vec := doc.Vector + if vec == nil { + vec = s.veczer.Vectorize(doc.Summary + " " + doc.Content) + } s.vec.Insert(doc.ID, doc.Summary, vec, nil) } diff --git a/internal/memory/document/document_test.go b/internal/memory/document/document_test.go index a3d1513..aa2e07b 100644 --- a/internal/memory/document/document_test.go +++ b/internal/memory/document/document_test.go @@ -82,7 +82,7 @@ func TestContextToDoc(t *testing.T) { {Timestamp: time.Now(), Source: "user", Content: "特别是Go语言", Response: "Go很棒"}, } - doc, err := s.ContextToDoc("test", entries) + doc, err := s.ContextToDoc("test", entries, nil) if err != nil { t.Fatal(err) }