mirror of
https://gitcode.com/JianFeeeee/HomeAgent.git
synced 2026-09-21 17:38:10 +00:00
685ed7e4a0
从一次真实的知识库更新里发现：在线实例更新一个已有条目之后， knowledge_count=32 而 vector_count=33——多出来的那一条是上一版的副本。 成因：vector.Store.Insert 是**追加**语义（s.docs = append + index.Add），不按 id 去重； 而 Store.Add 走的是「写 content.md + 覆盖 items[id] + Insert 向量」。 文件与内存条目都被正确替换了，只有向量索引多留了一份。 危害不在于多占内存：**检索可能命中已被替换掉的旧内容**，而且完全静默—— 条目数看起来是对的，只有向量数比条目数多。 修法：Insert 之前先 s.vec.Remove(id)（Remove 已按 id 过滤 docs 与倒排索引）。 回归测试 TestAddOverwriteReplacesVector 钉住 knowledge_count / vector_count / content.md 三者都必须只剩新版。 注：该文件在 origin/main 上本就有 32 行 gofmt 差异（结构体字段注释对齐）， 不属本次改动，按纪律不做整体重排。
457 lines
11 KiB
Go
457 lines
11 KiB
Go
package knowledge
|
||
|
||
import (
|
||
"encoding/json"
|
||
"fmt"
|
||
"log"
|
||
"os"
|
||
"path/filepath"
|
||
"sort"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
|
||
"gitcode.com/JianFeeeee/HomeAgent/internal/memory"
|
||
"gitcode.com/JianFeeeee/HomeAgent/internal/memory/vector"
|
||
)
|
||
|
||
type Knowledge struct {
|
||
Name string `json:"name"`
|
||
Content string `json:"content"`
|
||
Path string `json:"path"`
|
||
Category string `json:"category,omitempty"` // 父路径，如 "tech/go"
|
||
Tags []string `json:"tags"`
|
||
UpdatedAt time.Time `json:"updated_at"`
|
||
Meta map[string]string `json:"meta,omitempty"`
|
||
}
|
||
|
||
// IndexItem — 索引条目，包含向量特征和内容摘要
|
||
type IndexItem struct {
|
||
Name string `json:"name"`
|
||
Preview string `json:"preview"` // 前 200 字摘要
|
||
Tags []string `json:"tags"`
|
||
Vector map[string]float64 `json:"vector"` // TF-IDF 特征向量（top-N 特征）
|
||
Size int `json:"size"` // 内容总字节数
|
||
}
|
||
|
||
// TreeIndex — 树状索引节点
|
||
type TreeIndex struct {
|
||
Name string `json:"name"`
|
||
Children map[string]*TreeIndex `json:"children,omitempty"`
|
||
Items []IndexItem `json:"items,omitempty"` // 此节点下的知识条目（含向量）
|
||
}
|
||
|
||
func newTreeIndex(name string) *TreeIndex {
|
||
return &TreeIndex{Name: name, Children: make(map[string]*TreeIndex)}
|
||
}
|
||
|
||
// compressVector 压缩向量：保留 topN 个权重最高的特征
|
||
func compressVector(v vector.Vector, topN int) map[string]float64 {
|
||
if len(v) <= topN {
|
||
out := make(map[string]float64, len(v))
|
||
for k, w := range v {
|
||
out[k] = w
|
||
}
|
||
return out
|
||
}
|
||
type kv struct {
|
||
k string
|
||
v float64
|
||
}
|
||
sorted := make([]kv, 0, len(v))
|
||
for k, w := range v {
|
||
sorted = append(sorted, kv{k, w})
|
||
}
|
||
sort.Slice(sorted, func(i, j int) bool {
|
||
return sorted[i].v > sorted[j].v
|
||
})
|
||
if topN > len(sorted) {
|
||
topN = len(sorted)
|
||
}
|
||
sorted = sorted[:topN]
|
||
out := make(map[string]float64, topN)
|
||
for _, kv := range sorted {
|
||
out[kv.k] = kv.v
|
||
}
|
||
return out
|
||
}
|
||
|
||
type Store struct {
|
||
root string
|
||
vec *vector.Store
|
||
veczer *vector.TFIDFVectorizer
|
||
mu sync.RWMutex
|
||
items map[string]*Knowledge
|
||
|
||
indexPath string
|
||
summaries []string
|
||
vectorizer vector.Vectorizer // 可选：词嵌入向量化器，优先于 TF-IDF
|
||
}
|
||
|
||
func NewStore(root string) *Store {
|
||
return &Store{
|
||
root: root,
|
||
indexPath: filepath.Join(root, ".index.json"),
|
||
vec: vector.NewStore(),
|
||
veczer: vector.NewTFIDFVectorizer(memory.TokenizeWords),
|
||
items: make(map[string]*Knowledge),
|
||
}
|
||
}
|
||
|
||
// SetVectorizer 设置词嵌入向量化器，优先于 TF-IDF
|
||
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("[knowledge] reindex with vectorizer (%d items)", len(s.items))
|
||
s.vec = vector.NewStore()
|
||
for _, k := range s.items {
|
||
vec := v.Vectorize(k.Name + " " + k.Content)
|
||
s.vec.Insert(k.Name, k.Name+": "+k.Content, vec, map[string]string{
|
||
"name": k.Name, "path": k.Path,
|
||
})
|
||
}
|
||
log.Printf("[knowledge] reindex with vectorizer complete (%d vectors)", s.vec.Size())
|
||
}
|
||
|
||
// vectorize 优先使用词嵌入向量化器，不可用时回退到 TF-IDF
|
||
func (s *Store) vectorize(text string) vector.Vector {
|
||
if s.vectorizer != nil {
|
||
return s.vectorizer.Vectorize(text)
|
||
}
|
||
return s.veczer.Vectorize(text)
|
||
}
|
||
|
||
func (s *Store) Start() error {
|
||
if err := os.MkdirAll(s.root, 0755); err != nil {
|
||
return fmt.Errorf("knowledge root: %w", err)
|
||
}
|
||
if err := s.scanAll(); err != nil {
|
||
log.Printf("[knowledge] scan error: %v", err)
|
||
}
|
||
// 重建索引文件
|
||
if err := s.writeIndex(); err != nil {
|
||
log.Printf("[knowledge] write index error: %v", err)
|
||
}
|
||
log.Printf("[knowledge] started with %d items, %d vectors", len(s.items), s.vec.Size())
|
||
return nil
|
||
}
|
||
|
||
func (s *Store) Stop() {}
|
||
|
||
func (s *Store) Search(query string, topK int) []*Knowledge {
|
||
s.mu.RLock()
|
||
defer s.mu.RUnlock()
|
||
|
||
if topK <= 0 {
|
||
topK = 5
|
||
}
|
||
|
||
vec := s.vectorize(query)
|
||
results := s.vec.Search(vec, topK)
|
||
|
||
var out []*Knowledge
|
||
for _, r := range results {
|
||
if k, ok := s.items[r.ID]; ok {
|
||
out = append(out, k)
|
||
}
|
||
}
|
||
return out
|
||
}
|
||
|
||
func (s *Store) Add(name, content string) error {
|
||
s.mu.Lock()
|
||
defer s.mu.Unlock()
|
||
|
||
// 解析层级：将 "/" 作为路径分隔符
|
||
category := ""
|
||
leaf := name
|
||
if idx := strings.LastIndex(name, "/"); idx >= 0 {
|
||
category = name[:idx]
|
||
leaf = name[idx+1:]
|
||
}
|
||
dirName := sanitize(leaf)
|
||
if category != "" {
|
||
dirName = sanitize(category) + "/" + dirName
|
||
}
|
||
dir := filepath.Join(s.root, dirName)
|
||
if err := os.MkdirAll(dir, 0755); err != nil {
|
||
return fmt.Errorf("create knowledge dir: %w", err)
|
||
}
|
||
|
||
path := filepath.Join(dir, "content.md")
|
||
if err := os.WriteFile(path, []byte(content), 0644); err != nil {
|
||
return fmt.Errorf("write knowledge: %w", err)
|
||
}
|
||
|
||
now := time.Now()
|
||
id := sanitize(name)
|
||
k := &Knowledge{
|
||
Name: id,
|
||
Content: content,
|
||
Path: path,
|
||
Category: sanitize(category),
|
||
Tags: memory.ExtractKeywords(name + " " + content),
|
||
UpdatedAt: now,
|
||
}
|
||
s.items[id] = k
|
||
|
||
// 覆盖同名条目时必须先摘掉旧向量。
|
||
//
|
||
// vector.Store.Insert 是**追加**语义（s.docs = append + index.Add），不按 id
|
||
// 去重。少了这一步，更新一条知识会在向量索引里留下上一版的副本：条目数看起来
|
||
// 是对的，只有向量数比条目数多——而检索可能因此命中已被替换掉的旧内容。
|
||
s.vec.Remove(id)
|
||
|
||
vec := s.vectorize(name + " " + content)
|
||
s.vec.Insert(id, name+": "+content, vec, map[string]string{
|
||
"name": name, "path": path,
|
||
})
|
||
s.summaries = append(s.summaries, name+" "+content)
|
||
|
||
go func() {
|
||
if err := s.writeIndex(); err != nil {
|
||
log.Printf("[knowledge] write index error after adding %s: %v", name, err)
|
||
}
|
||
}()
|
||
log.Printf("[knowledge] added: %s (%d bytes)", name, len(content))
|
||
return nil
|
||
}
|
||
|
||
func (s *Store) SearchCategories(query string, topK int) []string {
|
||
s.mu.RLock()
|
||
defer s.mu.RUnlock()
|
||
|
||
if query == "" {
|
||
var names []string
|
||
for _, k := range s.items {
|
||
names = append(names, k.Name)
|
||
}
|
||
sort.Strings(names)
|
||
if len(names) > topK {
|
||
names = names[:topK]
|
||
}
|
||
return names
|
||
}
|
||
|
||
vec := s.vectorize(query)
|
||
results := s.vec.Search(vec, topK)
|
||
var names []string
|
||
for _, r := range results {
|
||
if k, ok := s.items[r.ID]; ok {
|
||
names = append(names, k.Name)
|
||
}
|
||
}
|
||
return names
|
||
}
|
||
|
||
func (s *Store) Remove(name string) error {
|
||
s.mu.Lock()
|
||
defer s.mu.Unlock()
|
||
|
||
id := sanitize(name)
|
||
dir := filepath.Join(s.root, id)
|
||
if err := os.RemoveAll(dir); err != nil {
|
||
return err
|
||
}
|
||
delete(s.items, id)
|
||
s.vec.Remove(id)
|
||
go func() {
|
||
if err := s.writeIndex(); err != nil {
|
||
log.Printf("[knowledge] write index error after removing %s: %v", name, err)
|
||
}
|
||
}()
|
||
return nil
|
||
}
|
||
|
||
func (s *Store) Stats() map[string]interface{} {
|
||
s.mu.RLock()
|
||
defer s.mu.RUnlock()
|
||
return map[string]interface{}{
|
||
"knowledge_count": len(s.items),
|
||
"vector_count": s.vec.Size(),
|
||
"root": s.root,
|
||
"index_file": s.indexPath,
|
||
}
|
||
}
|
||
|
||
func (s *Store) List() []string {
|
||
s.mu.RLock()
|
||
defer s.mu.RUnlock()
|
||
var names []string
|
||
for _, k := range s.items {
|
||
names = append(names, k.Name)
|
||
}
|
||
sort.Strings(names)
|
||
return names
|
||
}
|
||
|
||
// BuildTree 从当前知识库构建树状索引（含向量特征）
|
||
func (s *Store) BuildTree() *TreeIndex {
|
||
s.mu.RLock()
|
||
defer s.mu.RUnlock()
|
||
root := newTreeIndex("root")
|
||
for _, k := range s.items {
|
||
node := root
|
||
if k.Category != "" {
|
||
parts := strings.Split(k.Category, "/")
|
||
for _, part := range parts {
|
||
if part == "" {
|
||
continue
|
||
}
|
||
if _, ok := node.Children[part]; !ok {
|
||
node.Children[part] = newTreeIndex(part)
|
||
}
|
||
node = node.Children[part]
|
||
}
|
||
}
|
||
// 获取该条目的向量并压缩
|
||
vec := s.vectorize(k.Name + " " + k.Content)
|
||
preview := []rune(k.Content)
|
||
previewStr := ""
|
||
if len(preview) > 200 {
|
||
previewStr = string(preview[:200]) + "..."
|
||
} else {
|
||
previewStr = string(preview)
|
||
}
|
||
item := IndexItem{
|
||
Name: k.Name,
|
||
Preview: previewStr,
|
||
Tags: k.Tags,
|
||
Vector: compressVector(vec, 20),
|
||
Size: len(k.Content),
|
||
}
|
||
node.Items = append(node.Items, item)
|
||
}
|
||
return root
|
||
}
|
||
|
||
// SearchTree 树状搜索：在树节点下搜索，返回按分类聚合的结果
|
||
func (s *Store) SearchTree(query string, topK int) map[string][]*Knowledge {
|
||
s.mu.RLock()
|
||
defer s.mu.RUnlock()
|
||
|
||
if topK <= 0 {
|
||
topK = 10
|
||
}
|
||
|
||
vec := s.vectorize(query)
|
||
results := s.vec.Search(vec, topK*2)
|
||
|
||
categorized := make(map[string][]*Knowledge)
|
||
for _, r := range results {
|
||
if k, ok := s.items[r.ID]; ok {
|
||
cat := k.Category
|
||
if cat == "" {
|
||
cat = "未分类"
|
||
}
|
||
categorized[cat] = append(categorized[cat], k)
|
||
}
|
||
}
|
||
|
||
out := make(map[string][]*Knowledge)
|
||
for cat, items := range categorized {
|
||
if len(items) > topK {
|
||
items = items[:topK]
|
||
}
|
||
out[cat] = items
|
||
}
|
||
return out
|
||
}
|
||
|
||
// writeIndex 写入 .index.json 树状索引文件（含向量和摘要）
|
||
func (s *Store) writeIndex() error {
|
||
tree := s.BuildTree()
|
||
data, err := json.MarshalIndent(tree, "", " ")
|
||
if err != nil {
|
||
return err
|
||
}
|
||
return os.WriteFile(s.indexPath, data, 0644)
|
||
}
|
||
|
||
// ——— internal ———
|
||
|
||
func (s *Store) scanAll() error {
|
||
entries, err := os.ReadDir(s.root)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
|
||
for _, entry := range entries {
|
||
if !entry.IsDir() {
|
||
continue
|
||
}
|
||
// skip hidden dirs
|
||
if strings.HasPrefix(entry.Name(), ".") {
|
||
continue
|
||
}
|
||
s.scanDir("", entry.Name())
|
||
}
|
||
|
||
if len(s.summaries) > 0 {
|
||
s.veczer.Train(s.summaries)
|
||
}
|
||
|
||
for _, k := range s.items {
|
||
vec := s.vectorize(k.Name + " " + k.Content)
|
||
s.vec.Insert(k.Name, k.Name+": "+k.Content, vec, map[string]string{
|
||
"name": k.Name, "path": k.Path,
|
||
})
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
// scanDir 递归扫描目录
|
||
// category: 父级路径（从知识库根目录算起），如 "tech/go"
|
||
// dirName: 当前目录相对路径（从知识库根目录算起）
|
||
func (s *Store) scanDir(category, dirName string) {
|
||
dir := filepath.Join(s.root, dirName)
|
||
contentPath := filepath.Join(dir, "content.md")
|
||
data, err := os.ReadFile(contentPath)
|
||
if err == nil {
|
||
name := dirName
|
||
content := string(data)
|
||
now := time.Now()
|
||
k := &Knowledge{
|
||
Name: name,
|
||
Content: content,
|
||
Path: contentPath,
|
||
Category: category,
|
||
Tags: memory.ExtractKeywords(dirName + " " + content),
|
||
UpdatedAt: now,
|
||
}
|
||
s.items[name] = k
|
||
s.summaries = append(s.summaries, name+" "+content)
|
||
return
|
||
}
|
||
|
||
// 无 content.md => 是分类目录，递归子目录
|
||
subEntries, err := os.ReadDir(dir)
|
||
if err != nil {
|
||
return
|
||
}
|
||
for _, sub := range subEntries {
|
||
if !sub.IsDir() || strings.HasPrefix(sub.Name(), ".") {
|
||
continue
|
||
}
|
||
childDir := dirName + "/" + sub.Name()
|
||
s.scanDir(dirName, childDir)
|
||
}
|
||
}
|
||
|
||
func sanitize(name string) string {
|
||
name = strings.ToLower(name)
|
||
name = strings.TrimSpace(name)
|
||
name = strings.ReplaceAll(name, " ", "_")
|
||
name = strings.ReplaceAll(name, "\\", "_")
|
||
return name
|
||
}
|
||
|
||
|