新增 recall_memory 回忆工具：按 token 搜索长期记忆

internal/bot/bot.go

@@ -2,6 +2,7 @@ package bot
 
 import (
 	"context"
+	"database/sql"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -12,10 +13,10 @@ import (
 	"github.com/openai/openai-go/packages/param"
 	"github.com/openai/openai-go/shared"
 	"github.com/openai/openai-go/shared/constant"
-	"github.com/pkoukk/tiktoken-go"
 	"github.com/tidwall/gjson"
 	"myaibot/internal/config"
 	"myaibot/internal/store"
+	"myaibot/internal/tokens"
 	"myaibot/internal/tools"
 	"myaibot/internal/tools/builtin"
 )
@@ -41,12 +42,12 @@ type Bot struct {
 	toolRegistry   *tools.Registry
 }
 
-func New(cfg *config.Config) (*Bot, error) {
+func New(cfg *config.Config, db *sql.DB) (*Bot, error) {
 	b := &Bot{
 		clients:      make(map[string]*openai.Client),
 		cfg:          cfg,
 		systemPrompt: cfg.SystemPrompt,
-		toolRegistry: tools.NewRegistry(builtin.NewTimeTool(), builtin.NewCalculatorTool(), builtin.NewRandomTool()),
+		toolRegistry: tools.NewRegistry(builtin.NewTimeTool(), builtin.NewCalculatorTool(), builtin.NewRandomTool(), builtin.NewRecallTool(db)),
 	}
 	b.provider = config.FindProviderIn(cfg, cfg.DefaultProvider)
 	b.model = cfg.DefaultModel
@@ -153,7 +154,7 @@ func (b *Bot) ContextWindow() int64 {
 // ContextStats 统计当前上下文的 token 使用量与窗口总大小（0 表示未配置）。
 func (b *Bot) ContextStats() (used, total int64) {
 	total = b.ContextWindow()
-	used += estimateTokens(b.systemPrompt)
+	used += tokens.Count(b.systemPrompt)
 	for _, msg := range b.history {
 		var content string
 		switch {
@@ -164,54 +165,11 @@ func (b *Bot) ContextStats() (used, total int64) {
 		case msg.OfSystem != nil:
 			content = msg.OfSystem.Content.OfString.Value
 		}
-		used += estimateTokens(content)
+		used += tokens.Count(content)
 	}
 	return used, total
 }
 
-var tke *tiktoken.Tiktoken
-
-// Tokenize 将文本编码为 o200k_base token id 列表（去重）。
-func Tokenize(text string) []int64 {
-	if text == "" {
-		return nil
-	}
-	if tke == nil {
-		t, err := tiktoken.GetEncoding("o200k_base")
-		if err != nil {
-			return nil
-		}
-		tke = t
-	}
-	seen := make(map[int64]bool)
-	var out []int64
-	for _, id := range tke.Encode(text, nil, nil) {
-		tid := int64(id)
-		if seen[tid] {
-			continue
-		}
-		seen[tid] = true
-		out = append(out, tid)
-	}
-	return out
-}
-
-// estimateTokens 用 o200k_base 词表精确统计 token；
-// 初始化失败（如无法下载词表）时回退为字符数/2 估算。
-func estimateTokens(s string) int64 {
-	if s == "" {
-		return 0
-	}
-	if tke == nil {
-		t, err := tiktoken.GetEncoding("o200k_base")
-		if err != nil {
-			return int64(len([]rune(s)) / 2)
-		}
-		tke = t
-	}
-	return int64(len(tke.Encode(s, nil, nil)))
-}
-
 const memoryExtractPrompt = `你是记忆提取器。从下面的对话中提取值得长期记住的信息，包括：
 - 用户的个人偏好、习惯、兴趣
 - 用户的个人事实（职业、所在地、家庭等）

internal/bot/session_test.go

@@ -12,7 +12,7 @@ import (
 func newTestBot(t *testing.T) *Bot {
 	t.Helper()
 	t.Chdir(t.TempDir())
-	for _, name := range []string{"get_current_time", "calculate", "random_number"} {
+	for _, name := range []string{"get_current_time", "calculate", "random_number", "recall_memory"} {
 		if err := config.WriteDefaultToolConfig(name, map[string]any{"enabled": true, "prompt": "p"}); err != nil {
 			t.Fatalf("写入工具配置失败: %v", err)
 		}
@@ -26,7 +26,7 @@ func newTestBot(t *testing.T) *Bot {
 			{Name: "p", BaseURL: "x", Models: []config.ModelConfig{{Name: "m"}}},
 		},
 	}
-	b, err := New(cfg)
+	b, err := New(cfg, nil)
 	if err != nil {
 		t.Fatalf("New 出错: %v", err)
 	}

internal/bot/stats_test.go

@@ -3,13 +3,15 @@ package bot
 import (
 	"testing"
 
+	"myaibot/internal/tokens"
+
 	"github.com/openai/openai-go"
 
 	"myaibot/internal/config"
 )
 
 func TestEstimateTokensEmpty(t *testing.T) {
-	if n := estimateTokens(""); n != 0 {
+	if n := tokens.Count(""); n != 0 {
 		t.Errorf("空串应为 0, got %d", n)
 	}
 }
@@ -25,14 +27,14 @@ func TestEstimateTokensKnown(t *testing.T) {
 		{"你是一个乐于助人的 AI 助手。", 12},
 	}
 	for _, c := range cases {
-		if n := estimateTokens(c.text); n != c.want {
-			t.Errorf("estimateTokens(%q) = %d, want %d", c.text, n, c.want)
+		if n := tokens.Count(c.text); n != c.want {
+			t.Errorf("tokens.Count(%q) = %d, want %d", c.text, n, c.want)
 		}
 	}
 }
 
 func TestTokenize(t *testing.T) {
-	ids := Tokenize("hello hello world")
+	ids := tokens.Tokenize("hello hello world")
 	if len(ids) < 2 {
 		t.Errorf("应有多个 token, got %v", ids)
 	}
@@ -43,10 +45,10 @@ func TestTokenize(t *testing.T) {
 		}
 		seen[id] = true
 	}
-	if len(Tokenize("")) != 0 {
+	if len(tokens.Tokenize("")) != 0 {
 		t.Error("空串应返回空")
 	}
-	chinese := Tokenize("用户喜欢喝咖啡")
+	chinese := tokens.Tokenize("用户喜欢喝咖啡")
 	if len(chinese) == 0 {
 		t.Error("中文文本应产生 token")
 	}

internal/cli/cli.go

@@ -9,6 +9,7 @@ import (
 
 	"myaibot/internal/bot"
 	"myaibot/internal/store"
+	"myaibot/internal/tokens"
 )
 
 type Handler struct {
@@ -62,7 +63,7 @@ func (h *Handler) indexMemories(newIDs []int64) error {
 		if m == nil {
 			continue
 		}
-		if err := store.SaveMemoryTokens(h.db, id, bot.Tokenize(m.Content)); err != nil {
+		if err := store.SaveMemoryTokens(h.db, id, tokens.Tokenize(m.Content)); err != nil {
 			return fmt.Errorf("记忆 #%d 建索引失败: %w", id, err)
 		}
 	}

internal/store/memory.go

@@ -3,6 +3,7 @@ package store
 import (
 	"database/sql"
 	"fmt"
+	"strings"
 	"time"
 )
 
@@ -183,6 +184,51 @@ func UnindexedMemoryIDs(db *sql.DB) ([]int64, error) {
 	return out, rows.Err()
 }
 
+// SearchMemoriesByTokens 按 token id 搜索相关记忆，按命中 token 数从多到少排序。
+// limit 钳制在 1-10。
+func SearchMemoriesByTokens(db *sql.DB, tokenIDs []int64, limit int) ([]Memory, error) {
+	if len(tokenIDs) == 0 {
+		return nil, nil
+	}
+	if limit < 1 {
+		limit = 1
+	}
+	if limit > 10 {
+		limit = 10
+	}
+	placeholders := make([]string, len(tokenIDs))
+	args := make([]any, 0, len(tokenIDs)+1)
+	for i, tid := range tokenIDs {
+		placeholders[i] = "?"
+		args = append(args, tid)
+	}
+	args = append(args, limit)
+	query := `SELECT m.id, m.created_at, m.source_session_id, m.content, m.category, m.importance
+FROM memory_tokens mt JOIN memories m ON m.id = mt.memory_id
+WHERE mt.token_id IN (` + strings.Join(placeholders, ", ") + `)
+GROUP BY m.id
+ORDER BY COUNT(*) DESC, m.id DESC
+LIMIT ?`
+	rows, err := db.Query(query, args...)
+	if err != nil {
+		return nil, fmt.Errorf("搜索记忆失败: %w", err)
+	}
+	defer rows.Close()
+	var out []Memory
+	for rows.Next() {
+		var (
+			m       Memory
+			created string
+		)
+		if err := rows.Scan(&m.ID, &created, &m.SourceSessionID, &m.Content, &m.Category, &m.Importance); err != nil {
+			return nil, fmt.Errorf("读取记忆失败: %w", err)
+		}
+		m.CreatedAt = parseTime(created)
+		out = append(out, m)
+	}
+	return out, rows.Err()
+}
+
 func ListMemories(db *sql.DB) ([]Memory, error) {
 	rows, err := db.Query("SELECT id, created_at, source_session_id, content, category, importance FROM memories ORDER BY id DESC")
 	if err != nil {

internal/store/memory_test.go

@@ -146,3 +146,68 @@ func TestLoadMemory(t *testing.T) {
 		t.Errorf("不存在应返回 nil: %v, %v", missing, err)
 	}
 }
+
+func TestSearchMemoriesByTokens(t *testing.T) {
+	db := openMemDB(t)
+	ids, err := SaveMemories(db, []Memory{
+		{Content: "用户喜欢喝咖啡", Category: "preference"},
+		{Content: "用户喜欢喝咖啡和茶", Category: "preference"},
+		{Content: "用户是 Go 开发者", Category: "fact"},
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	// 手工建索引：记忆1 含 token 100,200；记忆2 含 100,200,300；记忆3 含 400
+	if err := SaveMemoryTokens(db, ids[0], []int64{100, 200}); err != nil {
+		t.Fatal(err)
+	}
+	if err := SaveMemoryTokens(db, ids[1], []int64{100, 200, 300}); err != nil {
+		t.Fatal(err)
+	}
+	if err := SaveMemoryTokens(db, ids[2], []int64{400}); err != nil {
+		t.Fatal(err)
+	}
+
+	// 命中数优先：记忆2 命中 3 个 token 排第一
+	res, err := SearchMemoriesByTokens(db, []int64{100, 200, 300, 400}, 5)
+	if err != nil {
+		t.Fatalf("搜索出错: %v", err)
+	}
+	if len(res) != 3 {
+		t.Fatalf("命中数量 = %d, want 3", len(res))
+	}
+	if res[0].ID != ids[1] {
+		t.Errorf("命中数最多的应排第一: %v", res[0])
+	}
+
+	// limit 钳制与命中子集
+	res, err = SearchMemoriesByTokens(db, []int64{100}, 10)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(res) != 2 {
+		t.Errorf("limit 10 应命中 2 条, got %d", len(res))
+	}
+	res, err = SearchMemoriesByTokens(db, []int64{100}, 1)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(res) != 1 {
+		t.Errorf("limit 1 应只返回 1 条, got %d", len(res))
+	}
+
+	// 无命中
+	res, err = SearchMemoriesByTokens(db, []int64{9999}, 5)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(res) != 0 {
+		t.Errorf("无命中应返回空, got %d", len(res))
+	}
+
+	// 空 token
+	res, err = SearchMemoriesByTokens(db, nil, 5)
+	if err != nil || res != nil {
+		t.Errorf("空 token 应返回 nil, %v %v", res, err)
+	}
+}

internal/tokens/tokens.go

@@ -0,0 +1,51 @@
+package tokens
+
+import "github.com/pkoukk/tiktoken-go"
+
+var tke *tiktoken.Tiktoken
+
+func getEncoding() *tiktoken.Tiktoken {
+	if tke == nil {
+		t, err := tiktoken.GetEncoding("o200k_base")
+		if err != nil {
+			return nil
+		}
+		tke = t
+	}
+	return tke
+}
+
+// Tokenize 将文本编码为 o200k_base token id 列表（去重）。
+// 编码器初始化失败时返回 nil。
+func Tokenize(text string) []int64 {
+	if text == "" {
+		return nil
+	}
+	enc := getEncoding()
+	if enc == nil {
+		return nil
+	}
+	seen := make(map[int64]bool)
+	var out []int64
+	for _, id := range enc.Encode(text, nil, nil) {
+		tid := int64(id)
+		if seen[tid] {
+			continue
+		}
+		seen[tid] = true
+		out = append(out, tid)
+	}
+	return out
+}
+
+// Count 统计文本的 token 数量；编码器初始化失败时回退为字符数/2 估算。
+func Count(text string) int64 {
+	if text == "" {
+		return 0
+	}
+	enc := getEncoding()
+	if enc == nil {
+		return int64(len([]rune(text)) / 2)
+	}
+	return int64(len(enc.Encode(text, nil, nil)))
+}

internal/tools/builtin/builtin_test.go

@@ -0,0 +1,84 @@
+package builtin
+
+import (
+	"database/sql"
+	"encoding/json"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"myaibot/internal/config"
+	"myaibot/internal/store"
+	"myaibot/internal/tokens"
+)
+
+func recallTestDB(t *testing.T) *sql.DB {
+	t.Helper()
+	cfg := &config.DatabaseConfig{
+		Driver: "sqlite3",
+		File:   filepath.Join(t.TempDir(), "memory.db"),
+	}
+	db, err := store.Open(cfg)
+	if err != nil {
+		t.Fatalf("Open 出错: %v", err)
+	}
+	t.Cleanup(func() { store.Close(db) })
+	if err := store.Migrate(db, "sqlite3"); err != nil {
+		t.Fatalf("Migrate 出错: %v", err)
+	}
+	ids, err := store.SaveMemories(db, []store.Memory{
+		{Content: "用户喜欢喝咖啡", Category: "preference", Importance: 7},
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := store.SaveMemoryTokens(db, ids[0], tokens.Tokenize("用户喜欢喝咖啡")); err != nil {
+		t.Fatal(err)
+	}
+	return db
+}
+
+func TestRecallToolFound(t *testing.T) {
+	tool := NewRecallTool(recallTestDB(t))
+	args, _ := json.Marshal(map[string]any{"query": "咖啡"})
+	out, err := tool.Execute(args)
+	if err != nil {
+		t.Fatalf("Execute 出错: %v", err)
+	}
+	if !strings.Contains(out, "找到 1 条相关记忆") || !strings.Contains(out, "用户喜欢喝咖啡") {
+		t.Errorf("输出异常: %q", out)
+	}
+}
+
+func TestRecallToolNotFound(t *testing.T) {
+	tool := NewRecallTool(recallTestDB(t))
+	args, _ := json.Marshal(map[string]any{"query": "不存在的关键词"})
+	out, err := tool.Execute(args)
+	if err != nil {
+		t.Fatalf("Execute 出错: %v", err)
+	}
+	if out != "未找到相关记忆" {
+		t.Errorf("应返回未找到, got %q", out)
+	}
+}
+
+func TestRecallToolEmptyQuery(t *testing.T) {
+	tool := NewRecallTool(recallTestDB(t))
+	args, _ := json.Marshal(map[string]any{"query": ""})
+	if _, err := tool.Execute(args); err == nil {
+		t.Error("空 query 应报错")
+	}
+}
+
+func TestRecallToolConfigure(t *testing.T) {
+	tool := NewRecallTool(nil)
+	if err := tool.Configure(map[string]any{"enabled": false, "prompt": "自定义提示"}); err != nil {
+		t.Fatalf("Configure 出错: %v", err)
+	}
+	if tool.Enabled() {
+		t.Error("应被禁用")
+	}
+	if tool.Description() != "自定义提示" {
+		t.Errorf("Description = %q", tool.Description())
+	}
+}

internal/tools/builtin/recall_tool.go

@@ -0,0 +1,84 @@
+package builtin
+
+import (
+	"database/sql"
+	"encoding/json"
+	"fmt"
+	"strings"
+
+	"myaibot/internal/store"
+	"myaibot/internal/tokens"
+)
+
+type recallTool struct {
+	db      *sql.DB
+	enabled bool
+	prompt  string
+}
+
+func NewRecallTool(db *sql.DB) *recallTool {
+	return &recallTool{
+		db:      db,
+		enabled: true,
+		prompt:  "从长期记忆中回忆与用户提问相关的信息。当用户询问个人偏好、个人信息、之前聊过的话题或需要回顾历史对话时调用此工具",
+	}
+}
+
+func (t *recallTool) Name() string        { return "recall_memory" }
+func (t *recallTool) Description() string { return t.prompt }
+func (t *recallTool) Enabled() bool       { return t.enabled }
+func (t *recallTool) DefaultConfig() map[string]any {
+	return map[string]any{"enabled": true, "prompt": t.prompt}
+}
+
+func (t *recallTool) Configure(cfg map[string]any) error {
+	var err error
+	if t.enabled, err = parseEnabled(cfg); err != nil {
+		return err
+	}
+	if p, ok := cfg["prompt"].(string); ok && p != "" {
+		t.prompt = p
+	}
+	return nil
+}
+
+func (t *recallTool) Parameters() map[string]any {
+	return map[string]any{
+		"type": "object",
+		"properties": map[string]any{
+			"query": map[string]any{"type": "string", "description": "要回忆的内容或用户的提问"},
+			"limit": map[string]any{"type": "integer", "description": "最多返回的记忆条数，默认 5，最大 10"},
+		},
+		"required": []string{"query"},
+	}
+}
+
+func (t *recallTool) Execute(args json.RawMessage) (string, error) {
+	var p struct {
+		Query string `json:"query"`
+		Limit int    `json:"limit"`
+	}
+	if err := json.Unmarshal(args, &p); err != nil {
+		return "", err
+	}
+	query := strings.TrimSpace(p.Query)
+	if query == "" {
+		return "", fmt.Errorf("query 不能为空")
+	}
+	if p.Limit == 0 {
+		p.Limit = 5
+	}
+	memories, err := store.SearchMemoriesByTokens(t.db, tokens.Tokenize(query), p.Limit)
+	if err != nil {
+		return "", err
+	}
+	if len(memories) == 0 {
+		return "未找到相关记忆", nil
+	}
+	var sb strings.Builder
+	fmt.Fprintf(&sb, "找到 %d 条相关记忆：\n", len(memories))
+	for _, m := range memories {
+		fmt.Fprintf(&sb, "- [%s %d] %s\n", m.Category, m.Importance, m.Content)
+	}
+	return strings.TrimSuffix(sb.String(), "\n"), nil
+}

internal/tools/builtin_test.go

@@ -1,115 +0,0 @@
-package tools_test
-
-import (
-	"encoding/json"
-	"strings"
-	"testing"
-
-	"myaibot/internal/tools/builtin"
-)
-
-func TestCalculator(t *testing.T) {
-	cases := []struct {
-		name string
-		expr string
-		want string
-	}{
-		{"四则运算", "(12 + 3) * 4", "60"},
-		{"幂运算", "2^10", "1024"},
-		{"浮点", "7 / 2", "3.5"},
-		{"小数精度", "1 / 3", "0.33333333"},
-		{"布尔", "2 > 1", "true"},
-	}
-	for _, c := range cases {
-		t.Run(c.name, func(t *testing.T) {
-			args, _ := json.Marshal(map[string]string{"expression": c.expr})
-			got, err := builtin.NewCalculatorTool().Execute(args)
-			if err != nil {
-				t.Fatalf("Execute 出错: %v", err)
-			}
-			if got != c.want {
-				t.Errorf("Execute(%q) = %q, want %q", c.expr, got, c.want)
-			}
-		})
-	}
-}
-
-func TestCalculatorInvalid(t *testing.T) {
-	args, _ := json.Marshal(map[string]string{"expression": "1 +"})
-	if _, err := builtin.NewCalculatorTool().Execute(args); err == nil {
-		t.Error("非法表达式应返回错误")
-	}
-}
-
-func TestRandom(t *testing.T) {
-	args, _ := json.Marshal(map[string]any{"min": 1, "max": 10})
-	for i := 0; i < 100; i++ {
-		out, err := builtin.NewRandomTool().Execute(args)
-		if err != nil {
-			t.Fatalf("Execute 出错: %v", err)
-		}
-		var v int
-		if err := json.Unmarshal([]byte(out), &v); err != nil {
-			t.Fatalf("结果 %q 解析失败: %v", out, err)
-		}
-		if v < 1 || v > 10 {
-			t.Fatalf("结果 %q 不在 [1,10] 内", out)
-		}
-	}
-	bad, _ := json.Marshal(map[string]any{"min": 10, "max": 1})
-	if _, err := builtin.NewRandomTool().Execute(bad); err == nil {
-		t.Error("min>max 应返回错误")
-	}
-}
-
-func TestTimeTool(t *testing.T) {
-	out, err := builtin.NewTimeTool().Execute(nil)
-	if err != nil {
-		t.Fatalf("Execute 出错: %v", err)
-	}
-	if !strings.Contains(out, "20") {
-		t.Errorf("时间输出异常: %q", out)
-	}
-}
-
-func TestConfigurePrompt(t *testing.T) {
-	tool := builtin.NewTimeTool()
-	err := tool.Configure(map[string]any{"enabled": true, "prompt": "自定义提示词"})
-	if err != nil {
-		t.Fatalf("Configure 出错: %v", err)
-	}
-	if tool.Description() != "自定义提示词" {
-		t.Errorf("Description = %q, want 自定义提示词", tool.Description())
-	}
-	if !tool.Enabled() {
-		t.Error("enabled 应为 true")
-	}
-}
-
-func TestConfigureDisable(t *testing.T) {
-	tool := builtin.NewCalculatorTool()
-	if err := tool.Configure(map[string]any{"enabled": false}); err != nil {
-		t.Fatalf("Configure 出错: %v", err)
-	}
-	if tool.Enabled() {
-		t.Error("enabled 应为 false")
-	}
-}
-
-func TestConfigureInvalid(t *testing.T) {
-	tool := builtin.NewRandomTool()
-	if err := tool.Configure(map[string]any{"enabled": "yes"}); err == nil {
-		t.Error("enabled 非布尔值应报错")
-	}
-}
-
-func TestDefaultConfig(t *testing.T) {
-	tool := builtin.NewTimeTool()
-	cfg := tool.DefaultConfig()
-	if cfg["enabled"] != true {
-		t.Errorf("默认 enabled 应为 true, got %v", cfg["enabled"])
-	}
-	if p, ok := cfg["prompt"].(string); !ok || p == "" {
-		t.Errorf("默认 prompt 缺失: %v", cfg["prompt"])
-	}
-}

main.go

@@ -33,7 +33,7 @@ func main() {
 	}
 	fmt.Printf("💾 数据库已连接 (%s)\n", cfg.Database.Driver)
 
-	b, err := bot.New(cfg)
+	b, err := bot.New(cfg, db)
 	if err != nil {
 		fmt.Printf("⚠️  %v\n", err)
 		fmt.Println("请填写工具配置文件后重新启动。")