kevin/ai-bot
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

新增/dream记忆提取与/memories查看：专用记忆AI、memories表、思考流输出

Browse Source
This commit is contained in:
kevin
2026-08-14 19:26:33 +08:00
parent ba2ee4b0a0
commit fc7dd39c54
8 changed files with 428 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/bot/bot.go
+99 -2
View File
@@ -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(&params, 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})
internal/bot/memory_test.go
+104
View File
@@ -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)
}
}
internal/cli/cli.go
+51 -2
View File
@@ -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 <low|high|max> 设置思考强度")
fmt.Println(" /context 打印当前聊天上下文")
fmt.Println(" /tools 列出可用工具")
fmt.Println(" /dream 从对话中提取长期记忆")
fmt.Println(" /memories 列出已提取的记忆")
fmt.Println(" /sessions 列出历史会话")
fmt.Println(" /session <id> 切换到历史会话，如 /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)
}
internal/cli/complete.go
+1 -1
View File
@@ -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)
internal/config/config.go
+6
View File
@@ -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)
internal/store/memory.go
+106
View File
@@ -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
}
internal/store/memory_test.go
+60
View File
@@ -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)
}
}
internal/store/session.go
+1 -1
View File
@@ -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) {