mirror of
https://gitcode.com/JianFeeeee/HomeAgent.git
synced 2026-09-23 10:28:06 +00:00
refactor: remove IO route mapping, add HTTP API tests, system prompt update
This commit is contained in:
root
@ -216,6 +216,18 @@ func (s *Store) RecentDocs(n int) []*Doc {
|
||||
return list
|
||||
}
|
||||
|
||||
// Remove 从文档存储中删除指定 ID 的文档
|
||||
func (s *Store) Remove(id string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
if _, ok := s.docs[id]; ok {
|
||||
delete(s.docs, id)
|
||||
s.vec.Remove(id)
|
||||
s.dirty = true
|
||||
}
|
||||
}
|
||||
|
||||
// ——— internal ———
|
||||
|
||||
func (s *Store) loadAll() error {
|
||||
|
||||
343
internal/memory/document/document_test.go
Normal file
343
internal/memory/document/document_test.go
Normal file
@ -0,0 +1,343 @@
|
||||
package document
|
||||
|
||||
import (
|
||||
"os"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestInsertAndQuery(t *testing.T) {
|
||||
dir, err := os.MkdirTemp("", "doc_test_*")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer os.RemoveAll(dir)
|
||||
|
||||
s := NewStore(dir)
|
||||
if err := s.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.Stop()
|
||||
|
||||
doc := &Doc{
|
||||
Summary: "用户喜欢喝咖啡",
|
||||
Content: "用户提到他每天早上都会喝一杯黑咖啡",
|
||||
Tags: []string{"咖啡", "习惯"},
|
||||
Source: "manual",
|
||||
}
|
||||
if err := s.Insert(doc); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if doc.ID == "" {
|
||||
t.Error("doc ID should be auto-generated")
|
||||
}
|
||||
|
||||
stats := s.Stats()
|
||||
if stats["doc_count"].(int) != 1 {
|
||||
t.Errorf("expected 1 doc, got %d", stats["doc_count"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestQuery(t *testing.T) {
|
||||
dir, err := os.MkdirTemp("", "doc_query_*")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer os.RemoveAll(dir)
|
||||
|
||||
s := NewStore(dir)
|
||||
s.Start()
|
||||
defer s.Stop()
|
||||
|
||||
s.Insert(&Doc{Summary: "咖啡是一种饮品", Content: "咖啡因提神", Source: "manual"})
|
||||
s.Insert(&Doc{Summary: "茶叶也有咖啡因", Content: "茶和咖啡都提神", Source: "manual"})
|
||||
s.Insert(&Doc{Summary: "今天天气很好", Content: "适合出去散步", Source: "manual"})
|
||||
|
||||
results := s.Query("咖啡", 5)
|
||||
if len(results) == 0 {
|
||||
t.Fatal("expected results for '咖啡'")
|
||||
}
|
||||
|
||||
if results[0].AccessCount <= 0 {
|
||||
t.Error("access count should be updated on query")
|
||||
}
|
||||
}
|
||||
|
||||
func TestContextToDoc(t *testing.T) {
|
||||
dir, err := os.MkdirTemp("", "doc_ctx_*")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer os.RemoveAll(dir)
|
||||
|
||||
s := NewStore(dir)
|
||||
s.Start()
|
||||
defer s.Stop()
|
||||
|
||||
entries := []ContextEntry{
|
||||
{Timestamp: time.Now(), Source: "user", Content: "我喜欢编程", Response: "很好"},
|
||||
{Timestamp: time.Now(), Source: "user", Content: "特别是Go语言", Response: "Go很棒"},
|
||||
}
|
||||
|
||||
doc, err := s.ContextToDoc("test", entries)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if doc == nil {
|
||||
t.Fatal("expected non-nil doc")
|
||||
}
|
||||
if doc.Summary == "" {
|
||||
t.Error("summary should not be empty")
|
||||
}
|
||||
if doc.Content == "" {
|
||||
t.Error("content should not be empty")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFindColdDocs(t *testing.T) {
|
||||
dir, err := os.MkdirTemp("", "doc_cold_*")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer os.RemoveAll(dir)
|
||||
|
||||
s := NewStore(dir)
|
||||
s.Start()
|
||||
defer s.Stop()
|
||||
|
||||
hot := &Doc{Summary: "常用的信息", Content: "经常被查询", Source: "manual"}
|
||||
hot.AccessCount = 10
|
||||
hot.LastAccess = time.Now()
|
||||
s.Insert(hot)
|
||||
|
||||
cold := &Doc{Summary: "很久没用的信息", Content: "几乎不被访问", Source: "manual"}
|
||||
s.Insert(cold)
|
||||
// Insert 会重置 LastAccess，手动改为过去的
|
||||
cold.LastAccess = time.Now().Add(-100 * time.Hour)
|
||||
cold.AccessCount = 1
|
||||
|
||||
// 应该只找到 cold（72h 前未访问且访问 ≤ 2）
|
||||
coldDocs := s.FindColdDocs(72*time.Hour, 2)
|
||||
if len(coldDocs) != 1 {
|
||||
t.Fatalf("expected 1 cold doc, got %d", len(coldDocs))
|
||||
}
|
||||
if coldDocs[0].Summary != "很久没用的信息" {
|
||||
t.Errorf("expected cold doc, got %s", coldDocs[0].Summary)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecentDocs(t *testing.T) {
|
||||
dir, err := os.MkdirTemp("", "doc_recent_*")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer os.RemoveAll(dir)
|
||||
|
||||
s := NewStore(dir)
|
||||
s.Start()
|
||||
defer s.Stop()
|
||||
|
||||
s.Insert(&Doc{Summary: "第一条", Content: "a", Source: "manual"})
|
||||
time.Sleep(time.Millisecond)
|
||||
s.Insert(&Doc{Summary: "第二条", Content: "b", Source: "manual"})
|
||||
|
||||
recent := s.RecentDocs(1)
|
||||
if len(recent) != 1 {
|
||||
t.Fatalf("expected 1 recent doc, got %d", len(recent))
|
||||
}
|
||||
if recent[0].Summary != "第二条" {
|
||||
t.Errorf("expected newest doc, got %s", recent[0].Summary)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReindex(t *testing.T) {
|
||||
dir, err := os.MkdirTemp("", "doc_reindex_*")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer os.RemoveAll(dir)
|
||||
|
||||
s := NewStore(dir)
|
||||
s.Start()
|
||||
defer s.Stop()
|
||||
|
||||
s.Insert(&Doc{Summary: "测试重索引", Content: "验证索引重建", Source: "manual"})
|
||||
s.Reindex()
|
||||
|
||||
results := s.Query("重索引", 5)
|
||||
if len(results) == 0 {
|
||||
t.Error("reindex should preserve searchability")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSummarizeEntries(t *testing.T) {
|
||||
entries := []ContextEntry{
|
||||
{Source: "user", Content: "今天天气如何"},
|
||||
{Source: "user", Content: "明天会下雨吗"},
|
||||
}
|
||||
summary := summarizeEntries(entries)
|
||||
if summary == "" {
|
||||
t.Error("summary should not be empty")
|
||||
}
|
||||
if !contains(summary, "2") {
|
||||
t.Errorf("summary should mention count, got: %s", summary)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractKeywords(t *testing.T) {
|
||||
kws := extractKeywords("今天天气很好")
|
||||
if len(kws) == 0 {
|
||||
t.Error("should extract keywords from Chinese text")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractTags(t *testing.T) {
|
||||
entries := []ContextEntry{
|
||||
{Content: "我喜欢喝咖啡和编程"},
|
||||
}
|
||||
tags := extractTags(entries)
|
||||
if len(tags) == 0 {
|
||||
t.Error("should extract tags")
|
||||
}
|
||||
}
|
||||
|
||||
func TestInsertEmptyDoc(t *testing.T) {
|
||||
dir, err := os.MkdirTemp("", "doc_empty_*")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer os.RemoveAll(dir)
|
||||
|
||||
s := NewStore(dir)
|
||||
s.Start()
|
||||
defer s.Stop()
|
||||
|
||||
doc := &Doc{Summary: "", Content: "", Source: "manual"}
|
||||
if err := s.Insert(doc); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if doc.ID == "" {
|
||||
t.Error("doc ID should be generated even for empty content")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPersistence(t *testing.T) {
|
||||
dir, err := os.MkdirTemp("", "doc_persist_*")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer os.RemoveAll(dir)
|
||||
|
||||
// 写
|
||||
s1 := NewStore(dir)
|
||||
s1.Start()
|
||||
s1.Insert(&Doc{Summary: "持久化测试", Content: "应该被保存到磁盘", Source: "manual"})
|
||||
s1.Stop()
|
||||
|
||||
// 读
|
||||
s2 := NewStore(dir)
|
||||
s2.Start()
|
||||
defer s2.Stop()
|
||||
|
||||
stats := s2.Stats()
|
||||
if stats["doc_count"].(int) != 1 {
|
||||
t.Errorf("expected 1 doc after reload, got %d", stats["doc_count"])
|
||||
}
|
||||
|
||||
results := s2.Query("持久化", 5)
|
||||
if len(results) == 0 {
|
||||
t.Error("search should work after reload")
|
||||
}
|
||||
}
|
||||
|
||||
func contains(s, substr string) bool {
|
||||
for i := 0; i <= len(s)-len(substr); i++ {
|
||||
if s[i:i+len(substr)] == substr {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func TestFlushNoDirty(t *testing.T) {
|
||||
dir, err := os.MkdirTemp("", "doc_flush_*")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer os.RemoveAll(dir)
|
||||
|
||||
s := NewStore(dir)
|
||||
s.Start()
|
||||
|
||||
// 不插任何文档，flush 不应报错
|
||||
s.Stop()
|
||||
}
|
||||
|
||||
func TestRemove(t *testing.T) {
|
||||
dir, err := os.MkdirTemp("", "doc_remove_*")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer os.RemoveAll(dir)
|
||||
|
||||
s := NewStore(dir)
|
||||
s.Start()
|
||||
defer s.Stop()
|
||||
|
||||
s.Insert(&Doc{Summary: "会被删除", Content: "a", Source: "manual"})
|
||||
s.Insert(&Doc{Summary: "会保留", Content: "b", Source: "manual"})
|
||||
|
||||
// 删除前应该有 2 个
|
||||
stats := s.Stats()
|
||||
if stats["doc_count"].(int) != 2 {
|
||||
t.Fatalf("expected 2 docs before remove, got %d", stats["doc_count"])
|
||||
}
|
||||
|
||||
// 遍历找到 "会被删除" 的 ID
|
||||
var rmID string
|
||||
for _, d := range s.docs {
|
||||
if d.Summary == "会被删除" {
|
||||
rmID = d.ID
|
||||
break
|
||||
}
|
||||
}
|
||||
if rmID == "" {
|
||||
t.Fatal("could not find test doc")
|
||||
}
|
||||
|
||||
s.Remove(rmID)
|
||||
|
||||
stats = s.Stats()
|
||||
if stats["doc_count"].(int) != 1 {
|
||||
t.Errorf("expected 1 doc after remove, got %d", stats["doc_count"])
|
||||
}
|
||||
|
||||
// 搜索不应再找到
|
||||
results := s.Query("删除", 5)
|
||||
if len(results) > 0 {
|
||||
t.Error("removed doc should not appear in search results")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRemoveNonexistent(t *testing.T) {
|
||||
dir, err := os.MkdirTemp("", "doc_rm_nonexist_*")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer os.RemoveAll(dir)
|
||||
|
||||
s := NewStore(dir)
|
||||
s.Start()
|
||||
defer s.Stop()
|
||||
|
||||
s.Insert(&Doc{Summary: "一个文档", Content: "x", Source: "manual"})
|
||||
|
||||
// 删除不存在的 ID 不应 panic
|
||||
s.Remove("nonexistent_id")
|
||||
|
||||
stats := s.Stats()
|
||||
if stats["doc_count"].(int) != 1 {
|
||||
t.Errorf("expected 1 doc after remove nonexistent, got %d", stats["doc_count"])
|
||||
}
|
||||
}
|
||||
@ -400,7 +400,7 @@ func (g *GraphDB) Purge(criteria map[string]string, mode string) (int, error) {
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
|
||||
conds := []string{"r.status = 'active'"}
|
||||
conds := []string{"status = 'active'"}
|
||||
args := []interface{}{}
|
||||
|
||||
if v, ok := criteria["subject_contains"]; ok {
|
||||
@ -416,7 +416,7 @@ func (g *GraphDB) Purge(criteria map[string]string, mode string) (int, error) {
|
||||
ids = append(ids, id)
|
||||
}
|
||||
if len(ids) > 0 {
|
||||
conds = append(conds, fmt.Sprintf("r.source_id IN (%s)", placeholders(len(ids))))
|
||||
conds = append(conds, fmt.Sprintf("source_id IN (%s)", placeholders(len(ids))))
|
||||
args = append(args, ids...)
|
||||
}
|
||||
}
|
||||
@ -434,18 +434,18 @@ func (g *GraphDB) Purge(criteria map[string]string, mode string) (int, error) {
|
||||
ids = append(ids, id)
|
||||
}
|
||||
if len(ids) > 0 {
|
||||
conds = append(conds, fmt.Sprintf("r.target_id IN (%s)", placeholders(len(ids))))
|
||||
conds = append(conds, fmt.Sprintf("target_id IN (%s)", placeholders(len(ids))))
|
||||
args = append(args, ids...)
|
||||
}
|
||||
}
|
||||
|
||||
if v, ok := criteria["relation_type"]; ok {
|
||||
conds = append(conds, "r.relation_type = ?")
|
||||
conds = append(conds, "relation_type = ?")
|
||||
args = append(args, v)
|
||||
}
|
||||
|
||||
if v, ok := criteria["session_id"]; ok {
|
||||
conds = append(conds, "r.session_id = ?")
|
||||
conds = append(conds, "session_id = ?")
|
||||
args = append(args, v)
|
||||
}
|
||||
|
||||
|
||||
227
internal/memory/graph_test.go
Normal file
227
internal/memory/graph_test.go
Normal file
@ -0,0 +1,227 @@
|
||||
package memory
|
||||
|
||||
import (
|
||||
"os"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func newTestGraph(t *testing.T) *GraphDB {
|
||||
t.Helper()
|
||||
f, err := os.CreateTemp("", "graph_test_*.db")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
f.Close()
|
||||
os.Remove(f.Name())
|
||||
|
||||
g, err := NewGraphDB(f.Name())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return g
|
||||
}
|
||||
|
||||
func TestNewGraphDB(t *testing.T) {
|
||||
g := newTestGraph(t)
|
||||
defer os.Remove(g.dbPath)
|
||||
defer g.Close()
|
||||
|
||||
stats, err := g.Introspect()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if stats["entity_count"].(int) != 0 {
|
||||
t.Errorf("expected 0 entities, got %d", stats["entity_count"])
|
||||
}
|
||||
if stats["relation_count"].(int) != 0 {
|
||||
t.Errorf("expected 0 relations, got %d", stats["relation_count"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestCommitTriples(t *testing.T) {
|
||||
g := newTestGraph(t)
|
||||
defer os.Remove(g.dbPath)
|
||||
defer g.Close()
|
||||
|
||||
triples := []Triple{
|
||||
{Subject: "张三", Relation: "喜欢", Object: "编程"},
|
||||
{Subject: "张三", Relation: "居住", Object: "北京"},
|
||||
}
|
||||
|
||||
ec, rc, err := g.Commit(triples, "test_session", 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if ec != 4 {
|
||||
t.Errorf("expected 4 entity ops (张三×2, 编程, 北京), got %d", ec)
|
||||
}
|
||||
if rc != 2 {
|
||||
t.Errorf("expected 2 relations, got %d", rc)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCommitEmptyTriples(t *testing.T) {
|
||||
g := newTestGraph(t)
|
||||
defer os.Remove(g.dbPath)
|
||||
defer g.Close()
|
||||
|
||||
ec, rc, err := g.Commit(nil, "test", 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if ec != 0 || rc != 0 {
|
||||
t.Errorf("expected 0,0 for nil triples, got %d,%d", ec, rc)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecallByKeywords(t *testing.T) {
|
||||
g := newTestGraph(t)
|
||||
defer os.Remove(g.dbPath)
|
||||
defer g.Close()
|
||||
|
||||
g.Commit([]Triple{
|
||||
{Subject: "咖啡", Relation: "属于", Object: "饮品"},
|
||||
{Subject: "咖啡", Relation: "含有", Object: "咖啡因"},
|
||||
}, "session1", 0)
|
||||
|
||||
result, err := g.Recall([]string{"咖啡"}, nil, 1, "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(result.Entities) == 0 {
|
||||
t.Error("expected entities for keyword '咖啡'")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecallBySeedEntity(t *testing.T) {
|
||||
g := newTestGraph(t)
|
||||
defer os.Remove(g.dbPath)
|
||||
defer g.Close()
|
||||
|
||||
g.Commit([]Triple{
|
||||
{Subject: "Go", Relation: "是", Object: "编程语言"},
|
||||
{Subject: "Go", Relation: "用于", Object: "后端开发"},
|
||||
}, "session2", 0)
|
||||
|
||||
result, err := g.Recall(nil, []string{"Go"}, 1, "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(result.Entities) == 0 {
|
||||
t.Error("expected entities for seed 'Go'")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecallWithDepth(t *testing.T) {
|
||||
g := newTestGraph(t)
|
||||
defer os.Remove(g.dbPath)
|
||||
defer g.Close()
|
||||
|
||||
g.Commit([]Triple{
|
||||
{Subject: "甲", Relation: "认识", Object: "乙"},
|
||||
{Subject: "乙", Relation: "认识", Object: "丙"},
|
||||
}, "session3", 0)
|
||||
|
||||
result, err := g.Recall(nil, []string{"甲"}, 2, "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(result.Relations) == 0 {
|
||||
t.Error("expected relations with depth search")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPurgeHard(t *testing.T) {
|
||||
g := newTestGraph(t)
|
||||
defer os.Remove(g.dbPath)
|
||||
defer g.Close()
|
||||
|
||||
g.Commit([]Triple{
|
||||
{Subject: "临时", Relation: "用于", Object: "测试"},
|
||||
}, "session4", 0)
|
||||
|
||||
n, err := g.Purge(map[string]string{"subject_contains": "临时"}, "hard")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if n != 1 {
|
||||
t.Errorf("expected 1 purged relation, got %d", n)
|
||||
}
|
||||
|
||||
stats, _ := g.Introspect()
|
||||
if stats["relation_count"].(int) != 0 {
|
||||
t.Errorf("expected 0 relations after purge, got %d", stats["relation_count"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestPurgeSoft(t *testing.T) {
|
||||
g := newTestGraph(t)
|
||||
defer os.Remove(g.dbPath)
|
||||
defer g.Close()
|
||||
|
||||
g.Commit([]Triple{
|
||||
{Subject: "可删除", Relation: "属于", Object: "测试"},
|
||||
}, "session5", 0)
|
||||
|
||||
n, err := g.Purge(map[string]string{"subject_contains": "可删除"}, "soft")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if n != 1 {
|
||||
t.Errorf("expected 1 soft-deleted relation, got %d", n)
|
||||
}
|
||||
|
||||
stats, _ := g.Introspect()
|
||||
if stats["relation_count"].(int) != 0 {
|
||||
t.Errorf("expected 0 active relations after soft-delete, got %d", stats["relation_count"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestArchive(t *testing.T) {
|
||||
g := newTestGraph(t)
|
||||
defer os.Remove(g.dbPath)
|
||||
defer g.Close()
|
||||
|
||||
// 直接插入一条旧记录
|
||||
g.db.Exec(`INSERT INTO entities (id, name, type) VALUES (1, '旧数据', 'Concept')`)
|
||||
g.db.Exec(`INSERT INTO relations (source_id, target_id, relation_type, created_at)
|
||||
VALUES (1, 1, '包含', datetime('now', '-1 day'))`)
|
||||
|
||||
n, err := g.Archive(0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if n != 1 {
|
||||
t.Errorf("expected 1 archived relation, got %d", n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIntrospectHotspots(t *testing.T) {
|
||||
g := newTestGraph(t)
|
||||
defer os.Remove(g.dbPath)
|
||||
defer g.Close()
|
||||
|
||||
g.Commit([]Triple{
|
||||
{Subject: "热门话题", Relation: "关于", Object: "AI"},
|
||||
{Subject: "热门话题", Relation: "关于", Object: "机器学习"},
|
||||
{Subject: "冷门话题", Relation: "关于", Object: "旧技术"},
|
||||
}, "session7", 0)
|
||||
|
||||
stats, _ := g.Introspect()
|
||||
hotspots := stats["memory_hotspots"].([]map[string]interface{})
|
||||
if len(hotspots) == 0 {
|
||||
t.Error("expected hotspots")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPlaceholders(t *testing.T) {
|
||||
if placeholders(0) != "NULL" {
|
||||
t.Errorf("expected NULL for n=0, got %s", placeholders(0))
|
||||
}
|
||||
if placeholders(1) != "?" {
|
||||
t.Errorf("expected '?' for n=1, got %s", placeholders(1))
|
||||
}
|
||||
if placeholders(3) != "?,?,?" {
|
||||
t.Errorf("expected '?,?,?' for n=3, got %s", placeholders(3))
|
||||
}
|
||||
}
|
||||
212
internal/memory/vector/store_test.go
Normal file
212
internal/memory/vector/store_test.go
Normal file
@ -0,0 +1,212 @@
|
||||
package vector
|
||||
|
||||
import (
|
||||
"math"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestExtractNGrams(t *testing.T) {
|
||||
tests := []struct {
|
||||
text string
|
||||
maxN int
|
||||
check []string // 应包含
|
||||
}{
|
||||
{"hello", 2, []string{"h", "e", "l", "o", "he", "el", "ll", "lo"}},
|
||||
{"中文测试", 2, []string{"中", "文", "测", "试", "中文", "文测", "测试"}},
|
||||
{"a b", 1, []string{"a", "b"}},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
got := extractNGrams(tt.text, tt.maxN)
|
||||
for _, want := range tt.check {
|
||||
found := false
|
||||
for _, g := range got {
|
||||
if g == want {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Errorf("extractNGrams(%q, %d) missing %q; got %v", tt.text, tt.maxN, want, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractNGramsNoDups(t *testing.T) {
|
||||
got := extractNGrams("aaaa", 2)
|
||||
seen := make(map[string]bool)
|
||||
for _, g := range got {
|
||||
if seen[g] {
|
||||
t.Errorf("duplicate ngram: %q", g)
|
||||
}
|
||||
seen[g] = true
|
||||
}
|
||||
}
|
||||
|
||||
func TestCosineSimilarity(t *testing.T) {
|
||||
a := Vector{"a": 1, "b": 2}
|
||||
b := Vector{"a": 2, "b": 4}
|
||||
sim := CosineSimilarity(a, b)
|
||||
if math.Abs(sim-1.0) > 1e-6 {
|
||||
t.Errorf("identical direction vectors should have cos=1, got %f", sim)
|
||||
}
|
||||
|
||||
c := Vector{"a": 1, "b": 0}
|
||||
d := Vector{"a": 0, "b": 1}
|
||||
sim = CosineSimilarity(c, d)
|
||||
if math.Abs(sim) > 1e-6 {
|
||||
t.Errorf("orthogonal vectors should have cos=0, got %f", sim)
|
||||
}
|
||||
|
||||
sim = CosineSimilarity(Vector{}, Vector{"a": 1})
|
||||
if sim != 0 {
|
||||
t.Errorf("zero vector should return 0, got %f", sim)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTFIDFVectorizer(t *testing.T) {
|
||||
v := NewTFIDFVectorizer(2)
|
||||
docs := []string{"今天天气很好", "今天心情不错", "明天要下雨"}
|
||||
v.Train(docs)
|
||||
|
||||
vec := v.Vectorize("今天")
|
||||
if len(vec) == 0 {
|
||||
t.Fatal("vector should not be empty")
|
||||
}
|
||||
if _, ok := vec["今天"]; !ok {
|
||||
t.Errorf("expected feature '今天' in vector")
|
||||
}
|
||||
|
||||
// 两个文档都有"今天"，idf 应该较低
|
||||
idfToday := vec["今天"]
|
||||
vecSun := v.Vectorize("下雨")
|
||||
idfRain := vecSun["下雨"]
|
||||
if idfRain <= idfToday {
|
||||
t.Errorf("expected rare '下雨' to have higher idf than common '今天', got today=%f rain=%f", idfToday, idfRain)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTFIDFVectorizerEmpty(t *testing.T) {
|
||||
v := NewTFIDFVectorizer(2)
|
||||
v.Train(nil)
|
||||
vec := v.Vectorize("test")
|
||||
if len(vec) == 0 {
|
||||
t.Error("should produce features even without training")
|
||||
}
|
||||
|
||||
// 未训练时所有 idf=1, 仅有 tf 归一化
|
||||
for _, w := range vec {
|
||||
if w < 0 {
|
||||
t.Errorf("weight should be non-negative, got %f", w)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestInvertedIndex(t *testing.T) {
|
||||
idx := NewInvertedIndex()
|
||||
|
||||
idx.Add("doc1", Vector{"a": 1, "b": 2})
|
||||
idx.Add("doc2", Vector{"b": 1, "c": 3})
|
||||
|
||||
results := idx.Search(Vector{"a": 1}, 10)
|
||||
if len(results) != 1 || results[0] != "doc1" {
|
||||
t.Errorf("search 'a' should return doc1 only, got %v", results)
|
||||
}
|
||||
|
||||
results = idx.Search(Vector{"b": 1}, 10)
|
||||
if len(results) != 2 {
|
||||
t.Errorf("search 'b' should return 2 docs, got %v", results)
|
||||
}
|
||||
|
||||
idx.Remove("doc1")
|
||||
results = idx.Search(Vector{"a": 1}, 10)
|
||||
if len(results) != 0 {
|
||||
t.Errorf("after remove, search 'a' should return empty, got %v", results)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStoreInsertAndSearch(t *testing.T) {
|
||||
s := NewStore()
|
||||
v := NewTFIDFVectorizer(2)
|
||||
v.Train([]string{"hello world", "goodbye world"})
|
||||
|
||||
s.Insert("1", "hello world", v.Vectorize("hello world"), nil)
|
||||
s.Insert("2", "goodbye world", v.Vectorize("goodbye world"), nil)
|
||||
|
||||
if s.Size() != 2 {
|
||||
t.Errorf("expected size 2, got %d", s.Size())
|
||||
}
|
||||
|
||||
results := s.Search(v.Vectorize("hello"), 5)
|
||||
if len(results) == 0 {
|
||||
t.Fatal("expected results")
|
||||
}
|
||||
if results[0].ID != "1" {
|
||||
t.Errorf("expected doc1 as top result for 'hello', got %s", results[0].ID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStoreRemove(t *testing.T) {
|
||||
s := NewStore()
|
||||
v := NewTFIDFVectorizer(1)
|
||||
v.Train([]string{"a"})
|
||||
|
||||
s.Insert("1", "a", v.Vectorize("a"), nil)
|
||||
s.Insert("2", "a", v.Vectorize("a"), nil)
|
||||
s.Remove("1")
|
||||
|
||||
if s.Size() != 1 {
|
||||
t.Errorf("after remove, size should be 1, got %d", s.Size())
|
||||
}
|
||||
|
||||
results := s.Search(v.Vectorize("a"), 5)
|
||||
if len(results) != 1 || results[0].ID != "2" {
|
||||
t.Errorf("only doc2 should remain, got %v", results)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStoreEmpty(t *testing.T) {
|
||||
s := NewStore()
|
||||
results := s.Search(Vector{"a": 1}, 5)
|
||||
if results != nil {
|
||||
t.Errorf("empty store should return nil, got %v", results)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStoreAll(t *testing.T) {
|
||||
s := NewStore()
|
||||
v := NewTFIDFVectorizer(1)
|
||||
v.Train([]string{"a", "b"})
|
||||
|
||||
s.Insert("1", "a", v.Vectorize("a"), map[string]string{"k": "v"})
|
||||
s.Insert("2", "b", v.Vectorize("b"), nil)
|
||||
|
||||
all := s.All()
|
||||
if len(all) != 2 {
|
||||
t.Errorf("All() should return 2 docs, got %d", len(all))
|
||||
}
|
||||
if all[0].Meta["k"] != "v" {
|
||||
t.Errorf("meta should be preserved")
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkCosineSimilarity(b *testing.B) {
|
||||
va := Vector{}
|
||||
vb := Vector{}
|
||||
for i := 0; i < 100; i++ {
|
||||
f := string(rune('a' + i%26))
|
||||
va[f] = float64(i)
|
||||
vb[f] = float64(100 - i)
|
||||
}
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
CosineSimilarity(va, vb)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkExtractNGrams(b *testing.B) {
|
||||
text := "今天天气很好，适合出去散步。明天可能下雨，记得带伞。"
|
||||
for i := 0; i < b.N; i++ {
|
||||
extractNGrams(text, 2)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user