mirror of
https://gitcode.com/JianFeeeee/HomeAgent.git
synced 2026-09-21 09:28:14 +00:00
feat: complete HomeAgent architecture v2
- IO abstraction layer with OutputChannel routing and capability validation - Three-layer memory (Context-Document-Graph) with TF-IDF relevance pruning - OneBot V11 QQ protocol plugin with Reverse WebSocket client - Plugin system with hot-reload (SKILL.md + native factories) - Knowledge system with TF-IDF vector indexing - Personality system (personal.md) - Text memory (JSONL with rotation) - Change tracker (overlayfs) with rollback - Lua adapter VM - Design document (DESIGN.md) Module: gitcode.com/JianFeeeee/HomeAgent
This commit is contained in:
root
388
internal/memory/document/document.go
Normal file
388
internal/memory/document/document.go
Normal file
@ -0,0 +1,388 @@
|
||||
package document
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"gitcode.com/JianFeeeee/HomeAgent/internal/memory/vector"
|
||||
)
|
||||
|
||||
// Doc — 记忆文档：由上下文提炼而来
|
||||
type Doc struct {
|
||||
ID string `json:"id"`
|
||||
Summary string `json:"summary"`
|
||||
Content string `json:"content"`
|
||||
Tags []string `json:"tags"`
|
||||
Entities []string `json:"entities"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
Source string `json:"source"` // context / graph / manual
|
||||
Meta map[string]string `json:"meta,omitempty"`
|
||||
AccessCount int `json:"access_count"` // 访问次数
|
||||
LastAccess time.Time `json:"last_access"` // 最后访问时间
|
||||
}
|
||||
|
||||
// Store — 文档记忆存储，包含向量索引
|
||||
type Store struct {
|
||||
dir string
|
||||
vec *vector.Store
|
||||
veczer *vector.TFIDFVectorizer
|
||||
mu sync.RWMutex
|
||||
|
||||
docs map[string]*Doc
|
||||
summaries []string // 用于训练向量化器
|
||||
|
||||
dirty bool
|
||||
}
|
||||
|
||||
func NewStore(dir string) *Store {
|
||||
return &Store{
|
||||
dir: dir,
|
||||
vec: vector.NewStore(),
|
||||
veczer: vector.NewTFIDFVectorizer(2),
|
||||
docs: make(map[string]*Doc),
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Store) Start() error {
|
||||
if err := os.MkdirAll(s.dir, 0755); err != nil {
|
||||
return fmt.Errorf("document store dir: %w", err)
|
||||
}
|
||||
if err := s.loadAll(); err != nil {
|
||||
log.Printf("[document memory] load error: %v", err)
|
||||
}
|
||||
log.Printf("[document memory] started with %d docs, %d vectors", len(s.docs), s.vec.Size())
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Store) Stop() {
|
||||
s.flush()
|
||||
}
|
||||
|
||||
// Insert 创建/更新文档
|
||||
func (s *Store) Insert(doc *Doc) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
if doc.ID == "" {
|
||||
doc.ID = fmt.Sprintf("doc_%d", time.Now().UnixNano())
|
||||
doc.CreatedAt = time.Now()
|
||||
}
|
||||
doc.UpdatedAt = time.Now()
|
||||
doc.LastAccess = time.Now()
|
||||
if doc.AccessCount == 0 {
|
||||
doc.AccessCount = 1
|
||||
}
|
||||
|
||||
s.docs[doc.ID] = doc
|
||||
|
||||
vec := s.veczer.Vectorize(doc.Summary + " " + doc.Content)
|
||||
s.vec.Insert(doc.ID, doc.Summary, vec, doc.Meta)
|
||||
|
||||
// 更新训练集
|
||||
s.summaries = append(s.summaries, doc.Summary)
|
||||
|
||||
s.dirty = true
|
||||
return nil
|
||||
}
|
||||
|
||||
// ContextToDoc — 将一段上下文对话历史提炼为文档
|
||||
func (s *Store) ContextToDoc(source string, entries []ContextEntry) (*Doc, error) {
|
||||
if len(entries) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
var parts []string
|
||||
for _, e := range entries {
|
||||
line := fmt.Sprintf("[%s] %s: %s", e.Timestamp.Format("15:04"), e.Source, e.Content)
|
||||
if e.Response != "" {
|
||||
line += fmt.Sprintf(" → %s", truncate(e.Response, 100))
|
||||
}
|
||||
parts = append(parts, line)
|
||||
}
|
||||
content := strings.Join(parts, "\n")
|
||||
|
||||
summary := summarizeEntries(entries)
|
||||
tags := extractTags(entries)
|
||||
entities := extractEntities(entries)
|
||||
|
||||
doc := &Doc{
|
||||
ID: fmt.Sprintf("doc_%d", time.Now().UnixNano()),
|
||||
Summary: summary,
|
||||
Content: content,
|
||||
Tags: tags,
|
||||
Entities: entities,
|
||||
CreatedAt: time.Now(),
|
||||
UpdatedAt: time.Now(),
|
||||
Source: source,
|
||||
}
|
||||
|
||||
if err := s.Insert(doc); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return doc, nil
|
||||
}
|
||||
|
||||
// Query — 向量相似度查询文档
|
||||
func (s *Store) Query(text string, topK int) []*Doc {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
|
||||
if topK <= 0 {
|
||||
topK = 5
|
||||
}
|
||||
|
||||
vec := s.veczer.Vectorize(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)
|
||||
}
|
||||
}
|
||||
return docs
|
||||
}
|
||||
|
||||
// Reindex — 重新训练并重建向量索引
|
||||
func (s *Store) Reindex() {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
log.Printf("[document memory] reindexing %d docs", len(s.docs))
|
||||
|
||||
s.veczer.Train(s.summaries)
|
||||
|
||||
s.vec = vector.NewStore()
|
||||
for _, doc := range s.docs {
|
||||
vec := s.veczer.Vectorize(doc.Summary + " " + doc.Content)
|
||||
s.vec.Insert(doc.ID, doc.Summary, vec, doc.Meta)
|
||||
}
|
||||
|
||||
log.Printf("[document memory] reindex complete (%d vectors)", s.vec.Size())
|
||||
}
|
||||
|
||||
func (s *Store) Stats() map[string]interface{} {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
|
||||
return map[string]interface{}{
|
||||
"doc_count": len(s.docs),
|
||||
"vector_count": s.vec.Size(),
|
||||
"summary_count": len(s.summaries),
|
||||
"dir": s.dir,
|
||||
}
|
||||
}
|
||||
|
||||
// FindColdDocs — 查找冷文档：超过 maxAge 未访问且访问次数 <= minAccess
|
||||
func (s *Store) FindColdDocs(maxAge time.Duration, minAccess int) []*Doc {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
|
||||
cutoff := time.Now().Add(-maxAge)
|
||||
var cold []*Doc
|
||||
for _, d := range s.docs {
|
||||
if d.AccessCount <= minAccess && d.LastAccess.Before(cutoff) {
|
||||
cold = append(cold, d)
|
||||
}
|
||||
}
|
||||
return cold
|
||||
}
|
||||
|
||||
func (s *Store) RecentDocs(n int) []*Doc {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
|
||||
var list []*Doc
|
||||
for _, d := range s.docs {
|
||||
list = append(list, d)
|
||||
}
|
||||
sort.Slice(list, func(i, j int) bool {
|
||||
return list[i].CreatedAt.After(list[j].CreatedAt)
|
||||
})
|
||||
if len(list) > n {
|
||||
list = list[:n]
|
||||
}
|
||||
return list
|
||||
}
|
||||
|
||||
// ——— internal ———
|
||||
|
||||
func (s *Store) loadAll() error {
|
||||
entries, err := os.ReadDir(s.dir)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
for _, e := range entries {
|
||||
if !strings.HasSuffix(e.Name(), ".json") || !strings.HasPrefix(e.Name(), "doc_") {
|
||||
continue
|
||||
}
|
||||
path := filepath.Join(s.dir, e.Name())
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
var doc Doc
|
||||
if err := json.Unmarshal(data, &doc); err != nil {
|
||||
continue
|
||||
}
|
||||
s.docs[doc.ID] = &doc
|
||||
s.summaries = append(s.summaries, doc.Summary)
|
||||
}
|
||||
|
||||
// 训练向量化器
|
||||
if len(s.summaries) > 0 {
|
||||
s.veczer.Train(s.summaries)
|
||||
}
|
||||
|
||||
// 重建向量索引
|
||||
for _, doc := range s.docs {
|
||||
vec := s.veczer.Vectorize(doc.Summary + " " + doc.Content)
|
||||
s.vec.Insert(doc.ID, doc.Summary, vec, nil)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Store) flush() {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
if !s.dirty {
|
||||
return
|
||||
}
|
||||
|
||||
for _, doc := range s.docs {
|
||||
path := filepath.Join(s.dir, doc.ID+".json")
|
||||
data, err := json.MarshalIndent(doc, "", " ")
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
os.WriteFile(path, data, 0644)
|
||||
}
|
||||
s.dirty = false
|
||||
}
|
||||
|
||||
type ContextEntry struct {
|
||||
Timestamp time.Time
|
||||
Source string
|
||||
Content string
|
||||
Response string
|
||||
}
|
||||
|
||||
func summarizeEntries(entries []ContextEntry) string {
|
||||
if len(entries) == 0 {
|
||||
return ""
|
||||
}
|
||||
sources := make(map[string]int)
|
||||
var topics []string
|
||||
for _, e := range entries {
|
||||
sources[e.Source]++
|
||||
words := extractKeywords(e.Content)
|
||||
topics = append(topics, words...)
|
||||
}
|
||||
|
||||
summary := fmt.Sprintf("来自 %d 个来源的 %d 条对话", len(sources), len(entries))
|
||||
var srcList []string
|
||||
for s := range sources {
|
||||
srcList = append(srcList, s)
|
||||
}
|
||||
summary += " (" + strings.Join(srcList, ", ") + ")"
|
||||
|
||||
if len(topics) > 0 {
|
||||
seen := make(map[string]bool)
|
||||
var uniq []string
|
||||
for _, t := range topics {
|
||||
if !seen[t] {
|
||||
seen[t] = true
|
||||
uniq = append(uniq, t)
|
||||
}
|
||||
}
|
||||
if len(uniq) > 5 {
|
||||
uniq = uniq[:5]
|
||||
}
|
||||
summary += " 涉及: " + strings.Join(uniq, ", ")
|
||||
}
|
||||
|
||||
return summary
|
||||
}
|
||||
|
||||
func extractTags(entries []ContextEntry) []string {
|
||||
tagSet := make(map[string]bool)
|
||||
for _, e := range entries {
|
||||
for _, kw := range extractKeywords(e.Content) {
|
||||
tagSet[kw] = true
|
||||
}
|
||||
}
|
||||
var tags []string
|
||||
for t := range tagSet {
|
||||
if len(tags) >= 10 {
|
||||
break
|
||||
}
|
||||
tags = append(tags, t)
|
||||
}
|
||||
return tags
|
||||
}
|
||||
|
||||
func extractEntities(entries []ContextEntry) []string {
|
||||
// 简易实体提取：提取引号内的内容、粗体/标记词
|
||||
var entities []string
|
||||
seen := make(map[string]bool)
|
||||
for _, e := range entries {
|
||||
for _, kw := range extractKeywords(e.Content) {
|
||||
if len(kw) >= 2 && !seen[kw] {
|
||||
seen[kw] = true
|
||||
entities = append(entities, kw)
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(entities) > 20 {
|
||||
entities = entities[:20]
|
||||
}
|
||||
return entities
|
||||
}
|
||||
|
||||
func extractKeywords(text string) []string {
|
||||
stopWords := map[string]bool{
|
||||
"的": true, "了": true, "是": true, "在": true, "有": true,
|
||||
"和": true, "就": true, "不": true, "人": true, "都": true,
|
||||
"一": true, "一个": true, "上": true, "也": true, "很": true,
|
||||
"到": true, "说": true, "要": true, "去": true, "你": true,
|
||||
"会": true, "着": true, "没有": true, "看": true, "好": true,
|
||||
"自己": true, "这": true, "他": true, "她": true, "它": true,
|
||||
"什么": true, "怎么": true, "为什么": true, "如何": true,
|
||||
"我": true, "我们": true, "你们": true, "他们": true, "这个": true,
|
||||
"那个": true, "可以": true, "吗": true, "吧": true, "啊": true,
|
||||
}
|
||||
|
||||
var keywords []string
|
||||
runes := []rune(text)
|
||||
|
||||
// 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) {
|
||||
keywords = append(keywords, word)
|
||||
}
|
||||
}
|
||||
return keywords
|
||||
}
|
||||
|
||||
func truncate(s string, max int) string {
|
||||
runes := []rune(s)
|
||||
if len(runes) > max {
|
||||
return string(runes[:max]) + "..."
|
||||
}
|
||||
return s
|
||||
}
|
||||
555
internal/memory/graph.go
Normal file
555
internal/memory/graph.go
Normal file
@ -0,0 +1,555 @@
|
||||
package memory
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
_ "github.com/mattn/go-sqlite3"
|
||||
)
|
||||
|
||||
type Entity struct {
|
||||
ID int64 `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Type string `json:"type"`
|
||||
MentionCount int `json:"mention_count"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
type Relation struct {
|
||||
ID int64 `json:"id"`
|
||||
SourceID int64 `json:"source_id"`
|
||||
TargetID int64 `json:"target_id"`
|
||||
SourceName string `json:"source_name"`
|
||||
TargetName string `json:"target_name"`
|
||||
RelationType string `json:"relation_type"`
|
||||
Confidence float64 `json:"confidence"`
|
||||
Status string `json:"status"`
|
||||
SessionID string `json:"session_id"`
|
||||
TurnID int `json:"turn_id"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
DateBucket string `json:"date_bucket"`
|
||||
}
|
||||
|
||||
type Triple struct {
|
||||
Subject string `json:"subject"`
|
||||
Relation string `json:"relation"`
|
||||
Object string `json:"object"`
|
||||
Confidence float64 `json:"confidence,omitempty"`
|
||||
SubjectType string `json:"subject_type,omitempty"`
|
||||
ObjectType string `json:"object_type,omitempty"`
|
||||
}
|
||||
|
||||
type GraphDB struct {
|
||||
db *sql.DB
|
||||
mu sync.RWMutex
|
||||
dbPath string
|
||||
}
|
||||
|
||||
func NewGraphDB(dbPath string) (*GraphDB, error) {
|
||||
db, err := sql.Open("sqlite3", dbPath+"?_journal_mode=WAL&_foreign_keys=on")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("open graph db: %w", err)
|
||||
}
|
||||
|
||||
g := &GraphDB{db: db, dbPath: dbPath}
|
||||
if err := g.initSchema(); err != nil {
|
||||
return nil, fmt.Errorf("init schema: %w", err)
|
||||
}
|
||||
|
||||
return g, nil
|
||||
}
|
||||
|
||||
func (g *GraphDB) initSchema() error {
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
|
||||
tx, err := g.db.Begin()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
|
||||
schemas := []string{
|
||||
`CREATE TABLE IF NOT EXISTS entities (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
name TEXT UNIQUE NOT NULL,
|
||||
type TEXT DEFAULT 'Concept',
|
||||
mention_count INTEGER DEFAULT 1,
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
)`,
|
||||
`CREATE TABLE IF NOT EXISTS relations (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
source_id INTEGER NOT NULL,
|
||||
target_id INTEGER NOT NULL,
|
||||
relation_type TEXT NOT NULL,
|
||||
confidence REAL DEFAULT 1.0,
|
||||
status TEXT DEFAULT 'active',
|
||||
session_id TEXT,
|
||||
turn_id INTEGER DEFAULT 0,
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
date_bucket TEXT,
|
||||
FOREIGN KEY (source_id) REFERENCES entities(id),
|
||||
FOREIGN KEY (target_id) REFERENCES entities(id)
|
||||
)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_entity_name ON entities(name)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_entity_type ON entities(type)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_relation_source ON relations(source_id)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_relation_target ON relations(target_id)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_relation_type ON relations(relation_type)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_relation_status ON relations(status)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_relation_session ON relations(session_id)`,
|
||||
}
|
||||
|
||||
for _, s := range schemas {
|
||||
if _, err := tx.Exec(s); err != nil {
|
||||
return fmt.Errorf("schema exec: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
func (g *GraphDB) Commit(triples []Triple, sessionID string, turnID int) (int, int, error) {
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
|
||||
tx, err := g.db.Begin()
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
|
||||
entitiesCreated := 0
|
||||
relationsCreated := 0
|
||||
dateBucket := time.Now().Format("2006-01-02")
|
||||
|
||||
for _, t := range triples {
|
||||
if t.Subject == "" || t.Relation == "" || t.Object == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
subjType := t.SubjectType
|
||||
if subjType == "" {
|
||||
subjType = "Concept"
|
||||
}
|
||||
objType := t.ObjectType
|
||||
if objType == "" {
|
||||
objType = "Concept"
|
||||
}
|
||||
confidence := t.Confidence
|
||||
if confidence == 0 {
|
||||
confidence = 1.0
|
||||
}
|
||||
|
||||
ec, err := g.upsertEntity(tx, t.Subject, subjType)
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
entitiesCreated += ec
|
||||
|
||||
ec, err = g.upsertEntity(tx, t.Object, objType)
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
entitiesCreated += ec
|
||||
|
||||
var sourceID, targetID int64
|
||||
err = tx.QueryRow("SELECT id FROM entities WHERE name = ?", t.Subject).Scan(&sourceID)
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
err = tx.QueryRow("SELECT id FROM entities WHERE name = ?", t.Object).Scan(&targetID)
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
|
||||
_, err = tx.Exec(
|
||||
`INSERT INTO relations (source_id, target_id, relation_type, confidence, session_id, turn_id, date_bucket)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?)`,
|
||||
sourceID, targetID, t.Relation, confidence, sessionID, turnID, dateBucket,
|
||||
)
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
relationsCreated++
|
||||
}
|
||||
|
||||
if err := tx.Commit(); err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
|
||||
return entitiesCreated, relationsCreated, nil
|
||||
}
|
||||
|
||||
func (g *GraphDB) upsertEntity(tx *sql.Tx, name string, entityType string) (int, error) {
|
||||
result, err := tx.Exec(
|
||||
`INSERT INTO entities (name, type) VALUES (?, ?)
|
||||
ON CONFLICT(name) DO UPDATE SET
|
||||
mention_count = mention_count + 1,
|
||||
updated_at = CURRENT_TIMESTAMP`,
|
||||
name, entityType,
|
||||
)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
rows, _ := result.RowsAffected()
|
||||
if rows > 0 {
|
||||
return 1, nil
|
||||
}
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
type RecallResult struct {
|
||||
Entities []Entity `json:"entities"`
|
||||
Relations []Relation `json:"relations"`
|
||||
}
|
||||
|
||||
func (g *GraphDB) Recall(keywords []string, seedEntities []string, depth int, sessionFilter string) (*RecallResult, error) {
|
||||
g.mu.RLock()
|
||||
defer g.mu.RUnlock()
|
||||
|
||||
result := &RecallResult{}
|
||||
|
||||
if len(keywords) == 0 && len(seedEntities) == 0 {
|
||||
rows, err := g.db.Query(
|
||||
`SELECT id, name, type, mention_count, created_at, updated_at
|
||||
FROM entities ORDER BY mention_count DESC LIMIT 50`,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
for rows.Next() {
|
||||
var e Entity
|
||||
if err := rows.Scan(&e.ID, &e.Name, &e.Type, &e.MentionCount, &e.CreatedAt, &e.UpdatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result.Entities = append(result.Entities, e)
|
||||
}
|
||||
|
||||
relRows, err := g.db.Query(
|
||||
`SELECT r.id, r.source_id, r.target_id, e1.name, e2.name,
|
||||
r.relation_type, r.confidence, r.status, r.session_id,
|
||||
r.turn_id, r.created_at, COALESCE(r.date_bucket, '')
|
||||
FROM relations r
|
||||
JOIN entities e1 ON r.source_id = e1.id
|
||||
JOIN entities e2 ON r.target_id = e2.id
|
||||
WHERE r.status = 'active'
|
||||
ORDER BY r.created_at DESC LIMIT 30`,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer relRows.Close()
|
||||
for relRows.Next() {
|
||||
var rel Relation
|
||||
if err := relRows.Scan(&rel.ID, &rel.SourceID, &rel.TargetID,
|
||||
&rel.SourceName, &rel.TargetName, &rel.RelationType,
|
||||
&rel.Confidence, &rel.Status, &rel.SessionID,
|
||||
&rel.TurnID, &rel.CreatedAt, &rel.DateBucket); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result.Relations = append(result.Relations, rel)
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
entityIDs := make(map[int64]bool)
|
||||
|
||||
for _, kw := range keywords {
|
||||
rows, err := g.db.Query(
|
||||
`SELECT id, name, type, mention_count, created_at, updated_at
|
||||
FROM entities WHERE LOWER(name) LIKE ?`,
|
||||
"%"+kw+"%",
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
for rows.Next() {
|
||||
var e Entity
|
||||
if err := rows.Scan(&e.ID, &e.Name, &e.Type, &e.MentionCount, &e.CreatedAt, &e.UpdatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !entityIDs[e.ID] {
|
||||
entityIDs[e.ID] = true
|
||||
result.Entities = append(result.Entities, e)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for _, se := range seedEntities {
|
||||
row := g.db.QueryRow(
|
||||
`SELECT id, name, type, mention_count, created_at, updated_at
|
||||
FROM entities WHERE name = ?`, se)
|
||||
var e Entity
|
||||
if err := row.Scan(&e.ID, &e.Name, &e.Type, &e.MentionCount, &e.CreatedAt, &e.UpdatedAt); err == nil {
|
||||
if !entityIDs[e.ID] {
|
||||
entityIDs[e.ID] = true
|
||||
result.Entities = append(result.Entities, e)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if len(entityIDs) == 0 {
|
||||
return result, nil
|
||||
}
|
||||
|
||||
for depthLevel := 0; depthLevel < depth; depthLevel++ {
|
||||
ids := make([]interface{}, 0, len(entityIDs))
|
||||
for id := range entityIDs {
|
||||
ids = append(ids, id)
|
||||
}
|
||||
|
||||
if len(ids) == 0 {
|
||||
break
|
||||
}
|
||||
|
||||
query := fmt.Sprintf(
|
||||
`SELECT r.id, r.source_id, r.target_id, e1.name, e2.name,
|
||||
r.relation_type, r.confidence, r.status, r.session_id,
|
||||
r.turn_id, r.created_at, COALESCE(r.date_bucket, '')
|
||||
FROM relations r
|
||||
JOIN entities e1 ON r.source_id = e1.id
|
||||
JOIN entities e2 ON r.target_id = e2.id
|
||||
WHERE (r.source_id IN (%s) OR r.target_id IN (%s))
|
||||
AND r.status = 'active'`,
|
||||
placeholders(len(ids)),
|
||||
placeholders(len(ids)),
|
||||
)
|
||||
allIDs := append(ids, ids...)
|
||||
|
||||
if sessionFilter != "" {
|
||||
query += " AND r.session_id = ?"
|
||||
allIDs = append(allIDs, sessionFilter)
|
||||
}
|
||||
|
||||
relRows, err := g.db.Query(query, allIDs...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer relRows.Close()
|
||||
|
||||
newIDs := make(map[int64]bool)
|
||||
for relRows.Next() {
|
||||
var rel Relation
|
||||
if err := relRows.Scan(&rel.ID, &rel.SourceID, &rel.TargetID,
|
||||
&rel.SourceName, &rel.TargetName, &rel.RelationType,
|
||||
&rel.Confidence, &rel.Status, &rel.SessionID,
|
||||
&rel.TurnID, &rel.CreatedAt, &rel.DateBucket); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result.Relations = append(result.Relations, rel)
|
||||
|
||||
if !entityIDs[rel.SourceID] {
|
||||
newIDs[rel.SourceID] = true
|
||||
}
|
||||
if !entityIDs[rel.TargetID] {
|
||||
newIDs[rel.TargetID] = true
|
||||
}
|
||||
}
|
||||
|
||||
if len(newIDs) == 0 {
|
||||
break
|
||||
}
|
||||
|
||||
ids2 := make([]interface{}, 0, len(newIDs))
|
||||
for id := range newIDs {
|
||||
ids2 = append(ids2, id)
|
||||
}
|
||||
|
||||
eRows, err := g.db.Query(
|
||||
fmt.Sprintf(
|
||||
`SELECT id, name, type, mention_count, created_at, updated_at
|
||||
FROM entities WHERE id IN (%s)`, placeholders(len(ids2))),
|
||||
ids2...,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer eRows.Close()
|
||||
|
||||
for eRows.Next() {
|
||||
var e Entity
|
||||
if err := eRows.Scan(&e.ID, &e.Name, &e.Type, &e.MentionCount, &e.CreatedAt, &e.UpdatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !entityIDs[e.ID] {
|
||||
entityIDs[e.ID] = true
|
||||
result.Entities = append(result.Entities, e)
|
||||
}
|
||||
}
|
||||
|
||||
for id := range newIDs {
|
||||
entityIDs[id] = true
|
||||
}
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (g *GraphDB) Purge(criteria map[string]string, mode string) (int, error) {
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
|
||||
conds := []string{"r.status = 'active'"}
|
||||
args := []interface{}{}
|
||||
|
||||
if v, ok := criteria["subject_contains"]; ok {
|
||||
rows, err := g.db.Query("SELECT id FROM entities WHERE name LIKE ?", "%"+v+"%")
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var ids []interface{}
|
||||
for rows.Next() {
|
||||
var id int64
|
||||
rows.Scan(&id)
|
||||
ids = append(ids, id)
|
||||
}
|
||||
if len(ids) > 0 {
|
||||
conds = append(conds, fmt.Sprintf("r.source_id IN (%s)", placeholders(len(ids))))
|
||||
args = append(args, ids...)
|
||||
}
|
||||
}
|
||||
|
||||
if v, ok := criteria["target_contains"]; ok {
|
||||
rows, err := g.db.Query("SELECT id FROM entities WHERE name LIKE ?", "%"+v+"%")
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var ids []interface{}
|
||||
for rows.Next() {
|
||||
var id int64
|
||||
rows.Scan(&id)
|
||||
ids = append(ids, id)
|
||||
}
|
||||
if len(ids) > 0 {
|
||||
conds = append(conds, fmt.Sprintf("r.target_id IN (%s)", placeholders(len(ids))))
|
||||
args = append(args, ids...)
|
||||
}
|
||||
}
|
||||
|
||||
if v, ok := criteria["relation_type"]; ok {
|
||||
conds = append(conds, "r.relation_type = ?")
|
||||
args = append(args, v)
|
||||
}
|
||||
|
||||
if v, ok := criteria["session_id"]; ok {
|
||||
conds = append(conds, "r.session_id = ?")
|
||||
args = append(args, v)
|
||||
}
|
||||
|
||||
if len(conds) == 1 {
|
||||
return 0, fmt.Errorf("no criteria provided")
|
||||
}
|
||||
|
||||
where := ""
|
||||
for i, c := range conds {
|
||||
if i == 0 {
|
||||
where = c
|
||||
} else {
|
||||
where += " AND " + c
|
||||
}
|
||||
}
|
||||
|
||||
if mode == "hard" {
|
||||
result, err := g.db.Exec(
|
||||
fmt.Sprintf(`DELETE FROM relations WHERE %s`, where), args...)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
n, _ := result.RowsAffected()
|
||||
|
||||
g.db.Exec(`DELETE FROM entities WHERE id NOT IN (
|
||||
SELECT DISTINCT source_id FROM relations
|
||||
UNION SELECT DISTINCT target_id FROM relations)`)
|
||||
|
||||
return int(n), nil
|
||||
}
|
||||
|
||||
result, err := g.db.Exec(
|
||||
fmt.Sprintf(`UPDATE relations SET status = 'deleted', updated_at = CURRENT_TIMESTAMP WHERE %s`, where),
|
||||
args...,
|
||||
)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
n, _ := result.RowsAffected()
|
||||
return int(n), nil
|
||||
}
|
||||
|
||||
func (g *GraphDB) Introspect() (map[string]interface{}, error) {
|
||||
g.mu.RLock()
|
||||
defer g.mu.RUnlock()
|
||||
|
||||
var entityCount, relationCount int
|
||||
g.db.QueryRow("SELECT COUNT(*) FROM entities").Scan(&entityCount)
|
||||
g.db.QueryRow("SELECT COUNT(*) FROM relations WHERE status = 'active'").Scan(&relationCount)
|
||||
|
||||
hotspots := []map[string]interface{}{}
|
||||
rows, err := g.db.Query(
|
||||
`SELECT name, mention_count, type FROM entities ORDER BY mention_count DESC LIMIT 10`,
|
||||
)
|
||||
if err == nil {
|
||||
defer rows.Close()
|
||||
for rows.Next() {
|
||||
var name, etype string
|
||||
var count int
|
||||
if err := rows.Scan(&name, &count, &etype); err == nil {
|
||||
hotspots = append(hotspots, map[string]interface{}{
|
||||
"name": name, "count": count, "type": etype,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return map[string]interface{}{
|
||||
"entity_count": entityCount,
|
||||
"relation_count": relationCount,
|
||||
"memory_hotspots": hotspots,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (g *GraphDB) Archive(days int) (int, error) {
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
|
||||
result, err := g.db.Exec(
|
||||
`UPDATE relations SET status = 'archived', updated_at = CURRENT_TIMESTAMP
|
||||
WHERE status = 'active' AND created_at < datetime('now', ?)`,
|
||||
fmt.Sprintf("-%d days", days),
|
||||
)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
n, _ := result.RowsAffected()
|
||||
return int(n), nil
|
||||
}
|
||||
|
||||
func (g *GraphDB) Close() error {
|
||||
return g.db.Close()
|
||||
}
|
||||
|
||||
func placeholders(n int) string {
|
||||
if n <= 0 {
|
||||
return "NULL"
|
||||
}
|
||||
b := make([]byte, 0, n*2-1)
|
||||
for i := 0; i < n; i++ {
|
||||
if i > 0 {
|
||||
b = append(b, ',')
|
||||
}
|
||||
b = append(b, '?')
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
329
internal/memory/indexer.go
Normal file
329
internal/memory/indexer.go
Normal file
@ -0,0 +1,329 @@
|
||||
package memory
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"log"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"gitcode.com/JianFeeeee/HomeAgent/internal/memory/vector"
|
||||
)
|
||||
|
||||
type Indexer struct {
|
||||
db *GraphDB
|
||||
vec *vector.Store
|
||||
veczer *vector.TFIDFVectorizer
|
||||
mu sync.RWMutex
|
||||
trained bool
|
||||
}
|
||||
|
||||
func NewIndexer(db *GraphDB) *Indexer {
|
||||
return &Indexer{
|
||||
db: db,
|
||||
vec: vector.NewStore(),
|
||||
veczer: vector.NewTFIDFVectorizer(2),
|
||||
}
|
||||
}
|
||||
|
||||
// Sync 从图数据库中同步实体名到向量索引
|
||||
func (idx *Indexer) Sync() error {
|
||||
idx.mu.Lock()
|
||||
defer idx.mu.Unlock()
|
||||
|
||||
if idx.db == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
result, err := idx.db.Recall(nil, nil, 1, "")
|
||||
if err != nil || result == nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// 收集实体名
|
||||
var names []string
|
||||
for _, e := range result.Entities {
|
||||
names = append(names, e.Name)
|
||||
}
|
||||
|
||||
if len(names) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// 训练向量化器
|
||||
idx.veczer.Train(names)
|
||||
|
||||
// 重建向量索引
|
||||
idx.vec = vector.NewStore()
|
||||
for _, e := range result.Entities {
|
||||
vec := idx.veczer.Vectorize(e.Name)
|
||||
idx.vec.Insert(fmt.Sprintf("entity_%d", e.ID), e.Name, vec, map[string]string{
|
||||
"type": "entity",
|
||||
"name": e.Name,
|
||||
})
|
||||
}
|
||||
|
||||
idx.trained = true
|
||||
log.Printf("[indexer] synced %d entities to vector index", len(names))
|
||||
return nil
|
||||
}
|
||||
|
||||
type InjectedContext struct {
|
||||
Entities []Entity `json:"entities"`
|
||||
Relations []Relation `json:"relations"`
|
||||
Summary string `json:"summary"`
|
||||
TokenEstimate int `json:"token_estimate"`
|
||||
}
|
||||
|
||||
func (idx *Indexer) BuildContext(userInput string) *InjectedContext {
|
||||
if idx.db == nil {
|
||||
return &InjectedContext{Summary: ""}
|
||||
}
|
||||
|
||||
// 1. 向量搜索：从实体名向量索引中找到相关实体
|
||||
vectorEntities := idx.vectorSearchEntities(userInput)
|
||||
|
||||
// 2. 关键词搜索：已有逻辑
|
||||
keywords := extractKeywords(userInput)
|
||||
if len(keywords) == 0 && len(vectorEntities) == 0 {
|
||||
keywords = []string{userInput}
|
||||
}
|
||||
|
||||
// 合并关键词和向量找到的实体名
|
||||
seedNames := make([]string, 0, len(vectorEntities))
|
||||
for _, e := range vectorEntities {
|
||||
seedNames = append(seedNames, e.Name)
|
||||
}
|
||||
allKeywords := append(keywords, seedNames...)
|
||||
|
||||
result, err := idx.db.Recall(allKeywords, nil, 2, "")
|
||||
if err != nil || result == nil {
|
||||
return &InjectedContext{Summary: ""}
|
||||
}
|
||||
|
||||
ctx := &InjectedContext{
|
||||
Entities: result.Entities,
|
||||
Relations: nil,
|
||||
}
|
||||
|
||||
if len(result.Entities) > 0 {
|
||||
summary := buildIndexSummary(result.Entities)
|
||||
ctx.Summary = summary
|
||||
ctx.TokenEstimate = estimateTokens(summary) + len(result.Entities)*8
|
||||
} else {
|
||||
ctx.Summary = ""
|
||||
}
|
||||
|
||||
return ctx
|
||||
}
|
||||
|
||||
// vectorSearchEntities 在实体名向量索引中搜索
|
||||
func (idx *Indexer) vectorSearchEntities(query string) []Entity {
|
||||
idx.mu.RLock()
|
||||
defer idx.mu.RUnlock()
|
||||
|
||||
if !idx.trained || idx.vec.Size() == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
queryVec := idx.veczer.Vectorize(query)
|
||||
results := idx.vec.Search(queryVec, 5)
|
||||
|
||||
var entities []Entity
|
||||
for _, r := range results {
|
||||
if r.Meta != nil && r.Meta["type"] == "entity" {
|
||||
entities = append(entities, Entity{Name: r.Meta["name"]})
|
||||
}
|
||||
}
|
||||
return entities
|
||||
}
|
||||
|
||||
func (idx *Indexer) BuildToolPrompt() string {
|
||||
return `## 图记忆工具
|
||||
|
||||
你有以下工具可以操作长期图记忆系统：
|
||||
|
||||
### memory_recall
|
||||
检索与关键词相关的实体和关系。
|
||||
参数:
|
||||
- query_intent: 查询关键词，逗号分隔
|
||||
- depth: 遍历深度（默认2）
|
||||
|
||||
### memory_commit
|
||||
将三元组写入图记忆。
|
||||
参数:
|
||||
- triples: [{"subject": "实体名", "relation": "关系类型", "object": "目标实体"}]
|
||||
|
||||
### memory_introspect
|
||||
查看记忆统计信息。
|
||||
|
||||
### memory_purge
|
||||
删除或修正记忆。
|
||||
参数:
|
||||
- criteria: {"subject_contains": "...", "relation_type": "..."}
|
||||
- mode: "soft" | "supersede"
|
||||
|
||||
使用方法：在推理过程中调用对应的 tool，系统会自动执行并返回结果。`
|
||||
}
|
||||
|
||||
func (idx *Indexer) FormatContext(ctx *InjectedContext) string {
|
||||
if ctx == nil || len(ctx.Entities) == 0 {
|
||||
return ""
|
||||
}
|
||||
|
||||
var b strings.Builder
|
||||
b.WriteString("【记忆索引】")
|
||||
|
||||
if ctx.Summary != "" {
|
||||
b.WriteString(" ")
|
||||
b.WriteString(ctx.Summary)
|
||||
}
|
||||
|
||||
b.WriteString(fmt.Sprintf(" 索引: "))
|
||||
for i, e := range ctx.Entities {
|
||||
if i >= 5 {
|
||||
b.WriteString("…")
|
||||
break
|
||||
}
|
||||
if i > 0 {
|
||||
b.WriteString(", ")
|
||||
}
|
||||
b.WriteString(e.Name)
|
||||
if e.Type != "Concept" {
|
||||
b.WriteString("(" + e.Type + ")")
|
||||
}
|
||||
}
|
||||
|
||||
b.WriteString(" | 需更多细节请用 memory_recall 查询")
|
||||
return b.String()
|
||||
}
|
||||
|
||||
func (idx *Indexer) GetToolDefinitions() []map[string]interface{} {
|
||||
return []map[string]interface{}{
|
||||
{
|
||||
"type": "function",
|
||||
"function": map[string]interface{}{
|
||||
"name": "memory_recall",
|
||||
"description": "检索图记忆。输入查询意图关键词，返回相关实体和关系。",
|
||||
"parameters": map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"query_intent": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "查询意图，支持逗号分隔多个关键词",
|
||||
},
|
||||
"depth": map[string]interface{}{
|
||||
"type": "integer",
|
||||
"description": "遍历深度，默认2",
|
||||
"default": 2,
|
||||
},
|
||||
},
|
||||
"required": []string{"query_intent"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": map[string]interface{}{
|
||||
"name": "memory_commit",
|
||||
"description": "写入图记忆。将三元组列表写入长期记忆。",
|
||||
"parameters": map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"triples": map[string]interface{}{
|
||||
"type": "array",
|
||||
"description": "三元组列表",
|
||||
"items": map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"subject": map[string]interface{}{"type": "string"},
|
||||
"relation": map[string]interface{}{"type": "string"},
|
||||
"object": map[string]interface{}{"type": "string"},
|
||||
},
|
||||
"required": []string{"subject", "relation", "object"},
|
||||
},
|
||||
},
|
||||
},
|
||||
"required": []string{"triples"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": map[string]interface{}{
|
||||
"name": "memory_introspect",
|
||||
"description": "查看图记忆统计信息：实体数量、关系数量、热点实体。",
|
||||
"parameters": map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func extractKeywords(input string) []string {
|
||||
stopWords := map[string]bool{
|
||||
"的": true, "了": true, "是": true, "在": true, "有": true,
|
||||
"和": true, "就": true, "不": true, "人": true, "都": true,
|
||||
"一": true, "一个": true, "上": true, "也": true, "很": true,
|
||||
"到": true, "说": true, "要": true, "去": true, "你": true,
|
||||
"会": true, "着": true, "没有": true, "看": true, "好": true,
|
||||
"自己": true, "这": true, "他": true, "她": true, "它": true,
|
||||
"什么": true, "怎么": true, "为什么": true, "如何": true,
|
||||
}
|
||||
|
||||
var keywords []string
|
||||
seen := make(map[string]bool)
|
||||
|
||||
runes := []rune(input)
|
||||
|
||||
bigram := []rune{}
|
||||
for _, r := range runes {
|
||||
bigram = append(bigram, r)
|
||||
if len(bigram) >= 2 {
|
||||
word := string(bigram)
|
||||
if !stopWords[word] && !seen[word] {
|
||||
seen[word] = true
|
||||
keywords = append(keywords, word)
|
||||
}
|
||||
bigram = bigram[1:]
|
||||
}
|
||||
}
|
||||
|
||||
if len(keywords) == 0 && len(runes) > 0 {
|
||||
keywords = []string{string(runes)}
|
||||
}
|
||||
|
||||
if len(keywords) > 5 {
|
||||
keywords = keywords[:5]
|
||||
}
|
||||
|
||||
return keywords
|
||||
}
|
||||
|
||||
func buildIndexSummary(entities []Entity) string {
|
||||
if len(entities) == 0 {
|
||||
return ""
|
||||
}
|
||||
|
||||
var b strings.Builder
|
||||
b.WriteString(fmt.Sprintf("关联 %d 个记忆实体", len(entities)))
|
||||
|
||||
topN := 3
|
||||
if len(entities) < topN {
|
||||
topN = len(entities)
|
||||
}
|
||||
b.WriteString("，高频：")
|
||||
for i := 0; i < topN; i++ {
|
||||
if i > 0 {
|
||||
b.WriteString("、")
|
||||
}
|
||||
b.WriteString(entities[i].Name)
|
||||
}
|
||||
|
||||
return b.String()
|
||||
}
|
||||
|
||||
func estimateTokens(s string) int {
|
||||
return len(s) / 2
|
||||
}
|
||||
291
internal/memory/pipeline/pipeline.go
Normal file
291
internal/memory/pipeline/pipeline.go
Normal file
@ -0,0 +1,291 @@
|
||||
package pipeline
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"gitcode.com/JianFeeeee/HomeAgent/internal/memory"
|
||||
)
|
||||
|
||||
type RawRecord struct {
|
||||
ID int64 `json:"id"`
|
||||
SessionID string `json:"session_id"`
|
||||
Role string `json:"role"`
|
||||
Content string `json:"content"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
Distilled bool `json:"distilled"`
|
||||
}
|
||||
|
||||
type DistillerConfig struct {
|
||||
Interval time.Duration `json:"interval"`
|
||||
RetentionDays int `json:"retention_days"`
|
||||
BatchSize int `json:"batch_size"`
|
||||
}
|
||||
|
||||
type Distiller struct {
|
||||
mu sync.Mutex
|
||||
db *memory.GraphDB
|
||||
rawPath string
|
||||
records []RawRecord
|
||||
nextID int64
|
||||
cfg DistillerConfig
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
onMemory func(input, response string)
|
||||
}
|
||||
|
||||
func NewDistiller(db *memory.GraphDB, dataDir string, cfg DistillerConfig) *Distiller {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
return &Distiller{
|
||||
db: db,
|
||||
rawPath: filepath.Join(dataDir, "memory", "raw"),
|
||||
cfg: cfg,
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
}
|
||||
}
|
||||
|
||||
func (d *Distiller) OnMemoryCandidate(fn func(input, response string)) {
|
||||
d.onMemory = fn
|
||||
}
|
||||
|
||||
func (d *Distiller) Start() {
|
||||
if err := os.MkdirAll(d.rawPath, 0755); err != nil {
|
||||
log.Printf("[memory] create raw path: %v", err)
|
||||
}
|
||||
d.loadExisting()
|
||||
log.Printf("[memory] distiller started (interval: %v, retention: %d days)", d.cfg.Interval, d.cfg.RetentionDays)
|
||||
go d.distillLoop()
|
||||
}
|
||||
|
||||
func (d *Distiller) Stop() {
|
||||
d.cancel()
|
||||
d.flush()
|
||||
}
|
||||
|
||||
func (d *Distiller) Append(sessionID string, role string, content string) {
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
d.nextID++
|
||||
d.records = append(d.records, RawRecord{
|
||||
ID: d.nextID, SessionID: sessionID, Role: role,
|
||||
Content: content, CreatedAt: time.Now(),
|
||||
})
|
||||
}
|
||||
|
||||
func (d *Distiller) flush() {
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
if len(d.records) == 0 {
|
||||
return
|
||||
}
|
||||
path := filepath.Join(d.rawPath, fmt.Sprintf("raw_%d.jsonl", time.Now().UnixNano()))
|
||||
f, err := os.Create(path)
|
||||
if err != nil {
|
||||
log.Printf("[memory] flush error: %v", err)
|
||||
return
|
||||
}
|
||||
defer f.Close()
|
||||
for _, r := range d.records {
|
||||
line := fmt.Sprintf("%d\t%s\t%s\t%s\t%d\n", r.ID, r.SessionID, r.Role, r.Content, r.CreatedAt.Unix())
|
||||
f.WriteString(line)
|
||||
}
|
||||
}
|
||||
|
||||
func (d *Distiller) loadExisting() {
|
||||
entries, err := os.ReadDir(d.rawPath)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
for _, entry := range entries {
|
||||
if filepath.Ext(entry.Name()) != ".jsonl" {
|
||||
continue
|
||||
}
|
||||
path := filepath.Join(d.rawPath, entry.Name())
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
for _, line := range parseLines(string(data)) {
|
||||
parts := splitLine(line)
|
||||
if len(parts) >= 4 {
|
||||
d.records = append(d.records, RawRecord{
|
||||
ID: d.nextID, SessionID: parts[1], Role: parts[2], Content: parts[3],
|
||||
})
|
||||
d.nextID++
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (d *Distiller) distillLoop() {
|
||||
ticker := time.NewTicker(d.cfg.Interval)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ticker.C:
|
||||
d.distillOnce()
|
||||
case <-d.ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (d *Distiller) distillOnce() {
|
||||
d.mu.Lock()
|
||||
cutoff := time.Now().AddDate(0, 0, -d.cfg.RetentionDays)
|
||||
var toDistill []RawRecord
|
||||
var remaining []RawRecord
|
||||
for _, r := range d.records {
|
||||
if r.CreatedAt.Before(cutoff) && !r.Distilled {
|
||||
toDistill = append(toDistill, r)
|
||||
} else {
|
||||
remaining = append(remaining, r)
|
||||
}
|
||||
}
|
||||
d.records = remaining
|
||||
d.mu.Unlock()
|
||||
|
||||
if len(toDistill) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
batchSize := d.cfg.BatchSize
|
||||
if batchSize <= 0 {
|
||||
batchSize = 50
|
||||
}
|
||||
for i := 0; i < len(toDistill); i += batchSize {
|
||||
end := i + batchSize
|
||||
if end > len(toDistill) {
|
||||
end = len(toDistill)
|
||||
}
|
||||
d.distillBatch(toDistill[i:end])
|
||||
}
|
||||
d.cleanupRawFiles()
|
||||
log.Printf("[memory] distilled %d records", len(toDistill))
|
||||
}
|
||||
|
||||
func (d *Distiller) distillBatch(batch []RawRecord) {
|
||||
var userContent, assistantContent string
|
||||
sessionIDs := make(map[string]bool)
|
||||
for _, r := range batch {
|
||||
sessionIDs[r.SessionID] = true
|
||||
if r.Role == "user" {
|
||||
userContent += r.Content + " "
|
||||
} else {
|
||||
assistantContent += r.Content + " "
|
||||
}
|
||||
}
|
||||
triples := extractKeyTriples(userContent, assistantContent)
|
||||
if len(triples) > 0 {
|
||||
sessionID := ""
|
||||
for sid := range sessionIDs {
|
||||
sessionID = sid
|
||||
break
|
||||
}
|
||||
if _, _, err := d.db.Commit(triples, sessionID, 0); err != nil {
|
||||
log.Printf("[memory] distill commit: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (d *Distiller) cleanupRawFiles() {
|
||||
entries, err := os.ReadDir(d.rawPath)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
cutoff := time.Now().AddDate(0, 0, -(d.cfg.RetentionDays + 1))
|
||||
for _, entry := range entries {
|
||||
info, err := entry.Info()
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
if info.ModTime().Before(cutoff) {
|
||||
os.Remove(filepath.Join(d.rawPath, entry.Name()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func extractKeyTriples(userContent, assistantContent string) []memory.Triple {
|
||||
var triples []memory.Triple
|
||||
if len(userContent) > 0 && len(userContent) < 500 {
|
||||
triples = append(triples, memory.Triple{Subject: "用户", Relation: "提及", Object: truncate(userContent, 200)})
|
||||
}
|
||||
if len(assistantContent) > 0 && len(assistantContent) < 500 {
|
||||
triples = append(triples, memory.Triple{Subject: "AI", Relation: "回应", Object: truncate(assistantContent, 200)})
|
||||
}
|
||||
return triples
|
||||
}
|
||||
|
||||
func truncate(s string, max int) string {
|
||||
if len(s) > max {
|
||||
return s[:max] + "..."
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func parseLines(data string) []string {
|
||||
var lines []string
|
||||
current := ""
|
||||
for _, ch := range data {
|
||||
if ch == '\n' {
|
||||
if current != "" {
|
||||
lines = append(lines, current)
|
||||
}
|
||||
current = ""
|
||||
} else {
|
||||
current += string(ch)
|
||||
}
|
||||
}
|
||||
if current != "" {
|
||||
lines = append(lines, current)
|
||||
}
|
||||
return lines
|
||||
}
|
||||
|
||||
func splitLine(line string) []string {
|
||||
var parts []string
|
||||
current := ""
|
||||
for _, ch := range line {
|
||||
if ch == '\t' {
|
||||
parts = append(parts, current)
|
||||
current = ""
|
||||
} else {
|
||||
current += string(ch)
|
||||
}
|
||||
}
|
||||
if current != "" {
|
||||
parts = append(parts, current)
|
||||
}
|
||||
return parts
|
||||
}
|
||||
|
||||
func (d *Distiller) GetRecentRecords(limit int) []RawRecord {
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
n := len(d.records)
|
||||
if n == 0 {
|
||||
return nil
|
||||
}
|
||||
if limit > 0 && limit < n {
|
||||
n = limit
|
||||
}
|
||||
result := make([]RawRecord, n)
|
||||
copy(result, d.records[len(d.records)-n:])
|
||||
return result
|
||||
}
|
||||
|
||||
func (d *Distiller) Stats() map[string]interface{} {
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
return map[string]interface{}{
|
||||
"raw_records": len(d.records),
|
||||
"interval": d.cfg.Interval.String(),
|
||||
"retention_days": d.cfg.RetentionDays,
|
||||
}
|
||||
}
|
||||
241
internal/memory/text/text.go
Normal file
241
internal/memory/text/text.go
Normal file
@ -0,0 +1,241 @@
|
||||
package text
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Event — 原始 I/O 事件记录，写入 JSONL
|
||||
type Event struct {
|
||||
Timestamp int64 `json:"ts"`
|
||||
Source string `json:"source"`
|
||||
Input string `json:"input"`
|
||||
Response string `json:"response,omitempty"`
|
||||
ToolsUsed []string `json:"tools_used,omitempty"`
|
||||
AgentID string `json:"agent_id,omitempty"`
|
||||
}
|
||||
|
||||
// Memory — 文本记忆：追加写 JSONL，按时间/大小旋转
|
||||
type Memory struct {
|
||||
dir string
|
||||
interval time.Duration
|
||||
maxSize int64
|
||||
|
||||
mu sync.Mutex
|
||||
current *os.File
|
||||
encoder *json.Encoder
|
||||
created time.Time
|
||||
size int64
|
||||
|
||||
stopCh chan struct{}
|
||||
}
|
||||
|
||||
type Option func(*Memory)
|
||||
|
||||
func WithRotationInterval(d time.Duration) Option {
|
||||
return func(m *Memory) { m.interval = d }
|
||||
}
|
||||
|
||||
func WithMaxSizeBytes(n int64) Option {
|
||||
return func(m *Memory) { m.maxSize = n }
|
||||
}
|
||||
|
||||
func New(dir string, opts ...Option) *Memory {
|
||||
m := &Memory{
|
||||
dir: dir,
|
||||
interval: 24 * time.Hour,
|
||||
maxSize: 10 * 1024 * 1024,
|
||||
stopCh: make(chan struct{}),
|
||||
}
|
||||
for _, opt := range opts {
|
||||
opt(m)
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
func (m *Memory) Start() error {
|
||||
if err := os.MkdirAll(m.dir, 0755); err != nil {
|
||||
return fmt.Errorf("text memory dir: %w", err)
|
||||
}
|
||||
if err := m.openCurrent(); err != nil {
|
||||
return err
|
||||
}
|
||||
go m.rotationLoop()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *Memory) Stop() {
|
||||
close(m.stopCh)
|
||||
m.mu.Lock()
|
||||
if m.current != nil {
|
||||
m.current.Close()
|
||||
}
|
||||
m.mu.Unlock()
|
||||
}
|
||||
|
||||
func (m *Memory) Append(evt Event) error {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
if m.needRotate() {
|
||||
m.rotateLocked()
|
||||
}
|
||||
|
||||
if err := m.encoder.Encode(evt); err != nil {
|
||||
return fmt.Errorf("encode event: %w", err)
|
||||
}
|
||||
m.current.Sync()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *Memory) needRotate() bool {
|
||||
return time.Since(m.created) > m.interval || m.size > m.maxSize
|
||||
}
|
||||
|
||||
func (m *Memory) rotateLocked() {
|
||||
if m.current != nil {
|
||||
m.current.Close()
|
||||
}
|
||||
m.openCurrent()
|
||||
}
|
||||
|
||||
func (m *Memory) openCurrent() error {
|
||||
now := time.Now()
|
||||
name := fmt.Sprintf("text_%s.jsonl", now.Format("2006-01-02_15-04-05"))
|
||||
path := filepath.Join(m.dir, name)
|
||||
|
||||
f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
|
||||
if err != nil {
|
||||
return fmt.Errorf("open text log %s: %w", path, err)
|
||||
}
|
||||
|
||||
stat, _ := f.Stat()
|
||||
m.current = f
|
||||
m.encoder = json.NewEncoder(f)
|
||||
m.created = now
|
||||
m.size = stat.Size()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *Memory) rotationLoop() {
|
||||
ticker := time.NewTicker(m.interval / 2)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ticker.C:
|
||||
m.mu.Lock()
|
||||
if m.needRotate() {
|
||||
m.rotateLocked()
|
||||
log.Printf("[text memory] rotated log file")
|
||||
}
|
||||
m.mu.Unlock()
|
||||
case <-m.stopCh:
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Replay — 从 JSONL 文件流式回放事件
|
||||
func (m *Memory) Replay(fn func(Event) error) error {
|
||||
m.mu.Lock()
|
||||
files, err := m.listFiles()
|
||||
m.mu.Unlock()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
for _, fpath := range files {
|
||||
if err := m.replayFile(fpath, fn); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *Memory) replayFile(path string, fn func(Event) error) error {
|
||||
f, err := os.Open(path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
scanner := bufio.NewScanner(f)
|
||||
scanner.Buffer(make([]byte, 1024*1024), 1024*1024)
|
||||
for scanner.Scan() {
|
||||
line := strings.TrimSpace(scanner.Text())
|
||||
if line == "" {
|
||||
continue
|
||||
}
|
||||
var evt Event
|
||||
if err := json.Unmarshal([]byte(line), &evt); err != nil {
|
||||
continue
|
||||
}
|
||||
if err := fn(evt); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return scanner.Err()
|
||||
}
|
||||
|
||||
func (m *Memory) listFiles() ([]string, error) {
|
||||
entries, err := os.ReadDir(m.dir)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var files []string
|
||||
for _, e := range entries {
|
||||
if strings.HasPrefix(e.Name(), "text_") && strings.HasSuffix(e.Name(), ".jsonl") {
|
||||
files = append(files, filepath.Join(m.dir, e.Name()))
|
||||
}
|
||||
}
|
||||
sort.Strings(files)
|
||||
return files, nil
|
||||
}
|
||||
|
||||
// RecentEvents — 返回最近 n 条事件（跨所有文件的最新事件）
|
||||
func (m *Memory) RecentEvents(n int) ([]Event, error) {
|
||||
var all []Event
|
||||
err := m.Replay(func(evt Event) error {
|
||||
all = append(all, evt)
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(all) > n {
|
||||
all = all[len(all)-n:]
|
||||
}
|
||||
return all, nil
|
||||
}
|
||||
|
||||
func (m *Memory) FileCount() int {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
files, err := m.listFiles()
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return len(files)
|
||||
}
|
||||
|
||||
func (m *Memory) Stats() map[string]interface{} {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
files, _ := m.listFiles()
|
||||
return map[string]interface{}{
|
||||
"file_count": len(files),
|
||||
"current_size": m.size,
|
||||
"rotation_bytes": m.maxSize,
|
||||
"rotation_interval": m.interval.String(),
|
||||
"dir": m.dir,
|
||||
}
|
||||
}
|
||||
302
internal/memory/vector/store.go
Normal file
302
internal/memory/vector/store.go
Normal file
@ -0,0 +1,302 @@
|
||||
package vector
|
||||
|
||||
import (
|
||||
"math"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// Vectorizer 接口：将文本转为向量
|
||||
type Vectorizer interface {
|
||||
Vectorize(text string) Vector
|
||||
}
|
||||
|
||||
// Vector 是带权特征映射：feature → weight
|
||||
type Vector map[string]float64
|
||||
|
||||
// Store 向量存储，支持近似查询
|
||||
type Store struct {
|
||||
mu sync.RWMutex
|
||||
docs []DocVector
|
||||
dim int
|
||||
index *InvertedIndex
|
||||
}
|
||||
|
||||
type DocVector struct {
|
||||
ID string
|
||||
Vector Vector
|
||||
Text string
|
||||
Meta map[string]string
|
||||
}
|
||||
|
||||
func NewStore() *Store {
|
||||
return &Store{
|
||||
index: NewInvertedIndex(),
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Store) Insert(id, text string, vec Vector, meta map[string]string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
s.docs = append(s.docs, DocVector{
|
||||
ID: id, Vector: vec, Text: text, Meta: meta,
|
||||
})
|
||||
s.index.Add(id, vec)
|
||||
}
|
||||
|
||||
func (s *Store) Remove(id string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
filtered := make([]DocVector, 0, len(s.docs))
|
||||
for _, d := range s.docs {
|
||||
if d.ID != id {
|
||||
filtered = append(filtered, d)
|
||||
}
|
||||
}
|
||||
s.docs = filtered
|
||||
s.index.Remove(id)
|
||||
}
|
||||
|
||||
func (s *Store) Search(query Vector, topK int) []DocVector {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
|
||||
if len(s.docs) == 0 || len(query) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
candidates := s.index.Search(query, len(s.docs))
|
||||
type scored struct {
|
||||
doc DocVector
|
||||
score float64
|
||||
}
|
||||
|
||||
var results []scored
|
||||
seen := make(map[string]bool)
|
||||
for _, id := range candidates {
|
||||
if seen[id] {
|
||||
continue
|
||||
}
|
||||
seen[id] = true
|
||||
for _, d := range s.docs {
|
||||
if d.ID == id {
|
||||
score := CosineSimilarity(query, d.Vector)
|
||||
if score > 0 {
|
||||
results = append(results, scored{d, score})
|
||||
}
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
sort.Slice(results, func(i, j int) bool {
|
||||
return results[i].score > results[j].score
|
||||
})
|
||||
|
||||
if len(results) > topK {
|
||||
results = results[:topK]
|
||||
}
|
||||
|
||||
out := make([]DocVector, len(results))
|
||||
for i, r := range results {
|
||||
out[i] = r.doc
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (s *Store) Size() int {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return len(s.docs)
|
||||
}
|
||||
|
||||
func (s *Store) All() []DocVector {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
out := make([]DocVector, len(s.docs))
|
||||
copy(out, s.docs)
|
||||
return out
|
||||
}
|
||||
|
||||
// TFIDFVectorizer 使用字符 bigram + TF-IDF
|
||||
type TFIDFVectorizer struct {
|
||||
mu sync.RWMutex
|
||||
docFreq map[string]float64 // feature → 文档频率
|
||||
totalDocs int
|
||||
maxNGram int
|
||||
}
|
||||
|
||||
func NewTFIDFVectorizer(maxNGram int) *TFIDFVectorizer {
|
||||
if maxNGram <= 0 {
|
||||
maxNGram = 2
|
||||
}
|
||||
return &TFIDFVectorizer{
|
||||
docFreq: make(map[string]float64),
|
||||
maxNGram: maxNGram,
|
||||
}
|
||||
}
|
||||
|
||||
func (v *TFIDFVectorizer) Train(docs []string) {
|
||||
v.mu.Lock()
|
||||
defer v.mu.Unlock()
|
||||
|
||||
v.docFreq = make(map[string]float64)
|
||||
v.totalDocs = len(docs)
|
||||
|
||||
seen := make(map[string]map[string]bool)
|
||||
for _, doc := range docs {
|
||||
features := extractNGrams(doc, v.maxNGram)
|
||||
key := doc
|
||||
if seen[key] == nil {
|
||||
seen[key] = make(map[string]bool)
|
||||
}
|
||||
for _, f := range features {
|
||||
if !seen[key][f] {
|
||||
seen[key][f] = true
|
||||
v.docFreq[f]++
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (v *TFIDFVectorizer) Vectorize(text string) Vector {
|
||||
v.mu.RLock()
|
||||
defer v.mu.RUnlock()
|
||||
|
||||
features := extractNGrams(text, v.maxNGram)
|
||||
tf := make(map[string]float64)
|
||||
for _, f := range features {
|
||||
tf[f]++
|
||||
}
|
||||
maxTF := 0.0
|
||||
for _, c := range tf {
|
||||
if c > maxTF {
|
||||
maxTF = c
|
||||
}
|
||||
}
|
||||
|
||||
vec := make(Vector)
|
||||
for f, count := range tf {
|
||||
tfNorm := count / maxTF
|
||||
idf := 1.0
|
||||
if v.totalDocs > 0 {
|
||||
df := v.docFreq[f]
|
||||
if df > 0 {
|
||||
idf = math.Log(float64(v.totalDocs+1)/df+1) + 1
|
||||
}
|
||||
}
|
||||
vec[f] = tfNorm * idf
|
||||
}
|
||||
return vec
|
||||
}
|
||||
|
||||
// extractNGrams 提取 n-gram 特征（主要用于中文）
|
||||
func extractNGrams(text string, maxN int) []string {
|
||||
runes := []rune(strings.ToLower(text))
|
||||
var features []string
|
||||
seen := make(map[string]bool)
|
||||
|
||||
for n := 1; n <= maxN; n++ {
|
||||
for i := 0; i <= len(runes)-n; i++ {
|
||||
gram := string(runes[i : i+n])
|
||||
gram = strings.TrimSpace(gram)
|
||||
if gram == "" {
|
||||
continue
|
||||
}
|
||||
if !seen[gram] {
|
||||
seen[gram] = true
|
||||
features = append(features, gram)
|
||||
}
|
||||
}
|
||||
}
|
||||
return features
|
||||
}
|
||||
|
||||
func CosineSimilarity(a, b Vector) float64 {
|
||||
var dot, normA, normB float64
|
||||
for f, va := range a {
|
||||
dot += va * b[f]
|
||||
normA += va * va
|
||||
}
|
||||
for _, vb := range b {
|
||||
normB += vb * vb
|
||||
}
|
||||
if normA == 0 || normB == 0 {
|
||||
return 0
|
||||
}
|
||||
return dot / (math.Sqrt(normA) * math.Sqrt(normB))
|
||||
}
|
||||
|
||||
// InvertedIndex 倒排索引，加速向量搜索
|
||||
type InvertedIndex struct {
|
||||
mu sync.RWMutex
|
||||
postings map[string]map[string]float64 // feature → {docID: weight}
|
||||
}
|
||||
|
||||
func NewInvertedIndex() *InvertedIndex {
|
||||
return &InvertedIndex{
|
||||
postings: make(map[string]map[string]float64),
|
||||
}
|
||||
}
|
||||
|
||||
func (idx *InvertedIndex) Add(docID string, vec Vector) {
|
||||
idx.mu.Lock()
|
||||
defer idx.mu.Unlock()
|
||||
|
||||
for feature, weight := range vec {
|
||||
if idx.postings[feature] == nil {
|
||||
idx.postings[feature] = make(map[string]float64)
|
||||
}
|
||||
idx.postings[feature][docID] = weight
|
||||
}
|
||||
}
|
||||
|
||||
func (idx *InvertedIndex) Remove(docID string) {
|
||||
idx.mu.Lock()
|
||||
defer idx.mu.Unlock()
|
||||
|
||||
for feature, postings := range idx.postings {
|
||||
delete(postings, docID)
|
||||
if len(postings) == 0 {
|
||||
delete(idx.postings, feature)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (idx *InvertedIndex) Search(query Vector, maxResults int) []string {
|
||||
idx.mu.RLock()
|
||||
defer idx.mu.RUnlock()
|
||||
|
||||
scores := make(map[string]float64)
|
||||
for feature, qw := range query {
|
||||
if postings, ok := idx.postings[feature]; ok {
|
||||
for docID, dw := range postings {
|
||||
scores[docID] += qw * dw
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type pair struct {
|
||||
id string
|
||||
score float64
|
||||
}
|
||||
var sorted []pair
|
||||
for id, score := range scores {
|
||||
sorted = append(sorted, pair{id, score})
|
||||
}
|
||||
sort.Slice(sorted, func(i, j int) bool {
|
||||
return sorted[i].score > sorted[j].score
|
||||
})
|
||||
|
||||
if len(sorted) > maxResults {
|
||||
sorted = sorted[:maxResults]
|
||||
}
|
||||
out := make([]string, len(sorted))
|
||||
for i, p := range sorted {
|
||||
out[i] = p.id
|
||||
}
|
||||
return out
|
||||
}
|
||||
Reference in New Issue
Block a user