127 lines
3.3 KiB
Go
127 lines
3.3 KiB
Go
package store
|
||
|
||
import (
|
||
"database/sql"
|
||
"sync"
|
||
|
||
_ "modernc.org/sqlite"
|
||
)
|
||
|
||
type Store struct {
|
||
db *sql.DB
|
||
mu sync.Mutex
|
||
}
|
||
|
||
func New(dbPath string) (*Store, error) {
|
||
db, err := sql.Open("sqlite", dbPath+"?_pragma=busy_timeout(5000)&_pragma=journal_mode(WAL)")
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
s := &Store{db: db}
|
||
if err := s.migrate(); err != nil {
|
||
db.Close()
|
||
return nil, err
|
||
}
|
||
return s, nil
|
||
}
|
||
|
||
func (s *Store) Close() error {
|
||
return s.db.Close()
|
||
}
|
||
|
||
func (s *Store) migrate() error {
|
||
_, err := s.db.Exec(`
|
||
CREATE TABLE IF NOT EXISTS messages (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
room_id TEXT NOT NULL,
|
||
agent TEXT NOT NULL,
|
||
role TEXT NOT NULL,
|
||
content TEXT NOT NULL DEFAULT '',
|
||
filename TEXT,
|
||
title TEXT,
|
||
group_id INTEGER,
|
||
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||
);
|
||
CREATE INDEX IF NOT EXISTS idx_msg_room ON messages(room_id, created_at);
|
||
CREATE INDEX IF NOT EXISTS idx_msg_group ON messages(group_id);
|
||
|
||
CREATE TABLE IF NOT EXISTS token_usage (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
room_id TEXT NOT NULL,
|
||
agent TEXT NOT NULL,
|
||
prompt_tokens INTEGER NOT NULL DEFAULT 0,
|
||
completion_tokens INTEGER NOT NULL DEFAULT 0,
|
||
total_tokens INTEGER NOT NULL DEFAULT 0,
|
||
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||
);
|
||
CREATE INDEX IF NOT EXISTS idx_tu_room ON token_usage(room_id);
|
||
|
||
CREATE TABLE IF NOT EXISTS schedules (
|
||
id TEXT PRIMARY KEY,
|
||
room_id TEXT NOT NULL,
|
||
cron TEXT NOT NULL,
|
||
message TEXT NOT NULL,
|
||
user_name TEXT NOT NULL DEFAULT 'scheduler',
|
||
enabled INTEGER NOT NULL DEFAULT 1,
|
||
once INTEGER NOT NULL DEFAULT 0,
|
||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||
last_run_at TEXT
|
||
);
|
||
CREATE INDEX IF NOT EXISTS idx_sch_room ON schedules(room_id);
|
||
|
||
CREATE TABLE IF NOT EXISTS file_versions (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
room_id TEXT NOT NULL,
|
||
filename TEXT NOT NULL,
|
||
content TEXT NOT NULL,
|
||
agent TEXT NOT NULL DEFAULT '',
|
||
version INTEGER NOT NULL DEFAULT 1,
|
||
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||
);
|
||
CREATE INDEX IF NOT EXISTS idx_fv_room_file ON file_versions(room_id, filename);
|
||
|
||
CREATE TABLE IF NOT EXISTS agent_configs (
|
||
room_id TEXT NOT NULL,
|
||
agent_name TEXT NOT NULL,
|
||
file_type TEXT NOT NULL,
|
||
content TEXT NOT NULL DEFAULT '',
|
||
updated_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||
PRIMARY KEY (room_id, agent_name, file_type)
|
||
);
|
||
CREATE INDEX IF NOT EXISTS idx_ac_room ON agent_configs(room_id);
|
||
`)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
|
||
// 安全添加 part_type 列（SQLite 不支持 ALTER TABLE ... IF NOT EXISTS）
|
||
var hasPartType bool
|
||
rows, err := s.db.Query(`PRAGMA table_info(messages)`)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer rows.Close()
|
||
for rows.Next() {
|
||
var cid int
|
||
var name, typ string
|
||
var notnull int
|
||
var dfltValue *string
|
||
var pk int
|
||
if err := rows.Scan(&cid, &name, &typ, ¬null, &dfltValue, &pk); err != nil {
|
||
return err
|
||
}
|
||
if name == "part_type" {
|
||
hasPartType = true
|
||
break
|
||
}
|
||
}
|
||
if !hasPartType {
|
||
_, err = s.db.Exec(`ALTER TABLE messages ADD COLUMN part_type TEXT NOT NULL DEFAULT 'text'`)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
}
|
||
|
||
return nil
|
||
}
|