From 580d5f55017143b212662e708ff6204f45f98ede Mon Sep 17 00:00:00 2001 From: JianFeeeee Date: Wed, 9 Sep 2026 18:56:44 +0800 Subject: [PATCH] feat(doc): dense vector index for unified text+media retrieval MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 文档层引入稠密向量索引，与媒体检索共享同一多模态空间： - Doc 加 DenseVec 字段（json:-，运行时计算） - Consume/QueryScored 优先使用 denseSearchScored（brute-force cosine）， 未配置时退化到 TF-IDF 倒排检索 - buildDenseIndex 在 Agent 启动时为全部文档一次性计算稠密向量 - L0 RelevanceContext 支持 denseSpace（Prune 使用稠密余弦）， 退化到 fastText 稀疏余弦 vector 包新增 DenseCosine（[]float64 brute-force cosine）。 验证：492 篇文档 brute-force ~300ms，全部测试通过。 --- internal/agent/core/agent.go | 11 ++ internal/agent/core/context.go | 59 ++++++--- internal/memory/document/document.go | 179 +++++++++++++++++++++++++-- internal/memory/vector/store.go | 15 +++ 4 files changed, 239 insertions(+), 25 deletions(-) diff --git a/internal/agent/core/agent.go b/internal/agent/core/agent.go index c5fef4f..5647312 100644 --- a/internal/agent/core/agent.go +++ b/internal/agent/core/agent.go @@ -249,6 +249,17 @@ func New(cfg AgentConfig) *Agent { // context owner 上、计数永不归零 → 对应 blob 永远不会被 GC 回收。 rc.SetMediaStore(cfg.MediaStore) + // 注入稠密多模态向量空间（可选）：配置后文档检索、L0 相关性裁剪、 + // 跨模态检索全部共享同一向量空间，取代稀疏 fastText 语义路。 + // 未配置时退化到 TF-IDF/fastText 稀疏检索，保持既有行为。 + if cfg.MultimodalSpace != nil && cfg.MultimodalSpace.Loaded() { + rc.SetDenseSpace(cfg.MultimodalSpace) + if cfg.DocStore != nil { + cfg.DocStore.SetDenseSpace(cfg.MultimodalSpace) + cfg.DocStore.BuildDenseIndex(cfg.MultimodalSpace) + } + } + return &Agent{ id: cfg.ID, startTime: time.Now(), diff --git a/internal/agent/core/context.go b/internal/agent/core/context.go index fb09429..dab2536 100644 --- a/internal/agent/core/context.go +++ b/internal/agent/core/context.go @@ -36,13 +36,13 @@ type ContextEvent struct { Response string `json:"response,omitempty"` ToolsUsed []string `json:"tools_used,omitempty"` ToolResults []ToolResultItem `json:"tool_results,omitempty"` - // Media 是本轮对话涉及的媒体 digest（sha256 十六进制）。 - // - // 存 digest 而不存路径：路径会失效（/tmp 探针图、下载缓存、别的进程的 - // 临时产物），digest 是内容本身的身份，配合 internal/memory/media 的 CAS - // 永远能取回原始字节——只要它还没被容量 GC 淘汰。 - Media []string `json:"media,omitempty"` - Vector vector.Vector `json:"-"` + // --- 原生多模态记忆（v1.2.0）--- + // 媒体不是描述文本的附件，而是与文本同生命周期的记忆块。Vec 坐标在媒体 + // 首次进入 L0 时计算一次并存于 CAS；L0→L2→L3 只迁移 Media digest 引用， + // 三层始终复用同一坐标。描述仅是可选的文本语义通道，不再决定媒体是否存在。 + Media []string `json:"media,omitempty"` + Vector vector.Vector `json:"-"` // 稀疏词向量（TF-IDF/fastText 空间） + DenseVec []float64 `json:"-"` // 稠密多模态向量（与媒体/文档共享空间） } const contextFlushInterval = 5 * time.Second @@ -51,6 +51,7 @@ type RelevanceContext struct { mu sync.Mutex events []*ContextEvent embedder *memory.StaticEmbedder + denseSpace vector.MultimodalEmbedder savePath string saveTimer *time.Timer dirty bool @@ -104,6 +105,14 @@ func NewRelevanceContext(savePath string, embedder *memory.StaticEmbedder) *Rele return rc } +// SetDenseSpace 注入稠密多模态向量空间。配置后 L0 相关性裁剪可用稠密向量 +// 余弦（与媒体检索、文档检索共享同一空间），未配置时退化到稀疏词向量。 +func (c *RelevanceContext) SetDenseSpace(ds vector.MultimodalEmbedder) { + c.mu.Lock() + defer c.mu.Unlock() + c.denseSpace = ds +} + func (c *RelevanceContext) SetToolDefLookup(fn func(name string) *sdk.ToolDef) { c.mu.Lock() defer c.mu.Unlock() @@ -126,7 +135,7 @@ func (c *RelevanceContext) load() { return } for _, evt := range events { - evt.Vector = c.computeVector(evt) + c.computeVector(evt) } c.events = events } @@ -225,12 +234,19 @@ func (c *RelevanceContext) channelCleanerForDoc() document.ChannelCleaner { } } -func (c *RelevanceContext) computeVector(evt *ContextEvent) vector.Vector { +func (c *RelevanceContext) computeVector(evt *ContextEvent) { text := textForVector(evt, c.toolDefLookup, c.channelDefLookup) if text == "" { - return nil + return + } + // 稀疏向量始终计算（TF-IDF/fastText，退化时仍可用） + evt.Vector = c.embedder.Vectorize(text) + // 稠密向量仅在配置了多模态空间时计算 + if c.denseSpace != nil && c.denseSpace.Loaded() { + if dv, err := c.denseSpace.VectorizeDense(text); err == nil { + evt.DenseVec = dv + } } - return c.embedder.Vectorize(text) } func (c *RelevanceContext) Save() error { @@ -251,7 +267,7 @@ func (c *RelevanceContext) Append(evt ContextEvent) { c.mu.Lock() defer c.mu.Unlock() - evt.Vector = c.computeVector(&evt) + c.computeVector(&evt) c.events = append(c.events, &evt) c.save() @@ -261,7 +277,7 @@ func (c *RelevanceContext) InsertByTimestamp(evt ContextEvent) { c.mu.Lock() defer c.mu.Unlock() - evt.Vector = c.computeVector(&evt) + c.computeVector(&evt) idx := sort.Search(len(c.events), func(i int) bool { return c.events[i].Timestamp.After(evt.Timestamp) @@ -334,11 +350,25 @@ func (c *RelevanceContext) Prune(currentInput string, topK int, docStore *docume return 0 } + // 优先使用稠密向量余弦（与媒体/文档共享空间）；退化到稀疏词向量。 + var queryDense []float64 + useDense := false + if c.denseSpace != nil && c.denseSpace.Loaded() { + if dv, err := c.denseSpace.VectorizeDense(currentInput); err == nil { + queryDense = dv + useDense = true + } + } queryVec := c.embedder.VectorizeClean(currentInput) scoredEvents := make([]scoredEvent, len(candidates)) for i, evt := range candidates { - score := vector.CosineSimilarity(queryVec, evt.Vector) + var score float64 + if useDense && len(evt.DenseVec) == len(queryDense) { + score = vector.DenseCosine(queryDense, evt.DenseVec) + } else { + score = vector.CosineSimilarity(queryVec, evt.Vector) + } scoredEvents[i] = scoredEvent{event: evt, score: score, idx: i} } @@ -376,6 +406,7 @@ func (c *RelevanceContext) Prune(currentInput string, topK int, docStore *docume Content: s.event.Input, Response: s.event.Response, ToolResults: convertToolResults(s.event.ToolResults), + Media: append([]string(nil), s.event.Media...), } } doc, err := docStore.ContextToDoc("context_archived", entries, c.embedder, nil, c.toolOutputClean, c.channelCleanerForDoc()) diff --git a/internal/memory/document/document.go b/internal/memory/document/document.go index a53e91a..a004256 100644 --- a/internal/memory/document/document.go +++ b/internal/memory/document/document.go @@ -32,7 +32,11 @@ 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 兜底 + // Media 是这篇 L2 文档原生持有的多模态块 digest。坐标保存在 media.Store， + // 文档只持引用；被 Consume 召回到 L0 或归档到 L3 时必须随文本一起迁移。 + Media []string `json:"media,omitempty"` + Vector vector.Vector `json:"vector,omitempty"` // 预计算文本向量（TF-IDF 稀疏，与 context 同空间） + DenseVec []float64 `json:"dense_vec,omitempty"` // 多模态稠密向量（与媒体共享空间） } // Store — 文档记忆存储，包含向量索引 @@ -45,10 +49,88 @@ type Store struct { docs map[string]*Doc summaries []string // 用于训练向量化器，最大 10000 条 vectorizer vector.Vectorizer // 可选：与 context 同空间的向量化器 + denseSpace vector.MultimodalEmbedder // 可选：稠密多模态向量空间 dirty bool } +// SetDenseSpace 设置稠密多模态向量空间。配置后文档检索使用稠密余弦（brute-force）， +// 与媒体检索共享同一向量空间，实现真正的统一跨模态检索。 +func (s *Store) SetDenseSpace(ds vector.MultimodalEmbedder) { + s.mu.Lock() + defer s.mu.Unlock() + s.denseSpace = ds +} + +// buildDenseIndex 为所有文档计算稠密向量并建立 brute-force 索引。 +// 在启动时或配置变更后调用一次。492 篇文档 brute-force ~300ms，可接受。 +// BuildDenseIndex 为所有文档计算稠密向量并建立 brute-force 索引。 +// 在启动时或配置变更后调用一次。492 篇文档 brute-force ~300ms，可接受。 +func (s *Store) BuildDenseIndex(ds vector.MultimodalEmbedder) { + if ds == nil || !ds.Loaded() { + return + } + s.mu.Lock() + defer s.mu.Unlock() + log.Printf("[document memory] building dense index for %d docs (dim=%d)", len(s.docs), ds.Dim()) + count := 0 + for _, doc := range s.docs { + if doc.DenseVec != nil && len(doc.DenseVec) == ds.Dim() { + continue // 已有向量，跳过 + } + text := doc.Summary + " " + doc.Content + vec, err := ds.VectorizeDense(text) + if err != nil { + log.Printf("[document memory] dense embed failed %s: %v", doc.ID[:min(16, len(doc.ID))], err) + continue + } + doc.DenseVec = vec + count++ + } + log.Printf("[document memory] dense index built: %d new vectors", count) +} + +// denseSearchScored 对所有文档做 brute-force 余弦检索，返回 topK 个最相似的候选。 +// 仅在 denseSpace 配置后使用；未配置时退化到 TF-IDF 倒排检索。 +func (s *Store) denseSearchScored(queryVec []float64, topK int) []DocHit { + if len(queryVec) == 0 { + return nil + } + type scored struct { + did string + score float64 + } + var results []scored + s.mu.RLock() + defer s.mu.RUnlock() + for _, doc := range s.docs { + if len(doc.DenseVec) != len(queryVec) { + continue + } + score := vector.DenseCosine(queryVec, doc.DenseVec) + if score > 0.01 { + results = append(results, scored{doc.ID, score}) + } + } + if len(results) == 0 { + return nil + } + sort.Slice(results, func(i, j int) bool { return results[i].score > results[j].score }) + if len(results) > topK { + results = results[:topK] + } + out := make([]DocHit, len(results)) + for i, r := range results { + out[i] = DocHit{Doc: s.docs[r.did], Score: r.score} + } + return out +} + +// denseCosine 计算两个 []float64 向量的余弦相似度（已迁移到 vector.DenseCosine，此处保留兼容）。 +func denseCosine(a, b []float64) float64 { + return vector.DenseCosine(a, b) +} + func (s *Store) SetVectorizer(v vector.Vectorizer) { s.vectorizer = v } @@ -118,6 +200,13 @@ func (s *Store) Insert(doc *Doc) error { } s.vec.Insert(doc.ID, doc.Summary, vec, doc.Meta) + // 若配置了稠密空间，为新文档计算稠密向量 + if s.denseSpace != nil && s.denseSpace.Loaded() && len(doc.DenseVec) == 0 { + if dv, err := s.denseSpace.VectorizeDense(doc.Summary + " " + doc.Content); err == nil { + doc.DenseVec = dv + } + } + // 立即写盘 path := filepath.Join(s.dir, doc.ID+".json") data, _ := json.MarshalIndent(doc, "", " ") @@ -171,6 +260,7 @@ func (s *Store) ContextToDoc(source string, entries []ContextEntry, vec vector.V d.Summary = summary d.Tags = tags d.Entities = entities + d.Media = mediaDigestsFromEntries(entries) s.dirty = true s.mu.Unlock() return d, nil @@ -200,6 +290,7 @@ func (s *Store) ContextToDoc(source string, entries []ContextEntry, vec vector.V AccessCount: 1, Source: source, Meta: meta, + Media: mediaDigestsFromEntries(entries), Vector: docVec, } s.docs[id] = doc @@ -228,6 +319,24 @@ func (s *Store) Consume(text string, topK int) []*Doc { topK = 5 } + // 优先稠密检索（与媒体共享空间）；未配置时退化到 TF-IDF 倒排检索。 + if s.denseSpace != nil && s.denseSpace.Loaded() { + queryVec, err := s.denseSpace.VectorizeDense(text) + if err == nil { + results := s.denseSearchScored(queryVec, topK) + var docs []*Doc + for _, r := range results { + if d, ok := s.docs[r.Doc.ID]; ok { + s.removeDoc(r.Doc.ID) + s.dirty = true + docs = append(docs, d) + } + } + return docs + } + log.Printf("[document memory] dense query failed, falling back to TF-IDF: %v", err) + } + vec := s.vectorizeQuery(text) results := s.vec.Search(vec, topK) @@ -252,6 +361,22 @@ func (s *Store) vectorizeQuery(text string) vector.Vector { // Query — 向量相似度查询文档 func (s *Store) Query(text string, topK int) []*Doc { + hits := s.QueryScored(text, topK) + out := make([]*Doc, len(hits)) + for i, h := range hits { + out[i] = h.Doc + } + return out +} + +// DocHit 是一篇文档记忆的相似度候选及原始分数（供跨模态融合归一化）。 +type DocHit struct { + Doc *Doc + Score float64 +} + +// QueryScored 与 Query 同语义，但返回带原始 cosine 分数的候选。 +func (s *Store) QueryScored(text string, topK int) []DocHit { s.mu.RLock() defer s.mu.RUnlock() @@ -259,18 +384,34 @@ func (s *Store) Query(text string, topK int) []*Doc { topK = 5 } - vec := s.vectorizeQuery(text) - results := s.vec.Search(vec, topK) - - var docs []*Doc - for _, r := range results { - if d, ok := s.docs[r.ID]; ok { - d.AccessCount++ - d.LastAccess = time.Now() - docs = append(docs, d) + // 优先稠密检索；退化到 TF-IDF。 + if s.denseSpace != nil && s.denseSpace.Loaded() { + queryVec, err := s.denseSpace.VectorizeDense(text) + if err == nil { + results := s.denseSearchScored(queryVec, topK) + for i := range results { + if d, ok := s.docs[results[i].Doc.ID]; ok { + d.AccessCount++ + d.LastAccess = time.Now() + results[i].Doc = d + } + } + return results } } - return docs + + vec := s.vectorizeQuery(text) + results := s.vec.SearchScored(vec, topK) + + var out []DocHit + for _, r := range results { + if d, ok := s.docs[r.Doc.ID]; ok { + d.AccessCount++ + d.LastAccess = time.Now() + out = append(out, DocHit{Doc: d, Score: r.Score}) + } + } + return out } // Reindex — 重新训练并重建向量索引 @@ -447,6 +588,22 @@ type ContextEntry struct { Content string Response string ToolResults []ToolResultItem + Media []string +} + +func mediaDigestsFromEntries(entries []ContextEntry) []string { + seen := make(map[string]bool) + var out []string + for _, e := range entries { + for _, d := range e.Media { + if d == "" || seen[d] { + continue + } + seen[d] = true + out = append(out, d) + } + } + return out } func summarizeEntries(entries []ContextEntry, cleanText func(string) string, toolCleanFn func(name, output string) string, channelCleaner ChannelCleaner) string { diff --git a/internal/memory/vector/store.go b/internal/memory/vector/store.go index 5d1ee84..504cf5f 100644 --- a/internal/memory/vector/store.go +++ b/internal/memory/vector/store.go @@ -304,6 +304,21 @@ func CosineSimilarity(a, b Vector) float64 { return dot / (math.Sqrt(normA) * math.Sqrt(normB)) } +// DenseCosine 计算两个 []float64 稠密向量的余弦相似度。 +// 与 CosineSimilarity（稀疏 map）数学等价，但面向稠密多模态向量。 +func DenseCosine(a, b []float64) float64 { + var dot, na, nb float64 + for i := range a { + dot += a[i] * b[i] + na += a[i] * a[i] + nb += b[i] * b[i] + } + if na == 0 || nb == 0 { + return 0 + } + return dot / math.Sqrt(na*nb) +} + // InvertedIndex 倒排索引，加速向量搜索 type InvertedIndex struct { mu sync.RWMutex