feat: SQLite session store and tool-calling agent loop
This commit is contained in:
OpenClaw Dev
parent
c1c7b0e005
commit
d7e138068e
6 changed files with 513 additions and 2 deletions
181
internal/store/store.go
Normal file
181
internal/store/store.go
Normal file
|
|
@ -0,0 +1,181 @@
|
|||
// Package store persists sessions and messages in SQLite (pure Go driver, no CGO).
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"database/sql"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/crypto/openclaw/internal/llm"
|
||||
_ "modernc.org/sqlite"
|
||||
)
|
||||
|
||||
var ErrNotFound = errors.New("not found")
|
||||
|
||||
type Session struct {
|
||||
ID string `json:"id"`
|
||||
Title string `json:"title"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
type StoredMessage struct {
|
||||
llm.Message
|
||||
ID int64 `json:"id"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
type Store struct{ db *sql.DB }
|
||||
|
||||
const schema = `
|
||||
PRAGMA journal_mode=WAL;
|
||||
PRAGMA foreign_keys=ON;
|
||||
CREATE TABLE IF NOT EXISTS sessions(
|
||||
id TEXT PRIMARY KEY, title TEXT NOT NULL DEFAULT '',
|
||||
created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL);
|
||||
CREATE TABLE IF NOT EXISTS messages(
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
|
||||
role TEXT NOT NULL, content TEXT NOT NULL DEFAULT '',
|
||||
tool_calls TEXT NOT NULL DEFAULT '', tool_call_id TEXT NOT NULL DEFAULT '',
|
||||
name TEXT NOT NULL DEFAULT '', created_at INTEGER NOT NULL);
|
||||
CREATE INDEX IF NOT EXISTS idx_messages_session ON messages(session_id, id);
|
||||
`
|
||||
|
||||
func Open(path string) (*Store, error) {
|
||||
db, err := sql.Open("sqlite", "file:"+path+"?_pragma=busy_timeout(5000)")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
db.SetMaxOpenConns(1) // SQLite: serialize writers, avoids SQLITE_BUSY
|
||||
if _, err := db.Exec(schema); err != nil {
|
||||
db.Close()
|
||||
return nil, fmt.Errorf("migrate: %w", err)
|
||||
}
|
||||
return &Store{db: db}, nil
|
||||
}
|
||||
|
||||
func (s *Store) Close() error { return s.db.Close() }
|
||||
|
||||
func newID() string {
|
||||
b := make([]byte, 12)
|
||||
rand.Read(b)
|
||||
return hex.EncodeToString(b)
|
||||
}
|
||||
|
||||
func (s *Store) CreateSession(ctx context.Context, title string) (Session, error) {
|
||||
now := time.Now().UTC()
|
||||
ss := Session{ID: newID(), Title: title, CreatedAt: now, UpdatedAt: now}
|
||||
_, err := s.db.ExecContext(ctx, `INSERT INTO sessions(id,title,created_at,updated_at) VALUES(?,?,?,?)`,
|
||||
ss.ID, ss.Title, now.UnixMilli(), now.UnixMilli())
|
||||
return ss, err
|
||||
}
|
||||
|
||||
func (s *Store) GetSession(ctx context.Context, id string) (Session, error) {
|
||||
var ss Session
|
||||
var c, u int64
|
||||
err := s.db.QueryRowContext(ctx, `SELECT id,title,created_at,updated_at FROM sessions WHERE id=?`, id).Scan(&ss.ID, &ss.Title, &c, &u)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return ss, ErrNotFound
|
||||
}
|
||||
ss.CreatedAt, ss.UpdatedAt = time.UnixMilli(c).UTC(), time.UnixMilli(u).UTC()
|
||||
return ss, err
|
||||
}
|
||||
|
||||
func (s *Store) ListSessions(ctx context.Context, limit int) ([]Session, error) {
|
||||
rows, err := s.db.QueryContext(ctx, `SELECT id,title,created_at,updated_at FROM sessions ORDER BY updated_at DESC LIMIT ?`, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
out := []Session{}
|
||||
for rows.Next() {
|
||||
var ss Session
|
||||
var c, u int64
|
||||
if err := rows.Scan(&ss.ID, &ss.Title, &c, &u); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ss.CreatedAt, ss.UpdatedAt = time.UnixMilli(c).UTC(), time.UnixMilli(u).UTC()
|
||||
out = append(out, ss)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
func (s *Store) RenameSession(ctx context.Context, id, title string) error {
|
||||
res, err := s.db.ExecContext(ctx, `UPDATE sessions SET title=? WHERE id=?`, title, id)
|
||||
return affected(res, err)
|
||||
}
|
||||
|
||||
func (s *Store) DeleteSession(ctx context.Context, id string) error {
|
||||
res, err := s.db.ExecContext(ctx, `DELETE FROM sessions WHERE id=?`, id)
|
||||
return affected(res, err)
|
||||
}
|
||||
|
||||
func affected(res sql.Result, err error) error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n, _ := res.RowsAffected(); n == 0 {
|
||||
return ErrNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Store) AppendMessage(ctx context.Context, sessionID string, m llm.Message) (int64, error) {
|
||||
tc := ""
|
||||
if len(m.ToolCalls) > 0 {
|
||||
b, _ := json.Marshal(m.ToolCalls)
|
||||
tc = string(b)
|
||||
}
|
||||
now := time.Now().UTC().UnixMilli()
|
||||
tx, err := s.db.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
res, err := tx.ExecContext(ctx, `INSERT INTO messages(session_id,role,content,tool_calls,tool_call_id,name,created_at) VALUES(?,?,?,?,?,?,?)`,
|
||||
sessionID, string(m.Role), m.Content, tc, m.ToolCallID, m.Name, now)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if _, err := tx.ExecContext(ctx, `UPDATE sessions SET updated_at=? WHERE id=?`, now, sessionID); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
id, _ := res.LastInsertId()
|
||||
return id, tx.Commit()
|
||||
}
|
||||
|
||||
// Messages returns the last `limit` messages of a session in chronological order (limit<=0: all).
|
||||
func (s *Store) Messages(ctx context.Context, sessionID string, limit int) ([]StoredMessage, error) {
|
||||
if limit <= 0 {
|
||||
limit = -1
|
||||
}
|
||||
rows, err := s.db.QueryContext(ctx, `SELECT * FROM (SELECT id,role,content,tool_calls,tool_call_id,name,created_at
|
||||
FROM messages WHERE session_id=? ORDER BY id DESC LIMIT ?) ORDER BY id ASC`, sessionID, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
out := []StoredMessage{}
|
||||
for rows.Next() {
|
||||
var m StoredMessage
|
||||
var role, tc string
|
||||
var c int64
|
||||
if err := rows.Scan(&m.ID, &role, &m.Content, &tc, &m.ToolCallID, &m.Name, &c); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
m.Role, m.CreatedAt = llm.Role(role), time.UnixMilli(c).UTC()
|
||||
if tc != "" {
|
||||
if err := json.Unmarshal([]byte(tc), &m.ToolCalls); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
out = append(out, m)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
50
internal/store/store_test.go
Normal file
50
internal/store/store_test.go
Normal file
|
|
@ -0,0 +1,50 @@
|
|||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/crypto/openclaw/internal/llm"
|
||||
)
|
||||
|
||||
func TestStoreRoundTrip(t *testing.T) {
|
||||
s, err := Open(filepath.Join(t.TempDir(), "t.db"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.Close()
|
||||
ctx := context.Background()
|
||||
ss, err := s.CreateSession(ctx, "first")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s.AppendMessage(ctx, ss.ID, llm.Message{Role: llm.RoleUser, Content: "hi"})
|
||||
s.AppendMessage(ctx, ss.ID, llm.Message{Role: llm.RoleAssistant, ToolCalls: []llm.ToolCall{{ID: "c1", Name: "time_now", Arguments: "{}"}}})
|
||||
s.AppendMessage(ctx, ss.ID, llm.Message{Role: llm.RoleTool, ToolCallID: "c1", Name: "time_now", Content: "now"})
|
||||
all, _ := s.Messages(ctx, ss.ID, 0)
|
||||
if len(all) != 3 || all[1].ToolCalls[0].Name != "time_now" || all[2].ToolCallID != "c1" {
|
||||
t.Fatalf("%+v", all)
|
||||
}
|
||||
last, _ := s.Messages(ctx, ss.ID, 2)
|
||||
if len(last) != 2 || last[0].Role != llm.RoleAssistant {
|
||||
t.Fatalf("limit: %+v", last)
|
||||
}
|
||||
if err := s.RenameSession(ctx, ss.ID, "renamed"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
list, _ := s.ListSessions(ctx, 10)
|
||||
if len(list) != 1 || list[0].Title != "renamed" {
|
||||
t.Fatalf("%+v", list)
|
||||
}
|
||||
if err := s.DeleteSession(ctx, ss.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := s.GetSession(ctx, ss.ID); !errors.Is(err, ErrNotFound) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if m, _ := s.Messages(ctx, ss.ID, 0); len(m) != 0 {
|
||||
t.Fatal("cascade delete failed")
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue