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:
2026-07-17 21:02:35 +08:00
parent d2aec1fd5f
commit 7892d7b0f2
11 changed files with 1734 additions and 69 deletions

View File

@ -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,

View File

@ -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
}
}

View File

@ -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
}