summaryrefslogtreecommitdiffhomepage
path: root/internal/db
diff options
context:
space:
mode:
Diffstat (limited to 'internal/db')
-rw-r--r--internal/db/connect.go4
-rw-r--r--internal/db/db.go10
-rw-r--r--internal/db/messages.sql.go89
-rw-r--r--internal/db/migrations/000001_initial.up.sql8
-rw-r--r--internal/db/models.go36
-rw-r--r--internal/db/querier.go1
-rw-r--r--internal/db/sessions.sql.go30
-rw-r--r--internal/db/sql/messages.sql19
-rw-r--r--internal/db/sql/sessions.sql3
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