133 lines
3.0 KiB
Go
133 lines
3.0 KiB
Go
package memory
|
|
|
|
import (
|
|
"database/sql"
|
|
"fmt"
|
|
"os"
|
|
"path/filepath"
|
|
"time"
|
|
|
|
"laodingbot/internal/logger"
|
|
|
|
_ "modernc.org/sqlite"
|
|
)
|
|
|
|
type Message struct {
|
|
ID int64
|
|
ChatID string
|
|
UserID string
|
|
Role string
|
|
Content string
|
|
CreatedAt time.Time
|
|
}
|
|
|
|
type SQLiteStore struct {
|
|
db *sql.DB
|
|
log *logger.Logger
|
|
}
|
|
|
|
func NewSQLiteStore(path string, log *logger.Logger) (*SQLiteStore, error) {
|
|
abs, err := filepath.Abs(path)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if err := os.MkdirAll(filepath.Dir(abs), 0o755); err != nil {
|
|
return nil, err
|
|
}
|
|
db, err := sql.Open("sqlite", abs)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
store := &SQLiteStore{db: db, log: log}
|
|
if err := store.migrate(); err != nil {
|
|
_ = db.Close()
|
|
return nil, err
|
|
}
|
|
if log != nil {
|
|
log.Infof("sqlite store initialized path=%s", abs)
|
|
}
|
|
return store, nil
|
|
}
|
|
|
|
func (s *SQLiteStore) Close() error {
|
|
return s.db.Close()
|
|
}
|
|
|
|
func (s *SQLiteStore) SaveMessage(chatID, userID, role, content string) error {
|
|
if s.log != nil {
|
|
s.log.Debugf("save message chat_id=%s role=%s content_len=%d", chatID, role, len(content))
|
|
}
|
|
_, err := s.db.Exec(`
|
|
INSERT INTO messages(chat_id, user_id, role, content, created_at)
|
|
VALUES (?, ?, ?, ?, ?)
|
|
`, chatID, userID, role, content, time.Now().UTC())
|
|
if err != nil && s.log != nil {
|
|
s.log.Errorf("save message failed chat_id=%s role=%s err=%v", chatID, role, err)
|
|
}
|
|
return err
|
|
}
|
|
|
|
func (s *SQLiteStore) LoadRecent(chatID string, limit int) ([]Message, error) {
|
|
if limit <= 0 {
|
|
limit = 20
|
|
}
|
|
rows, err := s.db.Query(`
|
|
SELECT id, chat_id, user_id, role, content, created_at
|
|
FROM messages
|
|
WHERE chat_id = ?
|
|
ORDER BY id DESC
|
|
LIMIT ?
|
|
`, chatID, limit)
|
|
if err != nil {
|
|
if s.log != nil {
|
|
s.log.Errorf("load recent query failed chat_id=%s err=%v", chatID, err)
|
|
}
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
|
|
messages := make([]Message, 0, limit)
|
|
for rows.Next() {
|
|
var m Message
|
|
if err := rows.Scan(&m.ID, &m.ChatID, &m.UserID, &m.Role, &m.Content, &m.CreatedAt); err != nil {
|
|
return nil, err
|
|
}
|
|
messages = append(messages, m)
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
if s.log != nil {
|
|
s.log.Errorf("load recent row iteration failed chat_id=%s err=%v", chatID, err)
|
|
}
|
|
return nil, err
|
|
}
|
|
|
|
for left, right := 0, len(messages)-1; left < right; left, right = left+1, right-1 {
|
|
messages[left], messages[right] = messages[right], messages[left]
|
|
}
|
|
if s.log != nil {
|
|
s.log.Debugf("load recent success chat_id=%s count=%d", chatID, len(messages))
|
|
}
|
|
return messages, nil
|
|
}
|
|
|
|
func (s *SQLiteStore) migrate() error {
|
|
stmt := `
|
|
CREATE TABLE IF NOT EXISTS messages (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
chat_id TEXT NOT NULL,
|
|
user_id TEXT NOT NULL,
|
|
role TEXT NOT NULL,
|
|
content TEXT NOT NULL,
|
|
created_at TIMESTAMP NOT NULL
|
|
);
|
|
CREATE INDEX IF NOT EXISTS idx_messages_chat_id_id ON messages(chat_id, id);
|
|
`
|
|
if _, err := s.db.Exec(stmt); err != nil {
|
|
return fmt.Errorf("migrate schema: %w", err)
|
|
}
|
|
if s.log != nil {
|
|
s.log.Infof("sqlite schema migration completed")
|
|
}
|
|
return nil
|
|
}
|