add initial message handling

4 files changed, 37 insertions, 18 deletions

diff --git a/internal/db/messages.sql.go b/internal/db/messages.sql.go
index 4309db181..0555b4330 100644
--- a/internal/db/messages.sql.go
+++ b/internal/db/messages.sql.go
@@ -7,6 +7,7 @@ package db
 
 import (
 	"context"
+	"database/sql"
 )
 
 const createMessage = `-- name: CreateMessage :one
@@ -15,19 +16,21 @@ INSERT INTO messages (
     session_id,
     role,
     parts,
+    model,
     created_at,
     updated_at
 ) VALUES (
-    ?, ?, ?, ?, strftime('%s', 'now'), strftime('%s', 'now')
+    ?, ?, ?, ?, ?, strftime('%s', 'now'), strftime('%s', 'now')
 )
-RETURNING id, session_id, role, parts, created_at, updated_at
+RETURNING id, session_id, role, parts, model, created_at, updated_at, finished_at
 `
 
 type CreateMessageParams struct {
-	ID        string `json:"id"`
-	SessionID string `json:"session_id"`
-	Role      string `json:"role"`
-	Parts     string `json:"parts"`
+	ID        string         `json:"id"`
+	SessionID string         `json:"session_id"`
+	Role      string         `json:"role"`
+	Parts     string         `json:"parts"`
+	Model     sql.NullString `json:"model"`
 }
 
 func (q *Queries) CreateMessage(ctx context.Context, arg CreateMessageParams) (Message, error) {
@@ -36,6 +39,7 @@ func (q *Queries) CreateMessage(ctx context.Context, arg CreateMessageParams) (M
 		arg.SessionID,
 		arg.Role,
 		arg.Parts,
+		arg.Model,
 	)
 	var i Message
 	err := row.Scan(
@@ -43,8 +47,10 @@ func (q *Queries) CreateMessage(ctx context.Context, arg CreateMessageParams) (M
 		&i.SessionID,
 		&i.Role,
 		&i.Parts,
+		&i.Model,
 		&i.CreatedAt,
 		&i.UpdatedAt,
+		&i.FinishedAt,
 	)
 	return i, err
 }
@@ -70,7 +76,7 @@ func (q *Queries) DeleteSessionMessages(ctx context.Context, sessionID string) e
 }
 
 const getMessage = `-- name: GetMessage :one
-SELECT id, session_id, role, parts, created_at, updated_at
+SELECT id, session_id, role, parts, model, created_at, updated_at, finished_at
 FROM messages
 WHERE id = ? LIMIT 1
 `
@@ -83,14 +89,16 @@ func (q *Queries) GetMessage(ctx context.Context, id string) (Message, error) {
 		&i.SessionID,
 		&i.Role,
 		&i.Parts,
+		&i.Model,
 		&i.CreatedAt,
 		&i.UpdatedAt,
+		&i.FinishedAt,
 	)
 	return i, err
 }
 
 const listMessagesBySession = `-- name: ListMessagesBySession :many
-SELECT id, session_id, role, parts, created_at, updated_at
+SELECT id, session_id, role, parts, model, created_at, updated_at, finished_at
 FROM messages
 WHERE session_id = ?
 ORDER BY created_at ASC
@@ -110,8 +118,10 @@ func (q *Queries) ListMessagesBySession(ctx context.Context, sessionID string) (
 			&i.SessionID,
 			&i.Role,
 			&i.Parts,
+			&i.Model,
 			&i.CreatedAt,
 			&i.UpdatedAt,
+			&i.FinishedAt,
 		); err != nil {
 			return nil, err
 		}
@@ -130,16 +140,18 @@ const updateMessage = `-- name: UpdateMessage :exec
 UPDATE messages
 SET
     parts = ?,
+    finished_at = ?,
     updated_at = strftime('%s', 'now')
 WHERE id = ?
 `
 
 type UpdateMessageParams struct {
-	Parts string `json:"parts"`
-	ID    string `json:"id"`
+	Parts      string        `json:"parts"`
+	FinishedAt sql.NullInt64 `json:"finished_at"`
+	ID         string        `json:"id"`
 }
 
 func (q *Queries) UpdateMessage(ctx context.Context, arg UpdateMessageParams) error {
-	_, err := q.exec(ctx, q.updateMessageStmt, updateMessage, arg.Parts, arg.ID)
+	_, err := q.exec(ctx, q.updateMessageStmt, updateMessage, arg.Parts, arg.FinishedAt, arg.ID)
 	return err
 }
diff --git a/internal/db/migrations/000001_initial.up.sql b/internal/db/migrations/000001_initial.up.sql
index 2fbe5547e..03479449d 100644
--- a/internal/db/migrations/000001_initial.up.sql
+++ b/internal/db/migrations/000001_initial.up.sql
@@ -24,8 +24,10 @@ CREATE TABLE IF NOT EXISTS messages (
     session_id TEXT NOT NULL,
     role TEXT NOT NULL,
     parts TEXT NOT NULL default '[]',
+    model TEXT,
     created_at INTEGER NOT NULL,  -- Unix timestamp in milliseconds
     updated_at INTEGER NOT NULL,  -- Unix timestamp in milliseconds
+    finished_at INTEGER,  -- 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 1ad8607a9..2fad913be 100644
--- a/internal/db/models.go
+++ b/internal/db/models.go
@@ -9,12 +9,14 @@ import (
 )
 
 type Message struct {
-	ID        string `json:"id"`
-	SessionID string `json:"session_id"`
-	Role      string `json:"role"`
-	Parts     string `json:"parts"`
-	CreatedAt int64  `json:"created_at"`
-	UpdatedAt int64  `json:"updated_at"`
+	ID         string         `json:"id"`
+	SessionID  string         `json:"session_id"`
+	Role       string         `json:"role"`
+	Parts      string         `json:"parts"`
+	Model      sql.NullString `json:"model"`
+	CreatedAt  int64          `json:"created_at"`
+	UpdatedAt  int64          `json:"updated_at"`
+	FinishedAt sql.NullInt64  `json:"finished_at"`
 }
 
 type Session struct {
diff --git a/internal/db/sql/messages.sql b/internal/db/sql/messages.sql
index 64571158f..a59cebe7d 100644
--- a/internal/db/sql/messages.sql
+++ b/internal/db/sql/messages.sql
@@ -15,10 +15,11 @@ INSERT INTO messages (
     session_id,
     role,
     parts,
+    model,
     created_at,
     updated_at
 ) VALUES (
-    ?, ?, ?, ?, strftime('%s', 'now'), strftime('%s', 'now')
+    ?, ?, ?, ?, ?, strftime('%s', 'now'), strftime('%s', 'now')
 )
 RETURNING *;
 
@@ -26,9 +27,11 @@ RETURNING *;
 UPDATE messages
 SET
     parts = ?,
+    finished_at = ?,
     updated_at = strftime('%s', 'now')
 WHERE id = ?;
 
+
 -- name: DeleteMessage :exec
 DELETE FROM messages
 WHERE id = ?;