mirror of
https://gitcode.com/JianFeeeee/HomeAgent.git
synced 2026-09-22 01:48:11 +00:00
fix: correct context pruning order and vector alignment
- Prune context BEFORE processing (LSTM forget gate pattern) so LLM only sees relevant context, instead of pruning after the fact - Fix ensureTrained() to recompute all event vectors after retraining vectorizer, fixing feature-space mismatch between stored vectors and query vector that made relevance scoring effectively random - Add read lock to knowledge BuildTree() (data race fix) - Log writeIndex() errors instead of discarding them - Fix TOCTOU race in document ContextToDoc() dedup - Fix healthcheck timing (measure elapsed before cleanup) - Refactor waiter CLI into separate files (state, conn, config, editor, history, builtin) for maintainability - Add CLI plugin API key authentication
This commit is contained in:
root
@ -1,6 +1,7 @@
|
||||
package knowledge
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
@ -13,17 +14,67 @@ import (
|
||||
"gitcode.com/JianFeeeee/HomeAgent/internal/memory/vector"
|
||||
)
|
||||
|
||||
// Knowledge — 单条知识
|
||||
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"`
|
||||
}
|
||||
|
||||
// Store — 知识库，文件系统 + 向量索引
|
||||
// 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
|
||||
@ -31,15 +82,17 @@ type Store struct {
|
||||
mu sync.RWMutex
|
||||
items map[string]*Knowledge
|
||||
|
||||
indexPath string
|
||||
summaries []string
|
||||
}
|
||||
|
||||
func NewStore(root string) *Store {
|
||||
return &Store{
|
||||
root: root,
|
||||
vec: vector.NewStore(),
|
||||
veczer: vector.NewTFIDFVectorizer(3),
|
||||
items: make(map[string]*Knowledge),
|
||||
root: root,
|
||||
indexPath: filepath.Join(root, ".index.json"),
|
||||
vec: vector.NewStore(),
|
||||
veczer: vector.NewTFIDFVectorizer(3),
|
||||
items: make(map[string]*Knowledge),
|
||||
}
|
||||
}
|
||||
|
||||
@ -50,13 +103,16 @@ func (s *Store) Start() error {
|
||||
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() {}
|
||||
|
||||
// Search — 向量查询知识
|
||||
func (s *Store) Search(query string, topK int) []*Knowledge {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
@ -77,48 +133,56 @@ func (s *Store) Search(query string, topK int) []*Knowledge {
|
||||
return out
|
||||
}
|
||||
|
||||
// Add — 添加或更新知识
|
||||
func (s *Store) Add(name, content string) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
// 创建知识目录
|
||||
dir := filepath.Join(s.root, sanitize(name))
|
||||
// 解析层级：将 "/" 作为路径分隔符
|
||||
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: name,
|
||||
Name: id,
|
||||
Content: content,
|
||||
Path: path,
|
||||
Category: sanitize(category),
|
||||
Tags: extractKeywords(name + " " + content),
|
||||
UpdatedAt: now,
|
||||
}
|
||||
|
||||
// 生成 ID = 目录名
|
||||
id := sanitize(name)
|
||||
s.items[id] = k
|
||||
|
||||
vec := s.veczer.Vectorize(name + " " + content)
|
||||
s.vec.Insert(id, name+": "+content, vec, map[string]string{
|
||||
"name": name, "path": path,
|
||||
})
|
||||
|
||||
s.summaries = append(s.summaries, name+" "+content)
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
// SearchCategories — 返回所有知识类别
|
||||
func (s *Store) SearchCategories(query string, topK int) []string {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
@ -157,6 +221,9 @@ func (s *Store) Remove(name string) error {
|
||||
}
|
||||
delete(s.items, id)
|
||||
s.vec.Remove(id)
|
||||
if err := s.writeIndex(); err != nil {
|
||||
log.Printf("[knowledge] write index error after removing %s: %v", name, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@ -167,6 +234,7 @@ func (s *Store) Stats() map[string]interface{} {
|
||||
"knowledge_count": len(s.items),
|
||||
"vector_count": s.vec.Size(),
|
||||
"root": s.root,
|
||||
"index_file": s.indexPath,
|
||||
}
|
||||
}
|
||||
|
||||
@ -181,6 +249,89 @@ func (s *Store) List() []string {
|
||||
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.veczer.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.veczer.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 {
|
||||
@ -193,34 +344,17 @@ func (s *Store) scanAll() error {
|
||||
if !entry.IsDir() {
|
||||
continue
|
||||
}
|
||||
dir := filepath.Join(s.root, entry.Name())
|
||||
contentPath := filepath.Join(dir, "content.md")
|
||||
data, err := os.ReadFile(contentPath)
|
||||
if err != nil {
|
||||
// skip hidden dirs
|
||||
if strings.HasPrefix(entry.Name(), ".") {
|
||||
continue
|
||||
}
|
||||
|
||||
name := entry.Name()
|
||||
content := string(data)
|
||||
now := time.Now()
|
||||
|
||||
k := &Knowledge{
|
||||
Name: name,
|
||||
Content: content,
|
||||
Path: contentPath,
|
||||
Tags: extractKeywords(name + " " + content),
|
||||
UpdatedAt: now,
|
||||
}
|
||||
s.items[name] = k
|
||||
s.summaries = append(s.summaries, name+" "+content)
|
||||
s.scanDir("", entry.Name())
|
||||
}
|
||||
|
||||
// 训练向量化器
|
||||
if len(s.summaries) > 0 {
|
||||
s.veczer.Train(s.summaries)
|
||||
}
|
||||
|
||||
// 构建向量索引
|
||||
for _, k := range s.items {
|
||||
vec := s.veczer.Vectorize(k.Name + " " + k.Content)
|
||||
s.vec.Insert(k.Name, k.Name+": "+k.Content, vec, map[string]string{
|
||||
@ -231,11 +365,48 @@ func (s *Store) scanAll() error {
|
||||
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: 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, "/", "_")
|
||||
name = strings.ReplaceAll(name, "\\", "_")
|
||||
return name
|
||||
}
|
||||
@ -255,10 +426,8 @@ func extractKeywords(text string) []string {
|
||||
|
||||
var keywords []string
|
||||
runes := []rune(text)
|
||||
|
||||
seen := make(map[string]bool)
|
||||
|
||||
// bi-gram
|
||||
for i := 0; i < len(runes)-1; i++ {
|
||||
word := string(runes[i : i+2])
|
||||
if !stopWords[word] && len(strings.TrimSpace(word)) == len(word) && !seen[word] {
|
||||
|
||||
Reference in New Issue
Block a user