CERBERUS/ollama
1
0
Fork 0
mirror of https://github.com/ollama/ollama.git synced 2026-10-24 04:21:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

agent: persist chat images

Browse source
This commit is contained in:
ParthSareen 2026-06-11 20:29:48 -07:00
parent 3be6794929
commit a7c8484bea
2 changed files with 185 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

58
agent/chatstore/sqlite.go
View file

@ -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 {

136
agent/chatstore/sqlite_test.go
View file

@ -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, &notNull, &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 {