add initial message handling

2 files changed, 25 insertions, 3 deletions

diff --git a/internal/message/content.go b/internal/message/content.go
index 2604cd68a..cd263798b 100644
--- a/internal/message/content.go
+++ b/internal/message/content.go
@@ -2,6 +2,7 @@ package message
 
 import (
 	"encoding/base64"
+	"time"
 )
 
 type MessageRole string
@@ -64,6 +65,7 @@ type ToolCall struct {
 	Name     string `json:"name"`
 	Input    string `json:"input"`
 	Type     string `json:"type"`
+	Metadata any    `json:"metadata"`
 	Finished bool   `json:"finished"`
 }
 
@@ -80,6 +82,7 @@ func (ToolResult) isPart() {}
 
 type Finish struct {
 	Reason string `json:"reason"`
+	Time   int64  `json:"time"`
 }
 
 func (Finish) isPart() {}
@@ -161,6 +164,15 @@ func (m *Message) IsFinished() bool {
 	return false
 }
 
+func (m *Message) FinishPart() *Finish {
+	for _, part := range m.Parts {
+		if c, ok := part.(Finish); ok {
+			return &c
+		}
+	}
+	return nil
+}
+
 func (m *Message) FinishReason() string {
 	for _, part := range m.Parts {
 		if c, ok := part.(Finish); ok {
@@ -232,7 +244,7 @@ func (m *Message) SetToolResults(tr []ToolResult) {
 }
 
 func (m *Message) AddFinish(reason string) {
-	m.Parts = append(m.Parts, Finish{Reason: reason})
+	m.Parts = append(m.Parts, Finish{Reason: reason, Time: time.Now().Unix()})
 }
 
 func (m *Message) AddImageURL(url, detail string) {
diff --git a/internal/message/message.go b/internal/message/message.go
index 13cf54048..eeeb83ed2 100644
--- a/internal/message/message.go
+++ b/internal/message/message.go
@@ -2,17 +2,20 @@ package message
 
 import (
 	"context"
+	"database/sql"
 	"encoding/json"
 	"fmt"
 
 	"github.com/google/uuid"
 	"github.com/kujtimiihoxha/termai/internal/db"
+	"github.com/kujtimiihoxha/termai/internal/llm/models"
 	"github.com/kujtimiihoxha/termai/internal/pubsub"
 )
 
 type CreateMessageParams struct {
 	Role  MessageRole
 	Parts []ContentPart
+	Model models.ModelID
 }
 
 type Service interface {
@@ -68,6 +71,7 @@ func (s *service) Create(sessionID string, params CreateMessageParams) (Message,
 		SessionID: sessionID,
 		Role:      string(params.Role),
 		Parts:     string(partsJSON),
+		Model:     sql.NullString{String: string(params.Model), Valid: true},
 	})
 	if err != nil {
 		return Message{}, err
@@ -101,9 +105,15 @@ func (s *service) Update(message Message) error {
 	if err != nil {
 		return err
 	}
+	finishedAt := sql.NullInt64{}
+	if f := message.FinishPart(); f != nil {
+		finishedAt.Int64 = f.Time
+		finishedAt.Valid = true
+	}
 	err = s.q.UpdateMessage(s.ctx, db.UpdateMessageParams{
-		ID:    message.ID,
-		Parts: string(parts),
+		ID:         message.ID,
+		Parts:      string(parts),
+		FinishedAt: finishedAt,
 	})
 	if err != nil {
 		return err