diff --git a/agent/chatstore/sqlite.go b/agent/chatstore/sqlite.go index 4e555d193f..fbb6c05533 100644 --- a/agent/chatstore/sqlite.go +++ b/agent/chatstore/sqlite.go @@ -100,10 +100,9 @@ func (s *Store) init(ctx context.Context) error { role TEXT NOT NULL, content TEXT NOT NULL DEFAULT '', thinking TEXT NOT NULL DEFAULT '', + images TEXT NOT NULL DEFAULT '[]', stream BOOLEAN NOT NULL DEFAULT 0, model_name TEXT, - model_cloud BOOLEAN, - model_ollama_host BOOLEAN, created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, thinking_time_start TIMESTAMP, @@ -152,6 +151,9 @@ func (s *Store) init(ctx context.Context) error { if err := ensureColumn(ctx, s.db, "messages", "tool_call_id", "TEXT NOT NULL DEFAULT ''"); err != nil { return err } + if err := ensureColumn(ctx, s.db, "messages", "images", "TEXT NOT NULL DEFAULT '[]'"); err != nil { + return err + } if _, err := s.db.ExecContext(ctx, `CREATE INDEX IF NOT EXISTS idx_messages_chat_id_archived ON messages(chat_id, archived, id)`); err != nil { return fmt.Errorf("create archived messages index: %w", err) } @@ -236,11 +238,16 @@ func (s *Store) UpdateLastMessageWithModel(ctx context.Context, chatID string, m modelName = sql.NullString{String: model, Valid: true} } + imagesJSON, err := marshalMessageImages(msg.Images) + if err != nil { + return err + } + _, err = tx.ExecContext(ctx, ` UPDATE messages - SET role = ?, content = ?, thinking = ?, tool_name = ?, tool_call_id = ?, model_name = ?, updated_at = ? + SET role = ?, content = ?, thinking = ?, images = ?, tool_name = ?, tool_call_id = ?, model_name = ?, updated_at = ? WHERE id = ? - `, msg.Role, msg.Content, msg.Thinking, msg.ToolName, msg.ToolCallID, modelName, now, messageID) + `, msg.Role, msg.Content, msg.Thinking, imagesJSON, msg.ToolName, msg.ToolCallID, modelName, now, messageID) if err != nil { return fmt.Errorf("update last message: %w", err) } @@ -281,7 +288,7 @@ func (s *Store) Chat(ctx context.Context, id string) (*Chat, error) { } rows, err := s.db.QueryContext(ctx, ` - SELECT id, role, content, thinking, tool_name, tool_call_id FROM messages WHERE chat_id = ? AND archived = 0 ORDER BY id ASC + SELECT id, role, content, thinking, images, tool_name, tool_call_id FROM messages WHERE chat_id = ? AND archived = 0 ORDER BY id ASC `, id) if err != nil { return nil, err @@ -291,9 +298,15 @@ func (s *Store) Chat(ctx context.Context, id string) (*Chat, error) { for rows.Next() { var messageID int64 var msg api.Message - if err := rows.Scan(&messageID, &msg.Role, &msg.Content, &msg.Thinking, &msg.ToolName, &msg.ToolCallID); err != nil { + var imagesJSON string + if err := rows.Scan(&messageID, &msg.Role, &msg.Content, &msg.Thinking, &imagesJSON, &msg.ToolName, &msg.ToolCallID); err != nil { return nil, err } + images, err := unmarshalMessageImages(imagesJSON) + if err != nil { + return nil, err + } + msg.Images = images toolCalls, err := getToolCalls(ctx, s.db, messageID) if err != nil { return nil, err @@ -564,10 +577,14 @@ func insertMessage(ctx context.Context, tx *sql.Tx, chatID string, msg api.Messa if model != "" { modelName = sql.NullString{String: model, Valid: true} } + imagesJSON, err := marshalMessageImages(msg.Images) + if err != nil { + return 0, err + } result, err := tx.ExecContext(ctx, ` - INSERT INTO messages (chat_id, role, content, thinking, stream, tool_name, tool_call_id, model_name, created_at, updated_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - `, chatID, msg.Role, msg.Content, msg.Thinking, false, msg.ToolName, msg.ToolCallID, modelName, now, now) + INSERT INTO messages (chat_id, role, content, thinking, images, stream, tool_name, tool_call_id, model_name, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + `, chatID, msg.Role, msg.Content, msg.Thinking, imagesJSON, false, msg.ToolName, msg.ToolCallID, modelName, now, now) if err != nil { return 0, fmt.Errorf("insert message: %w", err) } @@ -578,6 +595,29 @@ func insertMessage(ctx context.Context, tx *sql.Tx, chatID string, msg api.Messa return id, nil } +func marshalMessageImages(images []api.ImageData) (string, error) { + if len(images) == 0 { + return "[]", nil + } + data, err := json.Marshal(images) + if err != nil { + return "", fmt.Errorf("marshal message images: %w", err) + } + return string(data), nil +} + +func unmarshalMessageImages(value string) ([]api.ImageData, error) { + value = strings.TrimSpace(value) + if value == "" || value == "null" { + return nil, nil + } + var images []api.ImageData + if err := json.Unmarshal([]byte(value), &images); err != nil { + return nil, fmt.Errorf("unmarshal message images: %w", err) + } + return images, nil +} + func insertToolCall(ctx context.Context, tx *sql.Tx, messageID int64, call api.ToolCall) error { args, err := json.Marshal(call.Function.Arguments) if err != nil { diff --git a/agent/chatstore/sqlite_test.go b/agent/chatstore/sqlite_test.go index cb52beaa8c..170c1f4eff 100644 --- a/agent/chatstore/sqlite_test.go +++ b/agent/chatstore/sqlite_test.go @@ -1,6 +1,7 @@ package chatstore import ( + "bytes" "context" "database/sql" "encoding/json" @@ -96,6 +97,141 @@ func TestStoreRestoresToolMessageMetadata(t *testing.T) { } } +func TestStoreRoundTripsMessageImages(t *testing.T) { + store, err := New(filepath.Join(t.TempDir(), "db.sqlite")) + if err != nil { + t.Fatal(err) + } + defer store.Close() + + ctx := context.Background() + first := api.ImageData([]byte("first image")) + second := api.ImageData([]byte{0, 1, 2, 3}) + if err := store.AppendMessage(ctx, "chat-1", api.Message{ + Role: "user", + Content: "look", + Images: []api.ImageData{first, second}, + }); err != nil { + t.Fatal(err) + } + + chat, err := store.Chat(ctx, "chat-1") + if err != nil { + t.Fatal(err) + } + if len(chat.Messages) != 1 || len(chat.Messages[0].Images) != 2 { + t.Fatalf("messages = %#v", chat.Messages) + } + if !bytes.Equal(chat.Messages[0].Images[0], first) || !bytes.Equal(chat.Messages[0].Images[1], second) { + t.Fatalf("images = %#v, want %#v", chat.Messages[0].Images, []api.ImageData{first, second}) + } + + updated := api.ImageData([]byte("updated image")) + if err := store.UpdateLastMessage(ctx, "chat-1", api.Message{ + Role: "user", + Content: "updated", + Images: []api.ImageData{updated}, + }); err != nil { + t.Fatal(err) + } + chat, err = store.Chat(ctx, "chat-1") + if err != nil { + t.Fatal(err) + } + if len(chat.Messages) != 1 || len(chat.Messages[0].Images) != 1 || !bytes.Equal(chat.Messages[0].Images[0], updated) { + t.Fatalf("updated images = %#v", chat.Messages) + } +} + +func TestStoreMigratesOldMessagesSchemaForImages(t *testing.T) { + path := filepath.Join(t.TempDir(), "db.sqlite") + db, err := sql.Open("sqlite3", path) + if err != nil { + t.Fatal(err) + } + if _, err := db.Exec(` + CREATE TABLE chats ( + id TEXT PRIMARY KEY, + title TEXT NOT NULL DEFAULT '', + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + browser_state TEXT + ); + + CREATE TABLE messages ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + chat_id TEXT NOT NULL, + role TEXT NOT NULL, + content TEXT NOT NULL DEFAULT '', + thinking TEXT NOT NULL DEFAULT '', + stream BOOLEAN NOT NULL DEFAULT 0, + model_name TEXT, + model_cloud BOOLEAN, + model_ollama_host BOOLEAN, + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + thinking_time_start TIMESTAMP, + thinking_time_end TIMESTAMP, + tool_result TEXT, + tool_name TEXT NOT NULL DEFAULT '', + tool_call_id TEXT NOT NULL DEFAULT '', + archived BOOLEAN NOT NULL DEFAULT 0, + FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE + ); + `); err != nil { + db.Close() + t.Fatal(err) + } + if err := db.Close(); err != nil { + t.Fatal(err) + } + + store, err := New(path) + if err != nil { + t.Fatal(err) + } + defer store.Close() + + var found bool + rows, err := store.db.Query(`PRAGMA table_info(messages)`) + if err != nil { + t.Fatal(err) + } + defer rows.Close() + for rows.Next() { + var cid int + var name string + var typ string + var notNull int + var defaultValue any + var pk int + if err := rows.Scan(&cid, &name, &typ, ¬Null, &defaultValue, &pk); err != nil { + t.Fatal(err) + } + if name == "images" { + found = true + } + } + if err := rows.Err(); err != nil { + t.Fatal(err) + } + if !found { + t.Fatal("images column was not added") + } + + ctx := context.Background() + image := api.ImageData([]byte("after migration")) + if err := store.AppendMessage(ctx, "chat-1", api.Message{Role: "user", Content: "image", Images: []api.ImageData{image}}); err != nil { + t.Fatal(err) + } + chat, err := store.Chat(ctx, "chat-1") + if err != nil { + t.Fatal(err) + } + if len(chat.Messages) != 1 || len(chat.Messages[0].Images) != 1 || !bytes.Equal(chat.Messages[0].Images[0], image) { + t.Fatalf("migrated images = %#v", chat.Messages) + } +} + func TestStoreLatestChatForModel(t *testing.T) { store, err := New(filepath.Join(t.TempDir(), "db.sqlite")) if err != nil {