diff --git a/internal/bot/bot.go b/internal/bot/bot.go index 59424d9..a9d65b5 100644 --- a/internal/bot/bot.go +++ b/internal/bot/bot.go @@ -3,6 +3,7 @@ package bot import ( "context" "encoding/json" + "errors" "fmt" "strings" @@ -33,6 +34,8 @@ type Bot struct { toolModel string visionProvider *config.Provider visionModel string + memoryProvider *config.Provider + memoryModel string history []openai.ChatCompletionMessageParamUnion systemPrompt string toolRegistry *tools.Registry @@ -59,6 +62,12 @@ func New(cfg *config.Config) (*Bot, error) { b.visionProvider, b.visionModel = p, m } } + b.memoryProvider, b.memoryModel = b.provider, b.model + if cfg.MemoryModel != "" { + if p, m, err := config.ResolveModel(cfg.MemoryModel); err == nil { + b.memoryProvider, b.memoryModel = p, m + } + } if err := b.toolRegistry.InitConfigs(); err != nil { return nil, err } @@ -119,9 +128,10 @@ func (b *Bot) ThinkingConfig() (string, string) { return b.provider.Thinking, b.provider.ReasoningEffort } -func (b *Bot) CurrentRoles() (tool, vision string) { +func (b *Bot) CurrentRoles() (tool, vision, memory string) { return b.toolProvider.Name + "/" + b.toolModel, - b.visionProvider.Name + "/" + b.visionModel + b.visionProvider.Name + "/" + b.visionModel, + b.memoryProvider.Name + "/" + b.memoryModel } func (b *Bot) Tools() []string { @@ -172,6 +182,93 @@ func estimateTokens(s string) int64 { return int64(len(tke.Encode(s, nil, nil))) } +const memoryExtractPrompt = `你是记忆提取器。从下面的对话中提取值得长期记住的信息,包括: +- 用户的个人偏好、习惯、兴趣 +- 用户的个人事实(职业、所在地、家庭等) +- 项目或任务的背景信息 +- 用户的长期请求或承诺 +只提取确定的信息,忽略闲聊与一次性请求。 +输出 JSON:{"memories": [{"content": "记忆内容", "category": "preference|fact|background|other", "importance": 1到10的整数}]} +没有新记忆时输出 {"memories": []}` + +// ExtractMemories 用记忆 AI 从当前对话历史中提取新记忆。 +// existing 为已存记忆列表,注入 prompt 由 AI 判断去重; +// onReasoning 非空时流式输出 AI 的思考过程。 +func (b *Bot) ExtractMemories(ctx context.Context, existing []store.Memory, onReasoning func(string)) ([]store.Memory, error) { + if b.memoryProvider.APIKey == "" { + return nil, fmt.Errorf("记忆AI供应商 %s 未配置 api_key,请编辑 data/config.yaml", b.memoryProvider.Name) + } + if len(b.history) == 0 { + return nil, errors.New("没有可提取的对话历史") + } + sys := memoryExtractPrompt + if len(existing) > 0 { + var sb strings.Builder + sb.WriteString(sys) + sb.WriteString("\n已有记忆(请勿重复提取):\n") + for _, m := range existing { + fmt.Fprintf(&sb, "- %s (类别: %s, 重要度: %d)\n", m.Content, m.Category, m.Importance) + } + sys = sb.String() + } + messages := make([]openai.ChatCompletionMessageParamUnion, 0, len(b.history)+1) + messages = append(messages, openai.SystemMessage(sys)) + messages = append(messages, b.history...) + + params := openai.ChatCompletionNewParams{ + Model: b.memoryModel, + Messages: messages, + } + b.applyThinkingParams(¶ms, b.memoryProvider) + params.SetExtraFields(map[string]any{"response_format": map[string]string{"type": "json_object"}}) + stream := b.clientFor(b.memoryProvider).Chat.Completions.NewStreaming(ctx, params) + var content strings.Builder + for stream.Next() { + for _, choice := range stream.Current().Choices { + if onReasoning != nil { + if s := extraReasoning(choice.Delta.RawJSON()); s != "" { + onReasoning(s) + } + } + content.WriteString(choice.Delta.Content) + } + } + if err := stream.Err(); err != nil { + return nil, fmt.Errorf("记忆提取请求失败: %w", err) + } + text := strings.TrimSpace(content.String()) + if text == "" { + return nil, nil + } + raw := gjson.Get(text, "memories").Raw + if raw == "" { + raw = text + } + var extracted []struct { + Content string `json:"content"` + Category string `json:"category"` + Importance int `json:"importance"` + } + if err := json.Unmarshal([]byte(raw), &extracted); err != nil { + return nil, fmt.Errorf("解析记忆输出失败: %w", err) + } + var out []store.Memory + for _, e := range extracted { + c := strings.TrimSpace(e.Content) + if c == "" { + continue + } + if e.Importance < 0 { + e.Importance = 0 + } + if e.Importance > 10 { + e.Importance = 10 + } + out = append(out, store.Memory{Content: c, Category: e.Category, Importance: e.Importance}) + } + return out, nil +} + func (b *Bot) SessionMessages() []store.Message { out := make([]store.Message, 0, len(b.history)+1) out = append(out, store.Message{Role: "system", Content: b.systemPrompt}) diff --git a/internal/bot/memory_test.go b/internal/bot/memory_test.go new file mode 100644 index 0000000..82a0f7d --- /dev/null +++ b/internal/bot/memory_test.go @@ -0,0 +1,104 @@ +package bot + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/openai/openai-go" + + "myaibot/internal/config" + "myaibot/internal/store" +) + +func TestExtractMemories(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/chat/completions" { + t.Errorf("请求路径 = %q, want /chat/completions", r.URL.Path) + } + w.Header().Set("Content-Type", "text/event-stream") + w.Write([]byte(`data: {"choices":[{"delta":{"reasoning_content":"用户提到了喝咖啡和Go开发,值得长期记住。","content":"{\"memories\":[{\"content\":\"用户喜欢喝咖啡\",\"category\":\"preference\",\"importance\":7},{\"content\":\"\",\"category\":\"fact\",\"importance\":99},{\"content\":\"用户是Go开发者\",\"category\":\"fact\",\"importance\":-3}]}"}}]}` + "\n\n")) + w.Write([]byte("data: [DONE]\n\n")) + })) + defer srv.Close() + + b := &Bot{ + clients: make(map[string]*openai.Client), + memoryProvider: &config.Provider{Name: "mem", APIKey: "sk-test", BaseURL: srv.URL}, + memoryModel: "m", + history: []openai.ChatCompletionMessageParamUnion{ + openai.UserMessage("我喜欢喝咖啡,是个 Go 开发者"), + }, + } + var reasoning strings.Builder + ms, err := b.ExtractMemories(context.Background(), nil, func(text string) { + reasoning.WriteString(text) + }) + if err != nil { + t.Fatalf("ExtractMemories 出错: %v", err) + } + if len(ms) != 2 { + t.Fatalf("提取数量 = %d, want 2(空内容跳过)", len(ms)) + } + if ms[0].Content != "用户喜欢喝咖啡" || ms[0].Category != "preference" || ms[0].Importance != 7 { + t.Errorf("记忆0异常: %+v", ms[0]) + } + if ms[1].Content != "用户是Go开发者" || ms[1].Importance != 0 { + t.Errorf("重要度应钳制到 0-10: %+v", ms[1]) + } + if reasoning.Len() == 0 { + t.Error("思考流应被回调") + } +} + +func TestExtractMemoriesNoHistory(t *testing.T) { + b := &Bot{ + clients: make(map[string]*openai.Client), + memoryProvider: &config.Provider{Name: "mem", APIKey: "sk-test"}, + memoryModel: "m", + } + if _, err := b.ExtractMemories(context.Background(), nil, nil); err == nil || !strings.Contains(err.Error(), "没有可提取") { + t.Errorf("无历史应报错: %v", err) + } +} + +func TestExtractMemoriesNoAPIKey(t *testing.T) { + b := &Bot{ + clients: make(map[string]*openai.Client), + memoryProvider: &config.Provider{Name: "mem"}, + memoryModel: "m", + history: []openai.ChatCompletionMessageParamUnion{ + openai.UserMessage("hi"), + }, + } + if _, err := b.ExtractMemories(context.Background(), nil, nil); err == nil || !strings.Contains(err.Error(), "api_key") { + t.Errorf("无 api_key 应报错: %v", err) + } +} + +func TestExtractMemoriesEmptyResult(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"{\\\"memories\\\":[]}\"}}]}\n\n")) + w.Write([]byte("data: [DONE]\n\n")) + })) + defer srv.Close() + + b := &Bot{ + clients: make(map[string]*openai.Client), + memoryProvider: &config.Provider{Name: "mem", APIKey: "sk-test", BaseURL: srv.URL}, + memoryModel: "m", + history: []openai.ChatCompletionMessageParamUnion{ + openai.UserMessage("hello"), + }, + } + ms, err := b.ExtractMemories(context.Background(), []store.Memory{{Content: "已有记忆"}}, nil) + if err != nil { + t.Fatalf("ExtractMemories 出错: %v", err) + } + if len(ms) != 0 { + t.Errorf("空结果应返回空切片, got %+v", ms) + } +} diff --git a/internal/cli/cli.go b/internal/cli/cli.go index 9cd2ba4..2cb41dc 100644 --- a/internal/cli/cli.go +++ b/internal/cli/cli.go @@ -1,6 +1,7 @@ package cli import ( + "context" "database/sql" "fmt" "strconv" @@ -62,6 +63,8 @@ func (h *Handler) Handle(input string) bool { fmt.Println(" /effort 设置思考强度") fmt.Println(" /context 打印当前聊天上下文") fmt.Println(" /tools 列出可用工具") + fmt.Println(" /dream 从对话中提取长期记忆") + fmt.Println(" /memories 列出已提取的记忆") fmt.Println(" /sessions 列出历史会话") fmt.Println(" /session 切换到历史会话,如 /session 3") fmt.Println(" /info 显示当前供应商、模型和思考配置") @@ -115,6 +118,52 @@ func (h *Handler) Handle(input string) bool { for _, t := range h.bot.Tools() { fmt.Println(" " + t) } + case "/dream": + existing, err := store.ListMemories(h.db) + if err != nil { + fmt.Printf("⚠️ %v\n", err) + return true + } + thinkStyle := false + ms, err := h.bot.ExtractMemories(context.Background(), existing, func(text string) { + if !thinkStyle { + fmt.Print("\x1b[3;90m🧠 ") + thinkStyle = true + } + fmt.Print(text) + }) + if thinkStyle { + fmt.Print("\x1b[0m\n") + } + if err != nil { + fmt.Printf("⚠️ %v\n", err) + return true + } + if len(ms) == 0 { + fmt.Println("🧠 没有新的记忆") + return true + } + if _, err := store.SaveMemories(h.db, ms); err != nil { + fmt.Printf("⚠️ %v\n", err) + return true + } + fmt.Printf("🧠 已提取 %d 条新记忆\n", len(ms)) + for _, m := range ms { + fmt.Printf(" [%s %d] %s\n", m.Category, m.Importance, m.Content) + } + case "/memories": + list, err := store.ListMemories(h.db) + if err != nil { + fmt.Printf("⚠️ %v\n", err) + return true + } + if len(list) == 0 { + fmt.Println("暂无已提取的记忆") + return true + } + for _, m := range list { + fmt.Printf(" #%d %s [%s %d] %s\n", m.ID, m.CreatedAt.Format("2006-01-02 15:04"), m.Category, m.Importance, m.Content) + } case "/sessions": list, err := store.ListSessions(h.db) if err != nil { @@ -152,7 +201,7 @@ func (h *Handler) Handle(input string) bool { case "/info": provider, model := h.bot.Current() thinking, effort := h.bot.ThinkingConfig() - tool, vision := h.bot.CurrentRoles() + tool, vision, memory := h.bot.CurrentRoles() if thinking == "" { thinking = "enabled(默认)" } @@ -161,7 +210,7 @@ func (h *Handler) Handle(input string) bool { } fmt.Printf("供应商: %s, 模型: %s, 思考模式: %s, 思考强度: %s\n", provider, model, thinking, effort) fmt.Printf("上下文窗口: %s\n", formatWindow(h.bot.ContextWindow())) - fmt.Printf("工具调用AI: %s\n图片识别AI: %s\n", tool, vision) + fmt.Printf("工具调用AI: %s\n图片识别AI: %s\n记忆AI: %s\n", tool, vision, memory) default: fmt.Printf("未知命令: %s,输入 /help 查看命令列表\n", cmd) } diff --git a/internal/cli/complete.go b/internal/cli/complete.go index 78238c1..2f855ac 100644 --- a/internal/cli/complete.go +++ b/internal/cli/complete.go @@ -2,7 +2,7 @@ package cli import "strings" -var commands = []string{"/exit", "/quit", "/help", "/models", "/use", "/think", "/effort", "/context", "/tools", "/sessions", "/session", "/info"} +var commands = []string{"/exit", "/quit", "/help", "/models", "/use", "/think", "/effort", "/context", "/tools", "/dream", "/memories", "/sessions", "/session", "/info"} func Complete(line string, models []string) []string { fields := strings.Fields(line) diff --git a/internal/config/config.go b/internal/config/config.go index d4dbef3..69925bf 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -69,6 +69,7 @@ type Config struct { DefaultModel string `yaml:"default_model"` ToolModel string `yaml:"tool_model"` VisionModel string `yaml:"vision_model"` + MemoryModel string `yaml:"memory_model"` Database DatabaseConfig `yaml:"database"` } @@ -251,6 +252,11 @@ func validate(c *Config) error { return err } } + if c.MemoryModel != "" { + if err := validateModelRef("memory_model", c.MemoryModel, c); err != nil { + return err + } + } d := c.Database if !contains([]string{"sqlite3", "mysql"}, d.Driver) { return fmt.Errorf("database.driver 无效: %q(可选 sqlite3/mysql)", d.Driver) diff --git a/internal/store/memory.go b/internal/store/memory.go new file mode 100644 index 0000000..163be75 --- /dev/null +++ b/internal/store/memory.go @@ -0,0 +1,106 @@ +package store + +import ( + "database/sql" + "fmt" + "strings" + "time" +) + +type Memory struct { + ID int64 + CreatedAt time.Time + SourceSessionID int64 + Content string + Category string + Importance int +} + +const createMemoriesSQLite = ` +CREATE TABLE IF NOT EXISTS memories ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + source_session_id INTEGER NOT NULL DEFAULT 0, + content TEXT NOT NULL, + category TEXT NOT NULL DEFAULT '', + importance INTEGER NOT NULL DEFAULT 5 +)` + +const createMemoriesMySQL = ` +CREATE TABLE IF NOT EXISTS memories ( + id BIGINT AUTO_INCREMENT PRIMARY KEY, + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + source_session_id BIGINT NOT NULL DEFAULT 0, + content TEXT NOT NULL, + category VARCHAR(64) NOT NULL DEFAULT '', + importance INT NOT NULL DEFAULT 5 +)` + +const createMemoriesIndex = "CREATE INDEX IF NOT EXISTS idx_memories_created_at ON memories (created_at)" + +func migrateMemories(db *sql.DB, driver string) error { + var ddl string + switch driver { + case "sqlite3": + ddl = createMemoriesSQLite + case "mysql": + ddl = createMemoriesMySQL + default: + return fmt.Errorf("不支持的数据库驱动: %s", driver) + } + if _, err := db.Exec(ddl); err != nil { + return fmt.Errorf("创建 memories 表失败: %w", err) + } + if _, err := db.Exec(createMemoriesIndex); err != nil { + return fmt.Errorf("创建 memories 索引失败: %w", err) + } + return nil +} + +func SaveMemories(db *sql.DB, memories []Memory) (int64, error) { + if len(memories) == 0 { + return 0, nil + } + placeholders := make([]string, 0, len(memories)) + args := make([]any, 0, len(memories)*4) + for _, m := range memories { + placeholders = append(placeholders, "(?, ?, ?, ?)") + args = append(args, m.SourceSessionID, m.Content, m.Category, m.Importance) + } + query := "INSERT INTO memories (source_session_id, content, category, importance) VALUES " + + strings.Join(placeholders, ", ") + res, err := db.Exec(query, args...) + if err != nil { + return 0, fmt.Errorf("保存记忆失败: %w", err) + } + return res.LastInsertId() +} + +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 { + 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 MemoryCount(db *sql.DB) (int64, error) { + var n int64 + if err := db.QueryRow("SELECT COUNT(*) FROM memories").Scan(&n); err != nil { + return 0, fmt.Errorf("统计记忆失败: %w", err) + } + return n, nil +} diff --git a/internal/store/memory_test.go b/internal/store/memory_test.go new file mode 100644 index 0000000..3258364 --- /dev/null +++ b/internal/store/memory_test.go @@ -0,0 +1,60 @@ +package store + +import ( + "database/sql" + "path/filepath" + "testing" + + "myaibot/internal/config" +) + +func openMemDB(t *testing.T) *sql.DB { + t.Helper() + cfg := &config.DatabaseConfig{ + Driver: "sqlite3", + File: filepath.Join(t.TempDir(), "memory.db"), + } + db, err := Open(cfg) + if err != nil { + t.Fatalf("Open 出错: %v", err) + } + t.Cleanup(func() { Close(db) }) + if err := Migrate(db, "sqlite3"); err != nil { + t.Fatalf("Migrate 出错: %v", err) + } + return db +} + +func TestMemoriesRoundtrip(t *testing.T) { + db := openMemDB(t) + ms := []Memory{ + {Content: "用户喜欢喝咖啡", Category: "preference", Importance: 7}, + {Content: "用户是 Go 开发者", Category: "fact", Importance: 9, SourceSessionID: 3}, + } + if _, err := SaveMemories(db, ms); err != nil { + t.Fatalf("SaveMemories 出错: %v", err) + } + if n, err := MemoryCount(db); err != nil || n != 2 { + t.Errorf("MemoryCount = %d, %v; want 2", n, err) + } + list, err := ListMemories(db) + if err != nil { + t.Fatalf("ListMemories 出错: %v", err) + } + if len(list) != 2 { + t.Fatalf("列表数量 = %d, want 2", len(list)) + } + if list[0].Content != "用户是 Go 开发者" || list[0].Importance != 9 || list[0].SourceSessionID != 3 { + t.Errorf("最新记忆应为 Go 开发者: %+v", list[0]) + } + if list[1].Content != "用户喜欢喝咖啡" || list[1].Category != "preference" { + t.Errorf("记忆顺序/内容异常: %+v", list[1]) + } +} + +func TestSaveMemoriesEmpty(t *testing.T) { + db := openMemDB(t) + if _, err := SaveMemories(db, nil); err != nil { + t.Fatalf("空列表不应报错: %v", err) + } +} diff --git a/internal/store/session.go b/internal/store/session.go index fbdeab1..aab2fa8 100644 --- a/internal/store/session.go +++ b/internal/store/session.go @@ -62,7 +62,7 @@ func Migrate(db *sql.DB, driver string) error { if _, err := db.Exec(ddl); err != nil { return fmt.Errorf("创建 sessions 表失败: %w", err) } - return nil + return migrateMemories(db, driver) } func SaveSession(db *sql.DB, s *Session) (int64, error) {