390 lines
11 KiB
Go
390 lines
11 KiB
Go
package store
|
|
|
|
import (
|
|
"database/sql"
|
|
"fmt"
|
|
"strings"
|
|
"time"
|
|
|
|
"myaibot/internal/tokens"
|
|
)
|
|
|
|
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)"
|
|
|
|
const createTokensSQLite = `
|
|
CREATE TABLE IF NOT EXISTS tokens (
|
|
token_id INTEGER PRIMARY KEY,
|
|
token_text TEXT NOT NULL DEFAULT '',
|
|
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
|
|
)`
|
|
|
|
const createTokensMySQL = `
|
|
CREATE TABLE IF NOT EXISTS tokens (
|
|
token_id BIGINT PRIMARY KEY,
|
|
token_text TEXT NOT NULL,
|
|
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
|
|
)`
|
|
|
|
const createMemoryTokensSQLite = `
|
|
CREATE TABLE IF NOT EXISTS memory_tokens (
|
|
memory_id INTEGER NOT NULL,
|
|
token_id INTEGER NOT NULL,
|
|
PRIMARY KEY (memory_id, token_id)
|
|
)`
|
|
|
|
const createMemoryTokensMySQL = `
|
|
CREATE TABLE IF NOT EXISTS memory_tokens (
|
|
memory_id BIGINT NOT NULL,
|
|
token_id BIGINT NOT NULL,
|
|
PRIMARY KEY (memory_id, token_id)
|
|
)`
|
|
|
|
const createMemoryTokensIndex = "CREATE INDEX IF NOT EXISTS idx_memory_tokens_token_id ON memory_tokens (token_id)"
|
|
|
|
// migrateTokens 创建 tokens 表；旧表缺失 token_text 列时补列并回填文本。
|
|
func migrateTokens(db *sql.DB, driver string) error {
|
|
switch driver {
|
|
case "sqlite3":
|
|
if _, err := db.Exec(createTokensSQLite); err != nil {
|
|
return fmt.Errorf("创建 tokens 表失败: %w", err)
|
|
}
|
|
case "mysql":
|
|
if _, err := db.Exec(createTokensMySQL); err != nil {
|
|
return fmt.Errorf("创建 tokens 表失败: %w", err)
|
|
}
|
|
}
|
|
has, err := hasTokenTextColumn(db, driver)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if !has {
|
|
ddl := "ALTER TABLE tokens ADD COLUMN token_text TEXT NOT NULL DEFAULT ''"
|
|
if driver == "mysql" {
|
|
ddl = "ALTER TABLE tokens ADD COLUMN token_text TEXT NOT NULL"
|
|
}
|
|
if _, err := db.Exec(ddl); err != nil {
|
|
return fmt.Errorf("tokens 表增加 token_text 列失败: %w", err)
|
|
}
|
|
}
|
|
return backfillTokenTexts(db)
|
|
}
|
|
|
|
func hasTokenTextColumn(db *sql.DB, driver string) (bool, error) {
|
|
var query string
|
|
switch driver {
|
|
case "sqlite3":
|
|
query = "SELECT name FROM pragma_table_info('tokens')"
|
|
case "mysql":
|
|
query = "SELECT column_name FROM information_schema.COLUMNS WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME = 'tokens' AND COLUMN_NAME = 'token_text'"
|
|
default:
|
|
return false, fmt.Errorf("不支持的数据库驱动: %s", driver)
|
|
}
|
|
rows, err := db.Query(query)
|
|
if err != nil {
|
|
return false, fmt.Errorf("检查 tokens 表结构失败: %w", err)
|
|
}
|
|
defer rows.Close()
|
|
for rows.Next() {
|
|
var name string
|
|
if err := rows.Scan(&name); err != nil {
|
|
return false, fmt.Errorf("读取 tokens 表结构失败: %w", err)
|
|
}
|
|
if name == "token_text" {
|
|
return true, nil
|
|
}
|
|
}
|
|
return false, rows.Err()
|
|
}
|
|
|
|
// backfillTokenTexts 为 token_text 为空的 token 行回填符号文本。
|
|
func backfillTokenTexts(db *sql.DB) error {
|
|
rows, err := db.Query("SELECT token_id FROM tokens WHERE token_text = ''")
|
|
if err != nil {
|
|
return fmt.Errorf("查询待回填 token 失败: %w", err)
|
|
}
|
|
var ids []int64
|
|
for rows.Next() {
|
|
var id int64
|
|
if err := rows.Scan(&id); err != nil {
|
|
rows.Close()
|
|
return fmt.Errorf("读取 token id 失败: %w", err)
|
|
}
|
|
ids = append(ids, id)
|
|
}
|
|
rows.Close()
|
|
if err := rows.Err(); err != nil {
|
|
return err
|
|
}
|
|
for _, id := range ids {
|
|
text := tokens.Text(id)
|
|
if text == "" {
|
|
continue
|
|
}
|
|
if _, err := db.Exec("UPDATE tokens SET token_text = ? WHERE token_id = ? AND token_text = ''", text, id); err != nil {
|
|
return fmt.Errorf("回填 token 文本失败: %w", err)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
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)
|
|
}
|
|
switch driver {
|
|
case "sqlite3":
|
|
ddl = createTokensSQLite
|
|
case "mysql":
|
|
ddl = createTokensMySQL
|
|
}
|
|
if _, err := db.Exec(ddl); err != nil {
|
|
return fmt.Errorf("创建 tokens 表失败: %w", err)
|
|
}
|
|
if err := migrateTokens(db, driver); err != nil {
|
|
return err
|
|
}
|
|
switch driver {
|
|
case "sqlite3":
|
|
ddl = createMemoryTokensSQLite
|
|
case "mysql":
|
|
ddl = createMemoryTokensMySQL
|
|
}
|
|
if _, err := db.Exec(ddl); err != nil {
|
|
return fmt.Errorf("创建 memory_tokens 表失败: %w", err)
|
|
}
|
|
if _, err := db.Exec(createMemoryTokensIndex); err != nil {
|
|
return fmt.Errorf("创建 memory_tokens 索引失败: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// SaveMemories 逐条插入记忆，返回每条记忆的 id。
|
|
func SaveMemories(db *sql.DB, memories []Memory) ([]int64, error) {
|
|
if len(memories) == 0 {
|
|
return nil, nil
|
|
}
|
|
ids := make([]int64, 0, len(memories))
|
|
for _, m := range memories {
|
|
res, err := db.Exec(
|
|
"INSERT INTO memories (source_session_id, content, category, importance) VALUES (?, ?, ?, ?)",
|
|
m.SourceSessionID, m.Content, m.Category, m.Importance,
|
|
)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("保存记忆失败: %w", err)
|
|
}
|
|
id, err := res.LastInsertId()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("读取记忆 id 失败: %w", err)
|
|
}
|
|
ids = append(ids, id)
|
|
}
|
|
return ids, nil
|
|
}
|
|
|
|
// SaveMemoryTokens 为记忆建立 token 索引：token 去重、关联幂等。
|
|
func SaveMemoryTokens(db *sql.DB, memoryID int64, tokenList []tokens.Token) error {
|
|
if len(tokenList) == 0 {
|
|
return nil
|
|
}
|
|
tx, err := db.Begin()
|
|
if err != nil {
|
|
return fmt.Errorf("开启事务失败: %w", err)
|
|
}
|
|
defer tx.Rollback()
|
|
for _, t := range tokenList {
|
|
if _, err := tx.Exec("INSERT OR IGNORE INTO tokens (token_id, token_text) VALUES (?, ?)", t.ID, t.Text); err != nil {
|
|
return fmt.Errorf("保存 token 失败: %w", err)
|
|
}
|
|
if _, err := tx.Exec("INSERT OR IGNORE INTO memory_tokens (memory_id, token_id) VALUES (?, ?)", memoryID, t.ID); err != nil {
|
|
return fmt.Errorf("建立 token 关联失败: %w", err)
|
|
}
|
|
}
|
|
return tx.Commit()
|
|
}
|
|
|
|
func LoadMemory(db *sql.DB, id int64) (*Memory, error) {
|
|
row := db.QueryRow("SELECT id, created_at, source_session_id, content, category, importance FROM memories WHERE id = ?", id)
|
|
var (
|
|
m Memory
|
|
created string
|
|
)
|
|
if err := row.Scan(&m.ID, &created, &m.SourceSessionID, &m.Content, &m.Category, &m.Importance); err != nil {
|
|
if err == sql.ErrNoRows {
|
|
return nil, nil
|
|
}
|
|
return nil, fmt.Errorf("读取记忆失败: %w", err)
|
|
}
|
|
m.CreatedAt = parseTime(created)
|
|
return &m, nil
|
|
}
|
|
|
|
// UnindexedMemoryIDs 返回尚未建立 token 索引的记忆 id。
|
|
func UnindexedMemoryIDs(db *sql.DB) ([]int64, error) {
|
|
rows, err := db.Query("SELECT m.id FROM memories m LEFT JOIN memory_tokens mt ON m.id = mt.memory_id WHERE mt.memory_id IS NULL")
|
|
if err != nil {
|
|
return nil, fmt.Errorf("查询未索引记忆失败: %w", err)
|
|
}
|
|
defer rows.Close()
|
|
var out []int64
|
|
for rows.Next() {
|
|
var id int64
|
|
if err := rows.Scan(&id); err != nil {
|
|
return nil, fmt.Errorf("读取记忆 id 失败: %w", err)
|
|
}
|
|
out = append(out, id)
|
|
}
|
|
return out, rows.Err()
|
|
}
|
|
|
|
// RebuildTokenIndex 清空 token 索引并重新为全部记忆建立索引，返回处理的记忆数。
|
|
func RebuildTokenIndex(db *sql.DB) (int, error) {
|
|
memories, err := ListMemories(db)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
tx, err := db.Begin()
|
|
if err != nil {
|
|
return 0, fmt.Errorf("开启事务失败: %w", err)
|
|
}
|
|
defer tx.Rollback()
|
|
if _, err := tx.Exec("DELETE FROM memory_tokens"); err != nil {
|
|
return 0, fmt.Errorf("清空 memory_tokens 失败: %w", err)
|
|
}
|
|
if _, err := tx.Exec("DELETE FROM tokens"); err != nil {
|
|
return 0, fmt.Errorf("清空 tokens 失败: %w", err)
|
|
}
|
|
for _, m := range memories {
|
|
tokenList := tokens.Tokenize(m.Content)
|
|
if len(tokenList) == 0 {
|
|
continue
|
|
}
|
|
for _, t := range tokenList {
|
|
if _, err := tx.Exec("INSERT OR IGNORE INTO tokens (token_id, token_text) VALUES (?, ?)", t.ID, t.Text); err != nil {
|
|
return 0, fmt.Errorf("保存 token 失败: %w", err)
|
|
}
|
|
if _, err := tx.Exec("INSERT OR IGNORE INTO memory_tokens (memory_id, token_id) VALUES (?, ?)", m.ID, t.ID); err != nil {
|
|
return 0, fmt.Errorf("建立 token 关联失败: %w", err)
|
|
}
|
|
}
|
|
}
|
|
if err := tx.Commit(); err != nil {
|
|
return 0, fmt.Errorf("提交事务失败: %w", err)
|
|
}
|
|
return len(memories), nil
|
|
}
|
|
|
|
// 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 {
|
|
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
|
|
}
|