diff options
Diffstat (limited to 'internal/db')
| -rw-r--r-- | internal/db/connect.go | 4 | ||||
| -rw-r--r-- | internal/db/db.go | 10 | ||||
| -rw-r--r-- | internal/db/messages.sql.go | 89 | ||||
| -rw-r--r-- | internal/db/migrations/000001_initial.up.sql | 8 | ||||
| -rw-r--r-- | internal/db/models.go | 36 | ||||
| -rw-r--r-- | internal/db/querier.go | 1 | ||||
| -rw-r--r-- | internal/db/sessions.sql.go | 30 | ||||
| -rw-r--r-- | internal/db/sql/messages.sql | 19 | ||||
| -rw-r--r-- | internal/db/sql/sessions.sql | 3 |
9 files changed, 160 insertions, 40 deletions
diff --git a/internal/db/connect.go b/internal/db/connect.go index 50a30d0bc..aed04b986 100644 --- a/internal/db/connect.go +++ b/internal/db/connect.go @@ -12,14 +12,14 @@ import ( "github.com/golang-migrate/migrate/v4/database/sqlite3" _ "github.com/mattn/go-sqlite3" + "github.com/kujtimiihoxha/termai/internal/config" "github.com/kujtimiihoxha/termai/internal/logging" - "github.com/spf13/viper" ) var log = logging.Get() func Connect() (*sql.DB, error) { - dataDir := viper.GetString("data.dir") + dataDir := config.Get().Data.Directory if dataDir == "" { return nil, fmt.Errorf("data.dir is not set") } diff --git a/internal/db/db.go b/internal/db/db.go index e882106c4..75f626013 100644 --- a/internal/db/db.go +++ b/internal/db/db.go @@ -51,6 +51,9 @@ func Prepare(ctx context.Context, db DBTX) (*Queries, error) { if q.listSessionsStmt, err = db.PrepareContext(ctx, listSessions); err != nil { return nil, fmt.Errorf("error preparing query ListSessions: %w", err) } + if q.updateMessageStmt, err = db.PrepareContext(ctx, updateMessage); err != nil { + return nil, fmt.Errorf("error preparing query UpdateMessage: %w", err) + } if q.updateSessionStmt, err = db.PrepareContext(ctx, updateSession); err != nil { return nil, fmt.Errorf("error preparing query UpdateSession: %w", err) } @@ -104,6 +107,11 @@ func (q *Queries) Close() error { err = fmt.Errorf("error closing listSessionsStmt: %w", cerr) } } + if q.updateMessageStmt != nil { + if cerr := q.updateMessageStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing updateMessageStmt: %w", cerr) + } + } if q.updateSessionStmt != nil { if cerr := q.updateSessionStmt.Close(); cerr != nil { err = fmt.Errorf("error closing updateSessionStmt: %w", cerr) @@ -157,6 +165,7 @@ type Queries struct { getSessionByIDStmt *sql.Stmt listMessagesBySessionStmt *sql.Stmt listSessionsStmt *sql.Stmt + updateMessageStmt *sql.Stmt updateSessionStmt *sql.Stmt } @@ -173,6 +182,7 @@ func (q *Queries) WithTx(tx *sql.Tx) *Queries { getSessionByIDStmt: q.getSessionByIDStmt, listMessagesBySessionStmt: q.listMessagesBySessionStmt, listSessionsStmt: q.listSessionsStmt, + updateMessageStmt: q.updateMessageStmt, updateSessionStmt: q.updateSessionStmt, } } diff --git a/internal/db/messages.sql.go b/internal/db/messages.sql.go index d0f69458f..3f2846740 100644 --- a/internal/db/messages.sql.go +++ b/internal/db/messages.sql.go @@ -7,34 +7,56 @@ package db import ( "context" + "database/sql" ) const createMessage = `-- name: CreateMessage :one INSERT INTO messages ( id, session_id, - message_data, + role, + finished, + content, + tool_calls, + tool_results, created_at, updated_at ) VALUES ( - ?, ?, ?, strftime('%s', 'now'), strftime('%s', 'now') + ?, ?, ?, ?, ?, ?, ?, strftime('%s', 'now'), strftime('%s', 'now') ) -RETURNING id, session_id, message_data, created_at, updated_at +RETURNING id, session_id, role, content, thinking, finished, tool_calls, tool_results, created_at, updated_at ` type CreateMessageParams struct { - ID string `json:"id"` - SessionID string `json:"session_id"` - MessageData string `json:"message_data"` + ID string `json:"id"` + SessionID string `json:"session_id"` + Role string `json:"role"` + Finished bool `json:"finished"` + Content string `json:"content"` + ToolCalls sql.NullString `json:"tool_calls"` + ToolResults sql.NullString `json:"tool_results"` } func (q *Queries) CreateMessage(ctx context.Context, arg CreateMessageParams) (Message, error) { - row := q.queryRow(ctx, q.createMessageStmt, createMessage, arg.ID, arg.SessionID, arg.MessageData) + row := q.queryRow(ctx, q.createMessageStmt, createMessage, + arg.ID, + arg.SessionID, + arg.Role, + arg.Finished, + arg.Content, + arg.ToolCalls, + arg.ToolResults, + ) var i Message err := row.Scan( &i.ID, &i.SessionID, - &i.MessageData, + &i.Role, + &i.Content, + &i.Thinking, + &i.Finished, + &i.ToolCalls, + &i.ToolResults, &i.CreatedAt, &i.UpdatedAt, ) @@ -62,7 +84,7 @@ func (q *Queries) DeleteSessionMessages(ctx context.Context, sessionID string) e } const getMessage = `-- name: GetMessage :one -SELECT id, session_id, message_data, created_at, updated_at +SELECT id, session_id, role, content, thinking, finished, tool_calls, tool_results, created_at, updated_at FROM messages WHERE id = ? LIMIT 1 ` @@ -73,7 +95,12 @@ func (q *Queries) GetMessage(ctx context.Context, id string) (Message, error) { err := row.Scan( &i.ID, &i.SessionID, - &i.MessageData, + &i.Role, + &i.Content, + &i.Thinking, + &i.Finished, + &i.ToolCalls, + &i.ToolResults, &i.CreatedAt, &i.UpdatedAt, ) @@ -81,7 +108,7 @@ func (q *Queries) GetMessage(ctx context.Context, id string) (Message, error) { } const listMessagesBySession = `-- name: ListMessagesBySession :many -SELECT id, session_id, message_data, created_at, updated_at +SELECT id, session_id, role, content, thinking, finished, tool_calls, tool_results, created_at, updated_at FROM messages WHERE session_id = ? ORDER BY created_at ASC @@ -99,7 +126,12 @@ func (q *Queries) ListMessagesBySession(ctx context.Context, sessionID string) ( if err := rows.Scan( &i.ID, &i.SessionID, - &i.MessageData, + &i.Role, + &i.Content, + &i.Thinking, + &i.Finished, + &i.ToolCalls, + &i.ToolResults, &i.CreatedAt, &i.UpdatedAt, ); err != nil { @@ -115,3 +147,36 @@ func (q *Queries) ListMessagesBySession(ctx context.Context, sessionID string) ( } return items, nil } + +const updateMessage = `-- name: UpdateMessage :exec +UPDATE messages +SET + content = ?, + thinking = ?, + tool_calls = ?, + tool_results = ?, + finished = ?, + updated_at = strftime('%s', 'now') +WHERE id = ? +` + +type UpdateMessageParams struct { + Content string `json:"content"` + Thinking string `json:"thinking"` + ToolCalls sql.NullString `json:"tool_calls"` + ToolResults sql.NullString `json:"tool_results"` + Finished bool `json:"finished"` + ID string `json:"id"` +} + +func (q *Queries) UpdateMessage(ctx context.Context, arg UpdateMessageParams) error { + _, err := q.exec(ctx, q.updateMessageStmt, updateMessage, + arg.Content, + arg.Thinking, + arg.ToolCalls, + arg.ToolResults, + arg.Finished, + arg.ID, + ) + return err +} diff --git a/internal/db/migrations/000001_initial.up.sql b/internal/db/migrations/000001_initial.up.sql index b87592bb5..d4a5e39fe 100644 --- a/internal/db/migrations/000001_initial.up.sql +++ b/internal/db/migrations/000001_initial.up.sql @@ -1,6 +1,7 @@ -- Sessions CREATE TABLE IF NOT EXISTS sessions ( id TEXT PRIMARY KEY, + parent_session_id TEXT, title TEXT NOT NULL, message_count INTEGER NOT NULL DEFAULT 0 CHECK (message_count >= 0), prompt_tokens INTEGER NOT NULL DEFAULT 0 CHECK (prompt_tokens >= 0), @@ -21,7 +22,12 @@ END; CREATE TABLE IF NOT EXISTS messages ( id TEXT PRIMARY KEY, session_id TEXT NOT NULL, - message_data TEXT NOT NULL, -- JSON string of message content + role TEXT NOT NULL, + content TEXT NOT NULL, + thinking Text NOT NULL DEFAULT '', + finished BOOLEAN NOT NULL DEFAULT 0, + tool_calls TEXT, + tool_results TEXT, created_at INTEGER NOT NULL, -- Unix timestamp in milliseconds updated_at INTEGER NOT NULL, -- Unix timestamp in milliseconds FOREIGN KEY (session_id) REFERENCES sessions (id) ON DELETE CASCADE diff --git a/internal/db/models.go b/internal/db/models.go index de9cd3cdb..fc3140a13 100644 --- a/internal/db/models.go +++ b/internal/db/models.go @@ -4,21 +4,31 @@ package db +import ( + "database/sql" +) + type Message struct { - ID string `json:"id"` - SessionID string `json:"session_id"` - MessageData string `json:"message_data"` - CreatedAt int64 `json:"created_at"` - UpdatedAt int64 `json:"updated_at"` + ID string `json:"id"` + SessionID string `json:"session_id"` + Role string `json:"role"` + Content string `json:"content"` + Thinking string `json:"thinking"` + Finished bool `json:"finished"` + ToolCalls sql.NullString `json:"tool_calls"` + ToolResults sql.NullString `json:"tool_results"` + CreatedAt int64 `json:"created_at"` + UpdatedAt int64 `json:"updated_at"` } type Session struct { - ID string `json:"id"` - Title string `json:"title"` - MessageCount int64 `json:"message_count"` - PromptTokens int64 `json:"prompt_tokens"` - CompletionTokens int64 `json:"completion_tokens"` - Cost float64 `json:"cost"` - UpdatedAt int64 `json:"updated_at"` - CreatedAt int64 `json:"created_at"` + ID string `json:"id"` + ParentSessionID sql.NullString `json:"parent_session_id"` + Title string `json:"title"` + MessageCount int64 `json:"message_count"` + PromptTokens int64 `json:"prompt_tokens"` + CompletionTokens int64 `json:"completion_tokens"` + Cost float64 `json:"cost"` + UpdatedAt int64 `json:"updated_at"` + CreatedAt int64 `json:"created_at"` } diff --git a/internal/db/querier.go b/internal/db/querier.go index 626d206a4..c9d73ec39 100644 --- a/internal/db/querier.go +++ b/internal/db/querier.go @@ -18,6 +18,7 @@ type Querier interface { GetSessionByID(ctx context.Context, id string) (Session, error) ListMessagesBySession(ctx context.Context, sessionID string) ([]Message, error) ListSessions(ctx context.Context) ([]Session, error) + UpdateMessage(ctx context.Context, arg UpdateMessageParams) error UpdateSession(ctx context.Context, arg UpdateSessionParams) (Session, error) } diff --git a/internal/db/sessions.sql.go b/internal/db/sessions.sql.go index f3ee4ff42..18d70c3db 100644 --- a/internal/db/sessions.sql.go +++ b/internal/db/sessions.sql.go @@ -7,11 +7,13 @@ package db import ( "context" + "database/sql" ) const createSession = `-- name: CreateSession :one INSERT INTO sessions ( id, + parent_session_id, title, message_count, prompt_tokens, @@ -26,23 +28,26 @@ INSERT INTO sessions ( ?, ?, ?, + ?, strftime('%s', 'now'), strftime('%s', 'now') -) RETURNING id, title, message_count, prompt_tokens, completion_tokens, cost, updated_at, created_at +) RETURNING id, parent_session_id, title, message_count, prompt_tokens, completion_tokens, cost, updated_at, created_at ` type CreateSessionParams struct { - ID string `json:"id"` - Title string `json:"title"` - MessageCount int64 `json:"message_count"` - PromptTokens int64 `json:"prompt_tokens"` - CompletionTokens int64 `json:"completion_tokens"` - Cost float64 `json:"cost"` + ID string `json:"id"` + ParentSessionID sql.NullString `json:"parent_session_id"` + Title string `json:"title"` + MessageCount int64 `json:"message_count"` + PromptTokens int64 `json:"prompt_tokens"` + CompletionTokens int64 `json:"completion_tokens"` + Cost float64 `json:"cost"` } func (q *Queries) CreateSession(ctx context.Context, arg CreateSessionParams) (Session, error) { row := q.queryRow(ctx, q.createSessionStmt, createSession, arg.ID, + arg.ParentSessionID, arg.Title, arg.MessageCount, arg.PromptTokens, @@ -52,6 +57,7 @@ func (q *Queries) CreateSession(ctx context.Context, arg CreateSessionParams) (S var i Session err := row.Scan( &i.ID, + &i.ParentSessionID, &i.Title, &i.MessageCount, &i.PromptTokens, @@ -74,7 +80,7 @@ func (q *Queries) DeleteSession(ctx context.Context, id string) error { } const getSessionByID = `-- name: GetSessionByID :one -SELECT id, title, message_count, prompt_tokens, completion_tokens, cost, updated_at, created_at +SELECT id, parent_session_id, title, message_count, prompt_tokens, completion_tokens, cost, updated_at, created_at FROM sessions WHERE id = ? LIMIT 1 ` @@ -84,6 +90,7 @@ func (q *Queries) GetSessionByID(ctx context.Context, id string) (Session, error var i Session err := row.Scan( &i.ID, + &i.ParentSessionID, &i.Title, &i.MessageCount, &i.PromptTokens, @@ -96,8 +103,9 @@ func (q *Queries) GetSessionByID(ctx context.Context, id string) (Session, error } const listSessions = `-- name: ListSessions :many -SELECT id, title, message_count, prompt_tokens, completion_tokens, cost, updated_at, created_at +SELECT id, parent_session_id, title, message_count, prompt_tokens, completion_tokens, cost, updated_at, created_at FROM sessions +WHERE parent_session_id is NULL ORDER BY created_at DESC ` @@ -112,6 +120,7 @@ func (q *Queries) ListSessions(ctx context.Context) ([]Session, error) { var i Session if err := rows.Scan( &i.ID, + &i.ParentSessionID, &i.Title, &i.MessageCount, &i.PromptTokens, @@ -141,7 +150,7 @@ SET completion_tokens = ?, cost = ? WHERE id = ? -RETURNING id, title, message_count, prompt_tokens, completion_tokens, cost, updated_at, created_at +RETURNING id, parent_session_id, title, message_count, prompt_tokens, completion_tokens, cost, updated_at, created_at ` type UpdateSessionParams struct { @@ -163,6 +172,7 @@ func (q *Queries) UpdateSession(ctx context.Context, arg UpdateSessionParams) (S var i Session err := row.Scan( &i.ID, + &i.ParentSessionID, &i.Title, &i.MessageCount, &i.PromptTokens, diff --git a/internal/db/sql/messages.sql b/internal/db/sql/messages.sql index db5e192fc..0674e62c1 100644 --- a/internal/db/sql/messages.sql +++ b/internal/db/sql/messages.sql @@ -13,14 +13,29 @@ ORDER BY created_at ASC; INSERT INTO messages ( id, session_id, - message_data, + role, + finished, + content, + tool_calls, + tool_results, created_at, updated_at ) VALUES ( - ?, ?, ?, strftime('%s', 'now'), strftime('%s', 'now') + ?, ?, ?, ?, ?, ?, ?, strftime('%s', 'now'), strftime('%s', 'now') ) RETURNING *; +-- name: UpdateMessage :exec +UPDATE messages +SET + content = ?, + thinking = ?, + tool_calls = ?, + tool_results = ?, + finished = ?, + updated_at = strftime('%s', 'now') +WHERE id = ?; + -- name: DeleteMessage :exec DELETE FROM messages WHERE id = ?; diff --git a/internal/db/sql/sessions.sql b/internal/db/sql/sessions.sql index 2be8b7ccc..f065b5f56 100644 --- a/internal/db/sql/sessions.sql +++ b/internal/db/sql/sessions.sql @@ -1,6 +1,7 @@ -- name: CreateSession :one INSERT INTO sessions ( id, + parent_session_id, title, message_count, prompt_tokens, @@ -15,6 +16,7 @@ INSERT INTO sessions ( ?, ?, ?, + ?, strftime('%s', 'now'), strftime('%s', 'now') ) RETURNING *; @@ -27,6 +29,7 @@ WHERE id = ? LIMIT 1; -- name: ListSessions :many SELECT * FROM sessions +WHERE parent_session_id is NULL ORDER BY created_at DESC; -- name: UpdateSession :one |
