2 files changed, 35 insertions, 3 deletions
diff --git a/internal/db/sql/messages.sql b/internal/db/sql/messages.sql
new file mode 100644
index 000000000..db5e192fc
--- /dev/null
+++ b/internal/db/sql/messages.sql
@@ -0,0 +1,30 @@
+-- name: GetMessage :one
+SELECT *
+FROM messages
+WHERE id = ? LIMIT 1;
+
+-- name: ListMessagesBySession :many
+SELECT *
+FROM messages
+WHERE session_id = ?
+ORDER BY created_at ASC;
+
+-- name: CreateMessage :one
+INSERT INTO messages (
+    id,
+    session_id,
+    message_data,
+    created_at,
+    updated_at
+) VALUES (
+    ?, ?, ?, strftime('%s', 'now'), strftime('%s', 'now')
+)
+RETURNING *;
+
+-- name: DeleteMessage :exec
+DELETE FROM messages
+WHERE id = ?;
+
+-- name: DeleteSessionMessages :exec
+DELETE FROM messages
+WHERE session_id = ?;
diff --git a/internal/db/sql/sessions.sql b/internal/db/sql/sessions.sql
index e90e5e328..2be8b7ccc 100644
--- a/internal/db/sql/sessions.sql
+++ b/internal/db/sql/sessions.sql
@@ -1,10 +1,10 @@
--- sqlfluff:dialect:sqlite
 -- name: CreateSession :one
 INSERT INTO sessions (
     id,
     title,
     message_count,
-    tokens,
+    prompt_tokens,
+    completion_tokens,
     cost,
     updated_at,
     created_at
@@ -14,6 +14,7 @@ INSERT INTO sessions (
     ?,
     ?,
     ?,
+    ?,
     strftime('%s', 'now'),
     strftime('%s', 'now')
 ) RETURNING *;
@@ -32,7 +33,8 @@ ORDER BY created_at DESC;
 UPDATE sessions
 SET
     title = ?,
-    tokens = ?,
+    prompt_tokens = ?,
+    completion_tokens = ?,
     cost = ?
 WHERE id = ?
 RETURNING *;