feat: provider manager, cancellation, session summaries (store+agent)
This commit is contained in:
parent
406a523d54
commit
f1b2093b22
10 changed files with 883 additions and 37 deletions
|
|
@ -18,8 +18,11 @@ import (
|
|||
var ErrNotFound = errors.New("not found")
|
||||
|
||||
type Session struct {
|
||||
ID string `json:"id"`
|
||||
Title string `json:"title"`
|
||||
ID string `json:"id"`
|
||||
Title string `json:"title"`
|
||||
// Provider/Model override the server-wide active provider for this session (empty = default).
|
||||
Provider string `json:"provider,omitempty"`
|
||||
Model string `json:"model,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
|
@ -45,15 +48,46 @@ CREATE TABLE IF NOT EXISTS messages(
|
|||
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);
|
||||
CREATE TABLE IF NOT EXISTS summaries(
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
|
||||
upto_message_id INTEGER NOT NULL, content TEXT NOT NULL, created_at INTEGER NOT NULL);
|
||||
CREATE INDEX IF NOT EXISTS idx_summaries_session ON summaries(session_id, id);
|
||||
CREATE TABLE IF NOT EXISTS settings(key TEXT PRIMARY KEY, value TEXT NOT NULL);
|
||||
`
|
||||
|
||||
// migrations add columns to databases created by older versions.
|
||||
var migrations = []struct{ table, column, ddl string }{
|
||||
{"sessions", "provider", "ALTER TABLE sessions ADD COLUMN provider TEXT NOT NULL DEFAULT ''"},
|
||||
{"sessions", "model", "ALTER TABLE sessions ADD COLUMN model TEXT NOT NULL DEFAULT ''"},
|
||||
{"messages", "cancelled", "ALTER TABLE messages ADD COLUMN cancelled INTEGER NOT NULL DEFAULT 0"},
|
||||
}
|
||||
|
||||
func migrate(db *sql.DB) error {
|
||||
if _, err := db.Exec(schema); err != nil {
|
||||
return err
|
||||
}
|
||||
for _, m := range migrations {
|
||||
var n int
|
||||
if err := db.QueryRow(`SELECT COUNT(*) FROM pragma_table_info(?) WHERE name=?`, m.table, m.column).Scan(&n); err != nil {
|
||||
return err
|
||||
}
|
||||
if n == 0 {
|
||||
if _, err := db.Exec(m.ddl); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
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 {
|
||||
if err := migrate(db); err != nil {
|
||||
db.Close()
|
||||
return nil, fmt.Errorf("migrate: %w", err)
|
||||
}
|
||||
|
|
@ -79,7 +113,7 @@ func (s *Store) CreateSession(ctx context.Context, title string) (Session, error
|
|||
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)
|
||||
err := s.db.QueryRowContext(ctx, `SELECT id,title,provider,model,created_at,updated_at FROM sessions WHERE id=?`, id).Scan(&ss.ID, &ss.Title, &ss.Provider, &ss.Model, &c, &u)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return ss, ErrNotFound
|
||||
}
|
||||
|
|
@ -88,7 +122,7 @@ func (s *Store) GetSession(ctx context.Context, id string) (Session, error) {
|
|||
}
|
||||
|
||||
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)
|
||||
rows, err := s.db.QueryContext(ctx, `SELECT id,title,provider,model,created_at,updated_at FROM sessions ORDER BY updated_at DESC LIMIT ?`, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
|
@ -97,7 +131,7 @@ func (s *Store) ListSessions(ctx context.Context, limit int) ([]Session, error)
|
|||
for rows.Next() {
|
||||
var ss Session
|
||||
var c, u int64
|
||||
if err := rows.Scan(&ss.ID, &ss.Title, &c, &u); err != nil {
|
||||
if err := rows.Scan(&ss.ID, &ss.Title, &ss.Provider, &ss.Model, &c, &u); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ss.CreatedAt, ss.UpdatedAt = time.UnixMilli(c).UTC(), time.UnixMilli(u).UTC()
|
||||
|
|
@ -138,8 +172,8 @@ func (s *Store) AppendMessage(ctx context.Context, sessionID string, m llm.Messa
|
|||
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)
|
||||
res, err := tx.ExecContext(ctx, `INSERT INTO messages(session_id,role,content,tool_calls,tool_call_id,name,cancelled,created_at) VALUES(?,?,?,?,?,?,?,?)`,
|
||||
sessionID, string(m.Role), m.Content, tc, m.ToolCallID, m.Name, m.Cancelled, now)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
|
@ -152,11 +186,16 @@ func (s *Store) AppendMessage(ctx context.Context, sessionID string, m llm.Messa
|
|||
|
||||
// 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) {
|
||||
return s.MessagesAfter(ctx, sessionID, 0, limit)
|
||||
}
|
||||
|
||||
// MessagesAfter returns the last `limit` messages with id > afterID (limit<=0: all), oldest first.
|
||||
func (s *Store) MessagesAfter(ctx context.Context, sessionID string, afterID int64, 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)
|
||||
rows, err := s.db.QueryContext(ctx, `SELECT * FROM (SELECT id,role,content,tool_calls,tool_call_id,name,cancelled,created_at
|
||||
FROM messages WHERE session_id=? AND id>? ORDER BY id DESC LIMIT ?) ORDER BY id ASC`, sessionID, afterID, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
|
@ -166,7 +205,7 @@ func (s *Store) Messages(ctx context.Context, sessionID string, limit int) ([]St
|
|||
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 {
|
||||
if err := rows.Scan(&m.ID, &role, &m.Content, &tc, &m.ToolCallID, &m.Name, &m.Cancelled, &c); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
m.Role, m.CreatedAt = llm.Role(role), time.UnixMilli(c).UTC()
|
||||
|
|
@ -179,3 +218,58 @@ func (s *Store) Messages(ctx context.Context, sessionID string, limit int) ([]St
|
|||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
func (s *Store) SetSessionProvider(ctx context.Context, id, provider, model string) error {
|
||||
res, err := s.db.ExecContext(ctx, `UPDATE sessions SET provider=?, model=? WHERE id=?`, provider, model, id)
|
||||
return affected(res, err)
|
||||
}
|
||||
|
||||
// ---- summaries
|
||||
|
||||
type Summary struct {
|
||||
ID int64 `json:"id"`
|
||||
SessionID string `json:"session_id"`
|
||||
UptoMessageID int64 `json:"upto_message_id"`
|
||||
Content string `json:"content"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
func (s *Store) AddSummary(ctx context.Context, sessionID string, upto int64, content string) (Summary, error) {
|
||||
now := time.Now().UTC()
|
||||
res, err := s.db.ExecContext(ctx, `INSERT INTO summaries(session_id,upto_message_id,content,created_at) VALUES(?,?,?,?)`,
|
||||
sessionID, upto, content, now.UnixMilli())
|
||||
if err != nil {
|
||||
return Summary{}, err
|
||||
}
|
||||
id, _ := res.LastInsertId()
|
||||
return Summary{ID: id, SessionID: sessionID, UptoMessageID: upto, Content: content, CreatedAt: now}, nil
|
||||
}
|
||||
|
||||
// LatestSummary returns the newest summary of a session or ErrNotFound.
|
||||
func (s *Store) LatestSummary(ctx context.Context, sessionID string) (Summary, error) {
|
||||
var sm Summary
|
||||
var c int64
|
||||
err := s.db.QueryRowContext(ctx, `SELECT id,session_id,upto_message_id,content,created_at FROM summaries
|
||||
WHERE session_id=? ORDER BY id DESC LIMIT 1`, sessionID).Scan(&sm.ID, &sm.SessionID, &sm.UptoMessageID, &sm.Content, &c)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return sm, ErrNotFound
|
||||
}
|
||||
sm.CreatedAt = time.UnixMilli(c).UTC()
|
||||
return sm, err
|
||||
}
|
||||
|
||||
// ---- settings (key/value)
|
||||
|
||||
func (s *Store) GetSetting(ctx context.Context, key string) (string, error) {
|
||||
var v string
|
||||
err := s.db.QueryRowContext(ctx, `SELECT value FROM settings WHERE key=?`, key).Scan(&v)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return "", ErrNotFound
|
||||
}
|
||||
return v, err
|
||||
}
|
||||
|
||||
func (s *Store) SetSetting(ctx context.Context, key, value string) error {
|
||||
_, err := s.db.ExecContext(ctx, `INSERT INTO settings(key,value) VALUES(?,?) ON CONFLICT(key) DO UPDATE SET value=excluded.value`, key, value)
|
||||
return err
|
||||
}
|
||||
|
|
|
|||
|
|
@ -48,3 +48,43 @@ func TestStoreRoundTrip(t *testing.T) {
|
|||
t.Fatal("cascade delete failed")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSummariesSettingsProviderAndCancelled(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, _ := s.CreateSession(ctx, "x")
|
||||
if _, err := s.LatestSummary(ctx, ss.ID); !errors.Is(err, ErrNotFound) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
id1, _ := s.AppendMessage(ctx, ss.ID, llm.Message{Role: llm.RoleUser, Content: "a"})
|
||||
s.AppendMessage(ctx, ss.ID, llm.Message{Role: llm.RoleAssistant, Content: "part", Cancelled: true})
|
||||
s.AddSummary(ctx, ss.ID, id1, "old")
|
||||
s.AddSummary(ctx, ss.ID, id1, "new")
|
||||
sm, _ := s.LatestSummary(ctx, ss.ID)
|
||||
if sm.Content != "new" || sm.UptoMessageID != id1 {
|
||||
t.Fatalf("%+v", sm)
|
||||
}
|
||||
after, _ := s.MessagesAfter(ctx, ss.ID, id1, 0)
|
||||
if len(after) != 1 || !after[0].Cancelled || after[0].Content != "part" {
|
||||
t.Fatalf("%+v", after)
|
||||
}
|
||||
if all, _ := s.Messages(ctx, ss.ID, 0); len(all) != 2 {
|
||||
t.Fatal("original messages must be kept")
|
||||
}
|
||||
if err := s.SetSessionProvider(ctx, ss.ID, "dahl", "m"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, _ := s.GetSession(ctx, ss.ID)
|
||||
if got.Provider != "dahl" || got.Model != "m" {
|
||||
t.Fatalf("%+v", got)
|
||||
}
|
||||
s.SetSetting(ctx, "k", "v1")
|
||||
s.SetSetting(ctx, "k", "v2")
|
||||
if v, _ := s.GetSetting(ctx, "k"); v != "v2" {
|
||||
t.Fatal(v)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue