Skip to content

GitLab

Unverified Commit d3b4b997 authored by Daniel Hiltgen's avatar Daniel Hiltgen Committed by GitHub
Browse files

app: add code for macOS and Windows apps under 'app' (#12933) 



* app: add code for macOS and Windows apps under 'app'

* app: add readme

* app: windows and linux only for now

* ci: fix ui CI validation

---------
Co-authored-by: default avatarjmorganca <jmorganca@gmail.com>
parent a4770107
Hide whitespace changes
Inline Side-by-side
Showing with 4652 additions and 99 deletions
app/lifecycle/server_windows.go → app/server/server_windows.go
View file @ d3b4b997
package lifecycle
package server
import (
"context"
"fmt"
"log/slog"
"os"
"os/exec"
"path/filepath"
"strconv"
"strings"
"syscall"
"golang.org/x/sys/windows"
)
func getCmd(ctx context.Context, exePath string) *exec.Cmd {
cmd := exec.CommandContext(ctx, exePath, "serve")
var (
pidFile = filepath.Join(os.Getenv("LOCALAPPDATA"), "Ollama", "ollama.pid")
serverLogPath = filepath.Join(os.Getenv("LOCALAPPDATA"), "Ollama", "server.log")
)
func commandContext(ctx context.Context, name string, arg ...string) *exec.Cmd {
cmd := exec.CommandContext(ctx, name, arg...)
cmd.SysProcAttr = &syscall.SysProcAttr{
HideWindow: true,
CreationFlags: windows.CREATE_NEW_PROCESS_GROUP,
......@@ -19,15 +29,14 @@ func getCmd(ctx context.Context, exePath string) *exec.Cmd {
return cmd
}
func terminate(cmd *exec.Cmd) error {
func terminate(proc *os.Process) error {
dll, err := windows.LoadDLL("kernel32.dll")
if err != nil {
return err
}
//nolint:errcheck
defer dll.Release()
pid := cmd.Process.Pid
pid := proc.Pid
f, err := dll.FindProc("AttachConsole")
if err != nil {
......@@ -69,12 +78,14 @@ func terminate(cmd *exec.Cmd) error {
const STILL_ACTIVE = 259
func isProcessExited(pid int) (bool, error) {
func terminated(pid int) (bool, error) {
hProcess, err := windows.OpenProcess(windows.PROCESS_QUERY_INFORMATION, false, uint32(pid))
if err != nil {
if errno, ok := err.(windows.Errno); ok && errno == windows.ERROR_INVALID_PARAMETER {
return true, nil
}
return false, fmt.Errorf("failed to open process: %v", err)
}
//nolint:errcheck
defer windows.CloseHandle(hProcess)
var exitCode uint32
......@@ -89,3 +100,50 @@ func isProcessExited(pid int) (bool, error) {
return true, nil
}
// reapServers kills all ollama processes except our own
func reapServers() error {
// Get current process ID to avoid killing ourselves
currentPID := os.Getpid()
// Use wmic to find ollama processes
cmd := exec.Command("wmic", "process", "where", "name='ollama.exe'", "get", "ProcessId")
cmd.SysProcAttr = &syscall.SysProcAttr{HideWindow: true}
output, err := cmd.Output()
if err != nil {
// No ollama processes found
slog.Debug("no ollama processes found")
return nil //nolint:nilerr
}
lines := strings.Split(string(output), "\n")
var pids []string
for _, line := range lines {
line = strings.TrimSpace(line)
if line == "" || line == "ProcessId" {
continue
}
if _, err := strconv.Atoi(line); err == nil {
pids = append(pids, line)
}
}
for _, pidStr := range pids {
pid, err := strconv.Atoi(pidStr)
if err != nil {
continue
}
if pid == currentPID {
continue
}
cmd := exec.Command("taskkill", "/F", "/PID", pidStr)
if err := cmd.Run(); err != nil {
slog.Warn("failed to kill ollama process", "pid", pid, "err", err)
}
}
return nil
}
app/store/database.go 0 → 100644
View file @ d3b4b997
//go:build windows || darwin
package store
import (
"database/sql"
"encoding/json"
"fmt"
"strings"
"time"
sqlite3 "github.com/mattn/go-sqlite3"
)
// currentSchemaVersion defines the current database schema version.
// Increment this when making schema changes that require migrations.
const currentSchemaVersion = 12
// database wraps the SQLite connection.
// SQLite handles its own locking for concurrent access:
// - Multiple readers can access the database simultaneously
// - Writers are serialized (only one writer at a time)
// - WAL mode allows readers to not block writers
// This means we don't need application-level locks for database operations.
type database struct {
conn *sql.DB
}
func newDatabase(dbPath string) (*database, error) {
// Open database connection
conn, err := sql.Open("sqlite3", dbPath+"?_foreign_keys=on&_journal_mode=WAL&_busy_timeout=5000&_txlock=immediate")
if err != nil {
return nil, fmt.Errorf("open database: %w", err)
}
// Test the connection
if err := conn.Ping(); err != nil {
conn.Close()
return nil, fmt.Errorf("ping database: %w", err)
}
db := &database{conn: conn}
// Initialize schema
if err := db.init(); err != nil {
conn.Close()
return nil, fmt.Errorf("initialize database: %w", err)
}
return db, nil
}
func (db *database) Close() error {
_, _ = db.conn.Exec("PRAGMA wal_checkpoint(TRUNCATE);")
return db.conn.Close()
}
func (db *database) init() error {
if _, err := db.conn.Exec("PRAGMA foreign_keys = ON"); err != nil {
return fmt.Errorf("enable foreign keys: %w", err)
}
schema := fmt.Sprintf(`
CREATE TABLE IF NOT EXISTS settings (
id INTEGER PRIMARY KEY CHECK (id = 1),
device_id TEXT NOT NULL DEFAULT '',
has_completed_first_run BOOLEAN NOT NULL DEFAULT 0,
expose BOOLEAN NOT NULL DEFAULT 0,
survey BOOLEAN NOT NULL DEFAULT TRUE,
browser BOOLEAN NOT NULL DEFAULT 0,
models TEXT NOT NULL DEFAULT '',
agent BOOLEAN NOT NULL DEFAULT 0,
tools BOOLEAN NOT NULL DEFAULT 0,
working_dir TEXT NOT NULL DEFAULT '',
context_length INTEGER NOT NULL DEFAULT 4096,
window_width INTEGER NOT NULL DEFAULT 0,
window_height INTEGER NOT NULL DEFAULT 0,
config_migrated BOOLEAN NOT NULL DEFAULT 0,
airplane_mode BOOLEAN NOT NULL DEFAULT 0,
turbo_enabled BOOLEAN NOT NULL DEFAULT 0,
websearch_enabled BOOLEAN NOT NULL DEFAULT 0,
selected_model TEXT NOT NULL DEFAULT '',
sidebar_open BOOLEAN NOT NULL DEFAULT 0,
think_enabled BOOLEAN NOT NULL DEFAULT 0,
think_level TEXT NOT NULL DEFAULT '',
remote TEXT NOT NULL DEFAULT '', -- deprecated
schema_version INTEGER NOT NULL DEFAULT %d
);
-- Insert default settings row if it doesn't exist
INSERT OR IGNORE INTO settings (id) VALUES (1);
CREATE TABLE IF NOT EXISTS chats (
id TEXT PRIMARY KEY,
title TEXT NOT NULL DEFAULT '',
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
browser_state TEXT
);
CREATE TABLE IF NOT EXISTS 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, -- deprecated
model_ollama_host BOOLEAN, -- deprecated
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,
FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS idx_messages_chat_id ON messages(chat_id);
CREATE TABLE IF NOT EXISTS tool_calls (
id INTEGER PRIMARY KEY AUTOINCREMENT,
message_id INTEGER NOT NULL,
type TEXT NOT NULL,
function_name TEXT NOT NULL,
function_arguments TEXT NOT NULL,
function_result TEXT,
FOREIGN KEY (message_id) REFERENCES messages(id) ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS idx_tool_calls_message_id ON tool_calls(message_id);
CREATE TABLE IF NOT EXISTS attachments (
id INTEGER PRIMARY KEY AUTOINCREMENT,
message_id INTEGER NOT NULL,
filename TEXT NOT NULL,
data BLOB NOT NULL,
FOREIGN KEY (message_id) REFERENCES messages(id) ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS idx_attachments_message_id ON attachments(message_id);
CREATE TABLE IF NOT EXISTS users (
name TEXT NOT NULL DEFAULT '',
email TEXT NOT NULL DEFAULT '',
plan TEXT NOT NULL DEFAULT '',
cached_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
);
`, currentSchemaVersion)
_, err := db.conn.Exec(schema)
if err != nil {
return err
}
// Check and upgrade schema version if needed
if err := db.migrate(); err != nil {
return fmt.Errorf("migrate schema: %w", err)
}
// Clean up orphaned records created before foreign key constraints were properly enforced
// TODO: Can eventually be removed - cleans up data from foreign key bug (ollama/ollama#11785, ollama/app#476)
if err := db.cleanupOrphanedData(); err != nil {
return fmt.Errorf("cleanup orphaned data: %w", err)
}
return nil
}
// migrate handles database schema migrations
func (db *database) migrate() error {
// Get current schema version
version, err := db.getSchemaVersion()
if err != nil {
return fmt.Errorf("get schema version after migration attempt: %w", err)
}
// Run migrations for each version
for version < currentSchemaVersion {
switch version {
case 1:
// Migrate from version 1 to 2: add context_length column
if err := db.migrateV1ToV2(); err != nil {
return fmt.Errorf("migrate v1 to v2: %w", err)
}
version = 2
case 2:
// Migrate from version 2 to 3: create attachments table
if err := db.migrateV2ToV3(); err != nil {
return fmt.Errorf("migrate v2 to v3: %w", err)
}
version = 3
case 3:
// Migrate from version 3 to 4: add tool_result column to messages table
if err := db.migrateV3ToV4(); err != nil {
return fmt.Errorf("migrate v3 to v4: %w", err)
}
version = 4
case 4:
// add airplane_mode column to settings table
if err := db.migrateV4ToV5(); err != nil {
return fmt.Errorf("migrate v4 to v5: %w", err)
}
version = 5
case 5:
// add turbo_enabled column to settings table
if err := db.migrateV5ToV6(); err != nil {
return fmt.Errorf("migrate v5 to v6: %w", err)
}
version = 6
case 6:
// add missing index for attachments table
if err := db.migrateV6ToV7(); err != nil {
return fmt.Errorf("migrate v6 to v7: %w", err)
}
version = 7
case 7:
// add think_enabled and think_level columns to settings table
if err := db.migrateV7ToV8(); err != nil {
return fmt.Errorf("migrate v7 to v8: %w", err)
}
version = 8
case 8:
// add browser_state column to chats table
if err := db.migrateV8ToV9(); err != nil {
return fmt.Errorf("migrate v8 to v9: %w", err)
}
version = 9
case 9:
// add cached user table
if err := db.migrateV9ToV10(); err != nil {
return fmt.Errorf("migrate v9 to v10: %w", err)
}
version = 10
case 10:
// remove remote column from settings table
if err := db.migrateV10ToV11(); err != nil {
return fmt.Errorf("migrate v10 to v11: %w", err)
}
version = 11
case 11:
// bring back remote column for backwards compatibility (deprecated)
if err := db.migrateV11ToV12(); err != nil {
return fmt.Errorf("migrate v11 to v12: %w", err)
}
version = 12
default:
// If we have a version we don't recognize, just set it to current
// This might happen during development
version = currentSchemaVersion
}
}
return nil
}
// migrateV1ToV2 adds the context_length column to the settings table
func (db *database) migrateV1ToV2() error {
_, err := db.conn.Exec(`ALTER TABLE settings ADD COLUMN context_length INTEGER NOT NULL DEFAULT 4096;`)
if err != nil && !duplicateColumnError(err) {
return fmt.Errorf("add context_length column: %w", err)
}
_, err = db.conn.Exec(`ALTER TABLE settings ADD COLUMN survey BOOLEAN NOT NULL DEFAULT TRUE;`)
if err != nil && !duplicateColumnError(err) {
return fmt.Errorf("add survey column: %w", err)
}
_, err = db.conn.Exec(`UPDATE settings SET schema_version = 2;`)
if err != nil {
return fmt.Errorf("update schema version: %w", err)
}
return nil
}
// migrateV2ToV3 creates the attachments table
func (db *database) migrateV2ToV3() error {
_, err := db.conn.Exec(`
CREATE TABLE IF NOT EXISTS attachments (
id INTEGER PRIMARY KEY AUTOINCREMENT,
message_id INTEGER NOT NULL,
filename TEXT NOT NULL,
data BLOB NOT NULL,
FOREIGN KEY (message_id) REFERENCES messages(id) ON DELETE CASCADE
)
`)
if err != nil {
return fmt.Errorf("create attachments table: %w", err)
}
_, err = db.conn.Exec(`UPDATE settings SET schema_version = 3`)
if err != nil {
return fmt.Errorf("update schema version: %w", err)
}
return nil
}
func (db *database) migrateV3ToV4() error {
_, err := db.conn.Exec(`ALTER TABLE messages ADD COLUMN tool_result TEXT;`)
if err != nil && !duplicateColumnError(err) {
return fmt.Errorf("add tool_result column: %w", err)
}
_, err = db.conn.Exec(`UPDATE settings SET schema_version = 4;`)
if err != nil {
return fmt.Errorf("update schema version: %w", err)
}
return nil
}
// migrateV4ToV5 adds the airplane_mode column to the settings table
func (db *database) migrateV4ToV5() error {
_, err := db.conn.Exec(`ALTER TABLE settings ADD COLUMN airplane_mode BOOLEAN NOT NULL DEFAULT 0;`)
if err != nil && !duplicateColumnError(err) {
return fmt.Errorf("add airplane_mode column: %w", err)
}
_, err = db.conn.Exec(`UPDATE settings SET schema_version = 5;`)
if err != nil {
return fmt.Errorf("update schema version: %w", err)
}
return nil
}
// migrateV5ToV6 adds the turbo_enabled, websearch_enabled, selected_model, sidebar_open columns to the settings table
func (db *database) migrateV5ToV6() error {
_, err := db.conn.Exec(`ALTER TABLE settings ADD COLUMN turbo_enabled BOOLEAN NOT NULL DEFAULT 0;`)
if err != nil && !duplicateColumnError(err) {
return fmt.Errorf("add turbo_enabled column: %w", err)
}
_, err = db.conn.Exec(`ALTER TABLE settings ADD COLUMN websearch_enabled BOOLEAN NOT NULL DEFAULT 0;`)
if err != nil && !duplicateColumnError(err) {
return fmt.Errorf("add websearch_enabled column: %w", err)
}
_, err = db.conn.Exec(`ALTER TABLE settings ADD COLUMN selected_model TEXT NOT NULL DEFAULT '';`)
if err != nil && !duplicateColumnError(err) {
return fmt.Errorf("add selected_model column: %w", err)
}
_, err = db.conn.Exec(`ALTER TABLE settings ADD COLUMN sidebar_open BOOLEAN NOT NULL DEFAULT 0;`)
if err != nil && !duplicateColumnError(err) {
return fmt.Errorf("add sidebar_open column: %w", err)
}
_, err = db.conn.Exec(`UPDATE settings SET schema_version = 6;`)
if err != nil {
return fmt.Errorf("update schema version: %w", err)
}
return nil
}
// migrateV6ToV7 adds the missing index for the attachments table
func (db *database) migrateV6ToV7() error {
_, err := db.conn.Exec(`CREATE INDEX IF NOT EXISTS idx_attachments_message_id ON attachments(message_id);`)
if err != nil {
return fmt.Errorf("create attachments index: %w", err)
}
_, err = db.conn.Exec(`UPDATE settings SET schema_version = 7;`)
if err != nil {
return fmt.Errorf("update schema version: %w", err)
}
return nil
}
// migrateV7ToV8 adds the think_enabled and think_level columns to the settings table
func (db *database) migrateV7ToV8() error {
_, err := db.conn.Exec(`ALTER TABLE settings ADD COLUMN think_enabled BOOLEAN NOT NULL DEFAULT 0;`)
if err != nil && !duplicateColumnError(err) {
return fmt.Errorf("add think_enabled column: %w", err)
}
_, err = db.conn.Exec(`ALTER TABLE settings ADD COLUMN think_level TEXT NOT NULL DEFAULT '';`)
if err != nil && !duplicateColumnError(err) {
return fmt.Errorf("add think_level column: %w", err)
}
_, err = db.conn.Exec(`UPDATE settings SET schema_version = 8;`)
if err != nil {
return fmt.Errorf("update schema version: %w", err)
}
return nil
}
// migrateV8ToV9 adds browser_state to chats and bumps schema
func (db *database) migrateV8ToV9() error {
_, err := db.conn.Exec(`
ALTER TABLE chats ADD COLUMN browser_state TEXT;
UPDATE settings SET schema_version = 9;
`)
if err != nil && !duplicateColumnError(err) {
return fmt.Errorf("add browser_state column: %w", err)
}
return nil
}
// migrateV9ToV10 adds users table
func (db *database) migrateV9ToV10() error {
_, err := db.conn.Exec(`
CREATE TABLE IF NOT EXISTS users (
name TEXT NOT NULL DEFAULT '',
email TEXT NOT NULL DEFAULT '',
plan TEXT NOT NULL DEFAULT '',
cached_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
);
UPDATE settings SET schema_version = 10;
`)
if err != nil {
return fmt.Errorf("create users table: %w", err)
}
return nil
}
// migrateV10ToV11 removes the remote column from the settings table
func (db *database) migrateV10ToV11() error {
_, err := db.conn.Exec(`ALTER TABLE settings DROP COLUMN remote`)
if err != nil && !columnNotExists(err) {
return fmt.Errorf("drop remote column: %w", err)
}
_, err = db.conn.Exec(`UPDATE settings SET schema_version = 11`)
if err != nil {
return fmt.Errorf("update schema version: %w", err)
}
return nil
}
// migrateV11ToV12 brings back the remote column for backwards compatibility (deprecated)
func (db *database) migrateV11ToV12() error {
_, err := db.conn.Exec(`ALTER TABLE settings ADD COLUMN remote TEXT NOT NULL DEFAULT ''`)
if err != nil && !duplicateColumnError(err) {
return fmt.Errorf("add remote column: %w", err)
}
_, err = db.conn.Exec(`UPDATE settings SET schema_version = 12`)
if err != nil {
return fmt.Errorf("update schema version: %w", err)
}
return nil
}
// cleanupOrphanedData removes orphaned records that may exist due to the foreign key bug
func (db *database) cleanupOrphanedData() error {
_, err := db.conn.Exec(`
DELETE FROM tool_calls
WHERE message_id NOT IN (SELECT id FROM messages)
`)
if err != nil {
return fmt.Errorf("cleanup orphaned tool_calls: %w", err)
}
_, err = db.conn.Exec(`
DELETE FROM attachments
WHERE message_id NOT IN (SELECT id FROM messages)
`)
if err != nil {
return fmt.Errorf("cleanup orphaned attachments: %w", err)
}
_, err = db.conn.Exec(`
DELETE FROM messages
WHERE chat_id NOT IN (SELECT id FROM chats)
`)
if err != nil {
return fmt.Errorf("cleanup orphaned messages: %w", err)
}
return nil
}
func duplicateColumnError(err error) bool {
if sqlite3Err, ok := err.(sqlite3.Error); ok {
return sqlite3Err.Code == sqlite3.ErrError &&
strings.Contains(sqlite3Err.Error(), "duplicate column name")
}
return false
}
func columnNotExists(err error) bool {
if sqlite3Err, ok := err.(sqlite3.Error); ok {
return sqlite3Err.Code == sqlite3.ErrError &&
strings.Contains(sqlite3Err.Error(), "no such column")
}
return false
}
func (db *database) getAllChats() ([]Chat, error) {
// Query chats with their first user message and latest update time
query := `
SELECT
c.id,
c.title,
c.created_at,
COALESCE(first_msg.content, '') as first_user_content,
COALESCE(datetime(MAX(m.updated_at)), datetime(c.created_at)) as last_updated
FROM chats c
LEFT JOIN (
SELECT chat_id, content, MIN(id) as min_id
FROM messages
WHERE role = 'user'
GROUP BY chat_id
) first_msg ON c.id = first_msg.chat_id
LEFT JOIN messages m ON c.id = m.chat_id
GROUP BY c.id, c.title, c.created_at, first_msg.content
ORDER BY last_updated DESC
`
rows, err := db.conn.Query(query)
if err != nil {
return nil, fmt.Errorf("query chats: %w", err)
}
defer rows.Close()
var chats []Chat
for rows.Next() {
var chat Chat
var createdAt time.Time
var firstUserContent string
var lastUpdatedStr string
err := rows.Scan(
&chat.ID,
&chat.Title,
&createdAt,
&firstUserContent,
&lastUpdatedStr,
)
// Parse the last updated time
lastUpdated, _ := time.Parse("2006-01-02 15:04:05", lastUpdatedStr)
if err != nil {
return nil, fmt.Errorf("scan chat: %w", err)
}
chat.CreatedAt = createdAt
// Add a dummy first user message for the UI to display
// This is just for the excerpt, full messages are loaded when needed
chat.Messages = []Message{}
if firstUserContent != "" {
chat.Messages = append(chat.Messages, Message{
Role: "user",
Content: firstUserContent,
UpdatedAt: lastUpdated,
})
}
chats = append(chats, chat)
}
if err := rows.Err(); err != nil {
return nil, fmt.Errorf("iterate chats: %w", err)
}
return chats, nil
}
func (db *database) getChatWithOptions(id string, loadAttachmentData bool) (*Chat, error) {
query := `
SELECT id, title, created_at, browser_state
FROM chats
WHERE id = ?
`
var chat Chat
var createdAt time.Time
var browserState sql.NullString
err := db.conn.QueryRow(query, id).Scan(
&chat.ID,
&chat.Title,
&createdAt,
&browserState,
)
if err != nil {
if err == sql.ErrNoRows {
return nil, fmt.Errorf("chat not found")
}
return nil, fmt.Errorf("query chat: %w", err)
}
chat.CreatedAt = createdAt
if browserState.Valid && browserState.String != "" {
var raw json.RawMessage
if err := json.Unmarshal([]byte(browserState.String), &raw); err == nil {
chat.BrowserState = raw
}
}
messages, err := db.getMessages(id, loadAttachmentData)
if err != nil {
return nil, fmt.Errorf("get messages: %w", err)
}
chat.Messages = messages
return &chat, nil
}
func (db *database) saveChat(chat Chat) error {
tx, err := db.conn.Begin()
if err != nil {
return fmt.Errorf("begin transaction: %w", err)
}
defer tx.Rollback()
// Use COALESCE for browser_state to avoid wiping an existing
// chat-level browser_state when saving a chat that doesn't include a new state payload.
// Many code paths call SetChat to update metadata/messages only; without COALESCE the
// UPSERT would overwrite browser_state with NULL, breaking revisit rendering that relies
// on the last persisted full tool state.
query := `
INSERT INTO chats (id, title, created_at, browser_state)
VALUES (?, ?, ?, ?)
ON CONFLICT(id) DO UPDATE SET
title = excluded.title,
browser_state = COALESCE(excluded.browser_state, chats.browser_state)
`
var browserState sql.NullString
if chat.BrowserState != nil {
browserState = sql.NullString{String: string(chat.BrowserState), Valid: true}
}
_, err = tx.Exec(query,
chat.ID,
chat.Title,
chat.CreatedAt,
browserState,
)
if err != nil {
return fmt.Errorf("save chat: %w", err)
}
// Delete existing messages (we'll re-insert all)
_, err = tx.Exec("DELETE FROM messages WHERE chat_id = ?", chat.ID)
if err != nil {
return fmt.Errorf("delete messages: %w", err)
}
// Insert messages
for _, msg := range chat.Messages {
messageID, err := db.insertMessage(tx, chat.ID, msg)
if err != nil {
return fmt.Errorf("insert message: %w", err)
}
// Insert tool calls if any
for _, toolCall := range msg.ToolCalls {
err := db.insertToolCall(tx, messageID, toolCall)
if err != nil {
return fmt.Errorf("insert tool call: %w", err)
}
}
}
return tx.Commit()
}
// updateChatBrowserState updates only the browser_state for a chat
func (db *database) updateChatBrowserState(chatID string, state json.RawMessage) error {
_, err := db.conn.Exec(`UPDATE chats SET browser_state = ? WHERE id = ?`, string(state), chatID)
if err != nil {
return fmt.Errorf("update chat browser state: %w", err)
}
return nil
}
func (db *database) deleteChat(id string) error {
_, err := db.conn.Exec("DELETE FROM chats WHERE id = ?", id)
if err != nil {
return fmt.Errorf("delete chat: %w", err)
}
_, _ = db.conn.Exec("PRAGMA wal_checkpoint(TRUNCATE);")
return nil
}
func (db *database) updateLastMessage(chatID string, msg Message) error {
tx, err := db.conn.Begin()
if err != nil {
return fmt.Errorf("begin transaction: %w", err)
}
defer tx.Rollback()
// Get the ID of the last message
var messageID int64
err = tx.QueryRow(`
SELECT MAX(id) FROM messages WHERE chat_id = ?
`, chatID).Scan(&messageID)
if err != nil {
return fmt.Errorf("get last message id: %w", err)
}
query := `
UPDATE messages
SET content = ?, thinking = ?, model_name = ?, updated_at = ?, thinking_time_start = ?, thinking_time_end = ?, tool_result = ?
WHERE id = ?
`
var thinkingTimeStart, thinkingTimeEnd sql.NullTime
if msg.ThinkingTimeStart != nil {
thinkingTimeStart = sql.NullTime{Time: *msg.ThinkingTimeStart, Valid: true}
}
if msg.ThinkingTimeEnd != nil {
thinkingTimeEnd = sql.NullTime{Time: *msg.ThinkingTimeEnd, Valid: true}
}
var modelName sql.NullString
if msg.Model != "" {
modelName = sql.NullString{String: msg.Model, Valid: true}
}
var toolResultJSON sql.NullString
if msg.ToolResult != nil {
resultBytes, err := json.Marshal(msg.ToolResult)
if err != nil {
return fmt.Errorf("marshal tool result: %w", err)
}
toolResultJSON = sql.NullString{String: string(resultBytes), Valid: true}
}
result, err := tx.Exec(query,
msg.Content,
msg.Thinking,
modelName,
msg.UpdatedAt,
thinkingTimeStart,
thinkingTimeEnd,
toolResultJSON,
messageID,
)
if err != nil {
return fmt.Errorf("update last message: %w", err)
}
rowsAffected, err := result.RowsAffected()
if err != nil {
return fmt.Errorf("get rows affected: %w", err)
}
if rowsAffected == 0 {
return fmt.Errorf("no message found to update")
}
_, err = tx.Exec("DELETE FROM attachments WHERE message_id = ?", messageID)
if err != nil {
return fmt.Errorf("delete existing attachments: %w", err)
}
for _, att := range msg.Attachments {
err := db.insertAttachment(tx, messageID, att)
if err != nil {
return fmt.Errorf("insert attachment: %w", err)
}
}
_, err = tx.Exec("DELETE FROM tool_calls WHERE message_id = ?", messageID)
if err != nil {
return fmt.Errorf("delete existing tool calls: %w", err)
}
for _, toolCall := range msg.ToolCalls {
err := db.insertToolCall(tx, messageID, toolCall)
if err != nil {
return fmt.Errorf("insert tool call: %w", err)
}
}
return tx.Commit()
}
func (db *database) appendMessage(chatID string, msg Message) error {
tx, err := db.conn.Begin()
if err != nil {
return fmt.Errorf("begin transaction: %w", err)
}
defer tx.Rollback()
messageID, err := db.insertMessage(tx, chatID, msg)
if err != nil {
return fmt.Errorf("insert message: %w", err)
}
// Insert tool calls if any
for _, toolCall := range msg.ToolCalls {
err := db.insertToolCall(tx, messageID, toolCall)
if err != nil {
return fmt.Errorf("insert tool call: %w", err)
}
}
return tx.Commit()
}
func (db *database) getMessages(chatID string, loadAttachmentData bool) ([]Message, error) {
query := `
SELECT id, role, content, thinking, stream, model_name, created_at, updated_at, thinking_time_start, thinking_time_end, tool_result
FROM messages
WHERE chat_id = ?
ORDER BY id ASC
`
rows, err := db.conn.Query(query, chatID)
if err != nil {
return nil, fmt.Errorf("query messages: %w", err)
}
defer rows.Close()
var messages []Message
for rows.Next() {
var msg Message
var messageID int64
var thinkingTimeStart, thinkingTimeEnd sql.NullTime
var modelName sql.NullString
var toolResult sql.NullString
err := rows.Scan(
&messageID,
&msg.Role,
&msg.Content,
&msg.Thinking,
&msg.Stream,
&modelName,
&msg.CreatedAt,
&msg.UpdatedAt,
&thinkingTimeStart,
&thinkingTimeEnd,
&toolResult,
)
if err != nil {
return nil, fmt.Errorf("scan message: %w", err)
}
attachments, err := db.getAttachments(messageID, loadAttachmentData)
if err != nil {
return nil, fmt.Errorf("get attachments: %w", err)
}
msg.Attachments = attachments
if thinkingTimeStart.Valid {
msg.ThinkingTimeStart = &thinkingTimeStart.Time
}
if thinkingTimeEnd.Valid {
msg.ThinkingTimeEnd = &thinkingTimeEnd.Time
}
// Parse tool result from JSON if present
if toolResult.Valid && toolResult.String != "" {
var result json.RawMessage
if err := json.Unmarshal([]byte(toolResult.String), &result); err == nil {
msg.ToolResult = &result
}
}
// Set model if present
if modelName.Valid && modelName.String != "" {
msg.Model = modelName.String
}
// Get tool calls for this message
toolCalls, err := db.getToolCalls(messageID)
if err != nil {
return nil, fmt.Errorf("get tool calls: %w", err)
}
msg.ToolCalls = toolCalls
messages = append(messages, msg)
}
if err := rows.Err(); err != nil {
return nil, fmt.Errorf("iterate messages: %w", err)
}
return messages, nil
}
func (db *database) insertMessage(tx *sql.Tx, chatID string, msg Message) (int64, error) {
query := `
INSERT INTO messages (chat_id, role, content, thinking, stream, model_name, created_at, updated_at, thinking_time_start, thinking_time_end, tool_result)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
`
var thinkingTimeStart, thinkingTimeEnd sql.NullTime
if msg.ThinkingTimeStart != nil {
thinkingTimeStart = sql.NullTime{Time: *msg.ThinkingTimeStart, Valid: true}
}
if msg.ThinkingTimeEnd != nil {
thinkingTimeEnd = sql.NullTime{Time: *msg.ThinkingTimeEnd, Valid: true}
}
var modelName sql.NullString
if msg.Model != "" {
modelName = sql.NullString{String: msg.Model, Valid: true}
}
var toolResultJSON sql.NullString
if msg.ToolResult != nil {
resultBytes, err := json.Marshal(msg.ToolResult)
if err != nil {
return 0, fmt.Errorf("marshal tool result: %w", err)
}
toolResultJSON = sql.NullString{String: string(resultBytes), Valid: true}
}
result, err := tx.Exec(query,
chatID,
msg.Role,
msg.Content,
msg.Thinking,
msg.Stream,
modelName,
msg.CreatedAt,
msg.UpdatedAt,
thinkingTimeStart,
thinkingTimeEnd,
toolResultJSON,
)
if err != nil {
return 0, err
}
messageID, err := result.LastInsertId()
if err != nil {
return 0, err
}
for _, att := range msg.Attachments {
err := db.insertAttachment(tx, messageID, att)
if err != nil {
return 0, fmt.Errorf("insert attachment: %w", err)
}
}
return messageID, nil
}
func (db *database) getAttachments(messageID int64, loadData bool) ([]File, error) {
var query string
if loadData {
query = `
SELECT filename, data
FROM attachments
WHERE message_id = ?
ORDER BY id ASC
`
} else {
query = `
SELECT filename, '' as data
FROM attachments
WHERE message_id = ?
ORDER BY id ASC
`
}
rows, err := db.conn.Query(query, messageID)
if err != nil {
return nil, fmt.Errorf("query attachments: %w", err)
}
defer rows.Close()
var attachments []File
for rows.Next() {
var file File
err := rows.Scan(&file.Filename, &file.Data)
if err != nil {
return nil, fmt.Errorf("scan attachment: %w", err)
}
attachments = append(attachments, file)
}
if err := rows.Err(); err != nil {
return nil, fmt.Errorf("iterate attachments: %w", err)
}
return attachments, nil
}
func (db *database) getToolCalls(messageID int64) ([]ToolCall, error) {
query := `
SELECT type, function_name, function_arguments, function_result
FROM tool_calls
WHERE message_id = ?
ORDER BY id ASC
`
rows, err := db.conn.Query(query, messageID)
if err != nil {
return nil, fmt.Errorf("query tool calls: %w", err)
}
defer rows.Close()
var toolCalls []ToolCall
for rows.Next() {
var tc ToolCall
var functionResult sql.NullString
err := rows.Scan(
&tc.Type,
&tc.Function.Name,
&tc.Function.Arguments,
&functionResult,
)
if err != nil {
return nil, fmt.Errorf("scan tool call: %w", err)
}
if functionResult.Valid && functionResult.String != "" {
// Parse the JSON result
var result json.RawMessage
if err := json.Unmarshal([]byte(functionResult.String), &result); err == nil {
tc.Function.Result = &result
}
}
toolCalls = append(toolCalls, tc)
}
if err := rows.Err(); err != nil {
return nil, fmt.Errorf("iterate tool calls: %w", err)
}
return toolCalls, nil
}
func (db *database) insertAttachment(tx *sql.Tx, messageID int64, file File) error {
query := `
INSERT INTO attachments (message_id, filename, data)
VALUES (?, ?, ?)
`
_, err := tx.Exec(query, messageID, file.Filename, file.Data)
return err
}
func (db *database) insertToolCall(tx *sql.Tx, messageID int64, tc ToolCall) error {
query := `
INSERT INTO tool_calls (message_id, type, function_name, function_arguments, function_result)
VALUES (?, ?, ?, ?, ?)
`
var functionResult sql.NullString
if tc.Function.Result != nil {
// Convert result to JSON
resultJSON, err := json.Marshal(tc.Function.Result)
if err != nil {
return fmt.Errorf("marshal tool result: %w", err)
}
functionResult = sql.NullString{String: string(resultJSON), Valid: true}
}
_, err := tx.Exec(query,
messageID,
tc.Type,
tc.Function.Name,
tc.Function.Arguments,
functionResult,
)
return err
}
// Settings operations
func (db *database) getID() (string, error) {
var id string
err := db.conn.QueryRow("SELECT device_id FROM settings").Scan(&id)
if err != nil {
return "", fmt.Errorf("get device id: %w", err)
}
return id, nil
}
func (db *database) setID(id string) error {
_, err := db.conn.Exec("UPDATE settings SET device_id = ?", id)
if err != nil {
return fmt.Errorf("set device id: %w", err)
}
return nil
}
func (db *database) getHasCompletedFirstRun() (bool, error) {
var hasCompletedFirstRun bool
err := db.conn.QueryRow("SELECT has_completed_first_run FROM settings").Scan(&hasCompletedFirstRun)
if err != nil {
return false, fmt.Errorf("get has completed first run: %w", err)
}
return hasCompletedFirstRun, nil
}
func (db *database) setHasCompletedFirstRun(hasCompletedFirstRun bool) error {
_, err := db.conn.Exec("UPDATE settings SET has_completed_first_run = ?", hasCompletedFirstRun)
if err != nil {
return fmt.Errorf("set has completed first run: %w", err)
}
return nil
}
func (db *database) getSettings() (Settings, error) {
var s Settings
err := db.conn.QueryRow(`
SELECT expose, survey, browser, models, agent, tools, working_dir, context_length, airplane_mode, turbo_enabled, websearch_enabled, selected_model, sidebar_open, think_enabled, think_level
FROM settings
`).Scan(&s.Expose, &s.Survey, &s.Browser, &s.Models, &s.Agent, &s.Tools, &s.WorkingDir, &s.ContextLength, &s.AirplaneMode, &s.TurboEnabled, &s.WebSearchEnabled, &s.SelectedModel, &s.SidebarOpen, &s.ThinkEnabled, &s.ThinkLevel)
if err != nil {
return Settings{}, fmt.Errorf("get settings: %w", err)
}
return s, nil
}
func (db *database) setSettings(s Settings) error {
_, err := db.conn.Exec(`
UPDATE settings
SET expose = ?, survey = ?, browser = ?, models = ?, agent = ?, tools = ?, working_dir = ?, context_length = ?, airplane_mode = ?, turbo_enabled = ?, websearch_enabled = ?, selected_model = ?, sidebar_open = ?, think_enabled = ?, think_level = ?
`, s.Expose, s.Survey, s.Browser, s.Models, s.Agent, s.Tools, s.WorkingDir, s.ContextLength, s.AirplaneMode, s.TurboEnabled, s.WebSearchEnabled, s.SelectedModel, s.SidebarOpen, s.ThinkEnabled, s.ThinkLevel)
if err != nil {
return fmt.Errorf("set settings: %w", err)
}
return nil
}
func (db *database) getWindowSize() (int, int, error) {
var width, height int
err := db.conn.QueryRow("SELECT window_width, window_height FROM settings").Scan(&width, &height)
if err != nil {
return 0, 0, fmt.Errorf("get window size: %w", err)
}
return width, height, nil
}
func (db *database) setWindowSize(width, height int) error {
_, err := db.conn.Exec("UPDATE settings SET window_width = ?, window_height = ?", width, height)
if err != nil {
return fmt.Errorf("set window size: %w", err)
}
return nil
}
func (db *database) isConfigMigrated() (bool, error) {
var migrated bool
err := db.conn.QueryRow("SELECT config_migrated FROM settings").Scan(&migrated)
if err != nil {
return false, fmt.Errorf("get config migrated: %w", err)
}
return migrated, nil
}
func (db *database) setConfigMigrated(migrated bool) error {
_, err := db.conn.Exec("UPDATE settings SET config_migrated = ?", migrated)
if err != nil {
return fmt.Errorf("set config migrated: %w", err)
}
return nil
}
func (db *database) getSchemaVersion() (int, error) {
var version int
err := db.conn.QueryRow("SELECT schema_version FROM settings").Scan(&version)
if err != nil {
return 0, fmt.Errorf("get schema version: %w", err)
}
return version, nil
}
func (db *database) setSchemaVersion(version int) error {
_, err := db.conn.Exec("UPDATE settings SET schema_version = ?", version)
if err != nil {
return fmt.Errorf("set schema version: %w", err)
}
return nil
}
func (db *database) getUser() (*User, error) {
var user User
err := db.conn.QueryRow(`
SELECT name, email, plan, cached_at
FROM users
LIMIT 1
`).Scan(&user.Name, &user.Email, &user.Plan, &user.CachedAt)
if err != nil {
if err == sql.ErrNoRows {
return nil, nil // No user cached yet
}
return nil, fmt.Errorf("get user: %w", err)
}
return &user, nil
}
func (db *database) setUser(user User) error {
if err := db.clearUser(); err != nil {
return fmt.Errorf("before set: %w", err)
}
_, err := db.conn.Exec(`
INSERT INTO users (name, email, plan, cached_at)
VALUES (?, ?, ?, ?)
`, user.Name, user.Email, user.Plan, user.CachedAt)
if err != nil {
return fmt.Errorf("set user: %w", err)
}
return nil
}
func (db *database) clearUser() error {
_, err := db.conn.Exec("DELETE FROM users")
if err != nil {
return fmt.Errorf("clear user: %w", err)
}
return nil
}
app/store/database_test.go 0 → 100644
View file @ d3b4b997
//go:build windows || darwin
package store
import (
"database/sql"
"fmt"
"os"
"path/filepath"
"sort"
"strings"
"testing"
"time"
"github.com/google/go-cmp/cmp"
_ "github.com/mattn/go-sqlite3"
)
func TestSchemaMigrations(t *testing.T) {
t.Run("schema comparison after migration", func(t *testing.T) {
tmpDir := t.TempDir()
migratedDBPath := filepath.Join(tmpDir, "migrated.db")
migratedDB := loadV2Schema(t, migratedDBPath)
defer migratedDB.Close()
if err := migratedDB.migrate(); err != nil {
t.Fatalf("migration failed: %v", err)
}
// Create fresh database with current schema
freshDBPath := filepath.Join(tmpDir, "fresh.db")
freshDB, err := newDatabase(freshDBPath)
if err != nil {
t.Fatalf("failed to create fresh database: %v", err)
}
defer freshDB.Close()
// Extract tables and indexes from both databases, directly comparing their schemas won't work due to ordering
migratedSchema := schemaMap(migratedDB)
freshSchema := schemaMap(freshDB)
if !cmp.Equal(migratedSchema, freshSchema) {
t.Errorf("Schema difference found:\n%s", cmp.Diff(freshSchema, migratedSchema))
}
// Verify both databases have the same final schema version
migratedVersion, _ := migratedDB.getSchemaVersion()
freshVersion, _ := freshDB.getSchemaVersion()
if migratedVersion != freshVersion {
t.Errorf("schema version mismatch: migrated=%d, fresh=%d", migratedVersion, freshVersion)
}
})
t.Run("idempotent migrations", func(t *testing.T) {
tmpDir := t.TempDir()
dbPath := filepath.Join(tmpDir, "test.db")
db := loadV2Schema(t, dbPath)
defer db.Close()
// Run migration twice
if err := db.migrate(); err != nil {
t.Fatalf("first migration failed: %v", err)
}
if err := db.migrate(); err != nil {
t.Fatalf("second migration failed: %v", err)
}
// Verify schema version is still correct
version, err := db.getSchemaVersion()
if err != nil {
t.Fatalf("failed to get schema version: %v", err)
}
if version != currentSchemaVersion {
t.Errorf("expected schema version %d after double migration, got %d", currentSchemaVersion, version)
}
})
t.Run("init database has correct schema version", func(t *testing.T) {
tmpDir := t.TempDir()
dbPath := filepath.Join(tmpDir, "test.db")
db, err := newDatabase(dbPath)
if err != nil {
t.Fatalf("failed to create database: %v", err)
}
defer db.Close()
// Get the schema version from the newly initialized database
version, err := db.getSchemaVersion()
if err != nil {
t.Fatalf("failed to get schema version: %v", err)
}
// Verify it matches the currentSchemaVersion constant
if version != currentSchemaVersion {
t.Errorf("expected schema version %d in initialized database, got %d", currentSchemaVersion, version)
}
})
}
func TestChatDeletionWithCascade(t *testing.T) {
t.Run("chat deletion cascades to related messages", func(t *testing.T) {
tmpDir := t.TempDir()
dbPath := filepath.Join(tmpDir, "test.db")
db, err := newDatabase(dbPath)
if err != nil {
t.Fatalf("failed to create database: %v", err)
}
defer db.Close()
// Create test chat
testChatID := "test-chat-cascade-123"
testChat := Chat{
ID: testChatID,
Title: "Test Chat for Cascade Delete",
CreatedAt: time.Now(),
Messages: []Message{
{
Role: "user",
Content: "Hello, this is a test message",
CreatedAt: time.Now(),
UpdatedAt: time.Now(),
},
{
Role: "assistant",
Content: "Hi there! This is a response.",
CreatedAt: time.Now(),
UpdatedAt: time.Now(),
},
},
}
// Save the chat with messages
if err := db.saveChat(testChat); err != nil {
t.Fatalf("failed to save test chat: %v", err)
}
// Verify chat and messages exist
chatCount := countRows(t, db, "chats")
messageCount := countRows(t, db, "messages")
if chatCount != 1 {
t.Errorf("expected 1 chat, got %d", chatCount)
}
if messageCount != 2 {
t.Errorf("expected 2 messages, got %d", messageCount)
}
// Verify specific chat exists
var exists bool
err = db.conn.QueryRow("SELECT EXISTS(SELECT 1 FROM chats WHERE id = ?)", testChatID).Scan(&exists)
if err != nil {
t.Fatalf("failed to check chat existence: %v", err)
}
if !exists {
t.Error("test chat should exist before deletion")
}
// Verify messages exist for this chat
messageCountForChat := countRowsWithCondition(t, db, "messages", "chat_id = ?", testChatID)
if messageCountForChat != 2 {
t.Errorf("expected 2 messages for test chat, got %d", messageCountForChat)
}
// Delete the chat
if err := db.deleteChat(testChatID); err != nil {
t.Fatalf("failed to delete chat: %v", err)
}
// Verify chat is deleted
chatCountAfter := countRows(t, db, "chats")
if chatCountAfter != 0 {
t.Errorf("expected 0 chats after deletion, got %d", chatCountAfter)
}
// Verify messages are CASCADE deleted
messageCountAfter := countRows(t, db, "messages")
if messageCountAfter != 0 {
t.Errorf("expected 0 messages after CASCADE deletion, got %d", messageCountAfter)
}
// Verify specific chat no longer exists
err = db.conn.QueryRow("SELECT EXISTS(SELECT 1 FROM chats WHERE id = ?)", testChatID).Scan(&exists)
if err != nil {
t.Fatalf("failed to check chat existence after deletion: %v", err)
}
if exists {
t.Error("test chat should not exist after deletion")
}
// Verify no orphaned messages remain
orphanedCount := countRowsWithCondition(t, db, "messages", "chat_id = ?", testChatID)
if orphanedCount != 0 {
t.Errorf("expected 0 orphaned messages, got %d", orphanedCount)
}
})
t.Run("foreign keys are enabled", func(t *testing.T) {
tmpDir := t.TempDir()
dbPath := filepath.Join(tmpDir, "test.db")
db, err := newDatabase(dbPath)
if err != nil {
t.Fatalf("failed to create database: %v", err)
}
defer db.Close()
// Verify foreign keys are enabled
var foreignKeysEnabled int
err = db.conn.QueryRow("PRAGMA foreign_keys").Scan(&foreignKeysEnabled)
if err != nil {
t.Fatalf("failed to check foreign keys: %v", err)
}
if foreignKeysEnabled != 1 {
t.Errorf("expected foreign keys to be enabled (1), got %d", foreignKeysEnabled)
}
})
// This test is only relevant for v8 migrations, but we keep it here for now
// since it's a useful test to ensure that we don't introduce any new orphaned data
t.Run("cleanup orphaned data", func(t *testing.T) {
tmpDir := t.TempDir()
dbPath := filepath.Join(tmpDir, "test.db")
db, err := newDatabase(dbPath)
if err != nil {
t.Fatalf("failed to create database: %v", err)
}
defer db.Close()
// First disable foreign keys to simulate the bug from ollama/ollama#11785
_, err = db.conn.Exec("PRAGMA foreign_keys = OFF")
if err != nil {
t.Fatalf("failed to disable foreign keys: %v", err)
}
// Create a chat and message
testChatID := "orphaned-test-chat"
testMessageID := int64(999)
_, err = db.conn.Exec("INSERT INTO chats (id, title) VALUES (?, ?)", testChatID, "Orphaned Test Chat")
if err != nil {
t.Fatalf("failed to insert test chat: %v", err)
}
_, err = db.conn.Exec("INSERT INTO messages (id, chat_id, role, content) VALUES (?, ?, ?, ?)",
testMessageID, testChatID, "user", "test message")
if err != nil {
t.Fatalf("failed to insert test message: %v", err)
}
// Delete chat but keep message (simulating the bug from ollama/ollama#11785)
_, err = db.conn.Exec("DELETE FROM chats WHERE id = ?", testChatID)
if err != nil {
t.Fatalf("failed to delete chat: %v", err)
}
// Verify we have orphaned message
orphanedCount := countRowsWithCondition(t, db, "messages", "chat_id = ?", testChatID)
if orphanedCount != 1 {
t.Errorf("expected 1 orphaned message, got %d", orphanedCount)
}
// Run cleanup
if err := db.cleanupOrphanedData(); err != nil {
t.Fatalf("failed to cleanup orphaned data: %v", err)
}
// Verify orphaned message is gone
orphanedCountAfter := countRowsWithCondition(t, db, "messages", "chat_id = ?", testChatID)
if orphanedCountAfter != 0 {
t.Errorf("expected 0 orphaned messages after cleanup, got %d", orphanedCountAfter)
}
})
}
func countRows(t *testing.T, db *database, table string) int {
t.Helper()
var count int
err := db.conn.QueryRow(fmt.Sprintf("SELECT COUNT(*) FROM %s", table)).Scan(&count)
if err != nil {
t.Fatalf("failed to count rows in %s: %v", table, err)
}
return count
}
func countRowsWithCondition(t *testing.T, db *database, table, condition string, args ...interface{}) int {
t.Helper()
var count int
query := fmt.Sprintf("SELECT COUNT(*) FROM %s WHERE %s", table, condition)
err := db.conn.QueryRow(query, args...).Scan(&count)
if err != nil {
t.Fatalf("failed to count rows with condition: %v", err)
}
return count
}
// Test helpers for schema migration testing
// schemaMap returns both tables/columns and indexes (ignoring order)
func schemaMap(db *database) map[string]interface{} {
result := make(map[string]any)
result["tables"] = columnMap(db)
result["indexes"] = indexMap(db)
return result
}
// columnMap returns a map of table names to their column sets (ignoring order)
func columnMap(db *database) map[string][]string {
result := make(map[string][]string)
// Get all table names
tableQuery := `SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%' ORDER BY name`
rows, _ := db.conn.Query(tableQuery)
defer rows.Close()
for rows.Next() {
var tableName string
rows.Scan(&tableName)
// Get columns for this table
colQuery := fmt.Sprintf("PRAGMA table_info(%s)", tableName)
colRows, _ := db.conn.Query(colQuery)
var columns []string
for colRows.Next() {
var cid int
var name, dataType sql.NullString
var notNull, primaryKey int
var defaultValue sql.NullString
colRows.Scan(&cid, &name, &dataType, &notNull, &defaultValue, &primaryKey)
// Create a normalized column description
colDesc := fmt.Sprintf("%s %s", name.String, dataType.String)
if notNull == 1 {
colDesc += " NOT NULL"
}
if defaultValue.Valid && defaultValue.String != "" {
// Skip DEFAULT for schema_version as it doesn't get updated during migrations
if name.String != "schema_version" {
colDesc += " DEFAULT " + defaultValue.String
}
}
if primaryKey == 1 {
colDesc += " PRIMARY KEY"
}
columns = append(columns, colDesc)
}
colRows.Close()
// Sort columns to ignore order differences
sort.Strings(columns)
result[tableName] = columns
}
return result
}
// indexMap returns a map of index names to their definitions
func indexMap(db *database) map[string]string {
result := make(map[string]string)
// Get all indexes (excluding auto-created primary key indexes)
indexQuery := `SELECT name, sql FROM sqlite_master WHERE type='index' AND name NOT LIKE 'sqlite_%' AND sql IS NOT NULL ORDER BY name`
rows, _ := db.conn.Query(indexQuery)
defer rows.Close()
for rows.Next() {
var name, sql string
rows.Scan(&name, &sql)
// Normalize the SQL by removing extra whitespace
sql = strings.Join(strings.Fields(sql), " ")
result[name] = sql
}
return result
}
// loadV2Schema loads the version 2 schema from testdata/schema.sql
func loadV2Schema(t *testing.T, dbPath string) *database {
t.Helper()
// Read the v1 schema file
schemaFile := filepath.Join("testdata", "schema.sql")
schemaSQL, err := os.ReadFile(schemaFile)
if err != nil {
t.Fatalf("failed to read schema file: %v", err)
}
// Open database connection
conn, err := sql.Open("sqlite3", dbPath+"?_foreign_keys=on&_journal_mode=WAL&_busy_timeout=5000&_txlock=immediate")
if err != nil {
t.Fatalf("failed to open database: %v", err)
}
// Execute the v1 schema
_, err = conn.Exec(string(schemaSQL))
if err != nil {
conn.Close()
t.Fatalf("failed to execute v1 schema: %v", err)
}
return &database{conn: conn}
}
app/store/image.go 0 → 100644
View file @ d3b4b997
//go:build windows || darwin
package store
import (
"crypto/sha256"
"encoding/hex"
"fmt"
"os"
"path/filepath"
"strings"
)
type Image struct {
Filename string `json:"filename"`
Path string `json:"path"`
Size int64 `json:"size,omitempty"`
MimeType string `json:"mime_type,omitempty"`
}
// Bytes loads image data from disk for a given ImageData reference
func (i *Image) Bytes() ([]byte, error) {
return ImgBytes(i.Path)
}
// ImgBytes reads image data from the specified file path
func ImgBytes(path string) ([]byte, error) {
if path == "" {
return nil, fmt.Errorf("empty image path")
}
data, err := os.ReadFile(path)
if err != nil {
return nil, fmt.Errorf("read image file %s: %w", path, err)
}
return data, nil
}
// ImgDir returns the directory path for storing images for a specific chat
func (s *Store) ImgDir() string {
dbPath := s.DBPath
if dbPath == "" {
dbPath = defaultDBPath
}
storeDir := filepath.Dir(dbPath)
return filepath.Join(storeDir, "cache", "images")
}
// ImgToFile saves image data to disk and returns ImageData reference
func (s *Store) ImgToFile(chatID string, imageBytes []byte, filename, mimeType string) (Image, error) {
baseImageDir := s.ImgDir()
if err := os.MkdirAll(baseImageDir, 0o755); err != nil {
return Image{}, fmt.Errorf("create base image directory: %w", err)
}
// Root prevents path traversal issues
root, err := os.OpenRoot(baseImageDir)
if err != nil {
return Image{}, fmt.Errorf("open image root directory: %w", err)
}
defer root.Close()
// Create chat-specific subdirectory within the root
chatDir := sanitize(chatID)
if err := root.Mkdir(chatDir, 0o755); err != nil && !os.IsExist(err) {
return Image{}, fmt.Errorf("create chat directory: %w", err)
}
// Generate a unique filename to avoid conflicts
// Use hash of content + original filename for uniqueness
hash := sha256.Sum256(imageBytes)
hashStr := hex.EncodeToString(hash[:])[:16] // Use first 16 chars of hash
// Extract file extension from original filename or mime type
ext := filepath.Ext(filename)
if ext == "" {
switch mimeType {
case "image/jpeg":
ext = ".jpg"
case "image/png":
ext = ".png"
case "image/webp":
ext = ".webp"
default:
ext = ".img"
}
}
// Create unique filename: hash + original name + extension
baseFilename := sanitize(strings.TrimSuffix(filename, ext))
uniqueFilename := fmt.Sprintf("%s_%s%s", hashStr, baseFilename, ext)
relativePath := filepath.Join(chatDir, uniqueFilename)
file, err := root.Create(relativePath)
if err != nil {
return Image{}, fmt.Errorf("create image file: %w", err)
}
defer file.Close()
if _, err := file.Write(imageBytes); err != nil {
return Image{}, fmt.Errorf("write image data: %w", err)
}
return Image{
Filename: uniqueFilename,
Path: filepath.Join(baseImageDir, relativePath),
Size: int64(len(imageBytes)),
MimeType: mimeType,
}, nil
}
// sanitize removes unsafe characters from filenames
func sanitize(filename string) string {
// Convert to safe characters only
safe := strings.Map(func(r rune) rune {
if (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') || (r >= '0' && r <= '9') || r == '-' {
return r
}
return '_'
}, filename)
// Clean up and validate
safe = strings.Trim(safe, "_")
if safe == "" {
return "image"
}
return safe
}
app/store/migration_test.go 0 → 100644
View file @ d3b4b997
//go:build windows || darwin
package store
import (
"database/sql"
"encoding/json"
"os"
"path/filepath"
"testing"
)
func TestConfigMigration(t *testing.T) {
tmpDir := t.TempDir()
// Create a legacy config.json
legacyConfig := legacyData{
ID: "test-device-id-12345",
FirstTimeRun: true, // In old system, true meant "has completed first run"
}
configData, err := json.MarshalIndent(legacyConfig, "", " ")
if err != nil {
t.Fatal(err)
}
configPath := filepath.Join(tmpDir, "config.json")
if err := os.WriteFile(configPath, configData, 0o644); err != nil {
t.Fatal(err)
}
// Override the legacy config path for testing
oldLegacyConfigPath := legacyConfigPath
legacyConfigPath = configPath
defer func() { legacyConfigPath = oldLegacyConfigPath }()
// Create store with database in same directory
s := Store{DBPath: filepath.Join(tmpDir, "db.sqlite")}
defer s.Close()
// First access should trigger migration
id, err := s.ID()
if err != nil {
t.Fatalf("failed to get ID: %v", err)
}
if id != "test-device-id-12345" {
t.Errorf("expected migrated ID 'test-device-id-12345', got '%s'", id)
}
// Check HasCompletedFirstRun
hasCompleted, err := s.HasCompletedFirstRun()
if err != nil {
t.Fatalf("failed to get has completed first run: %v", err)
}
if !hasCompleted {
t.Error("expected has completed first run to be true after migration")
}
// Verify migration is marked as complete
migrated, err := s.db.isConfigMigrated()
if err != nil {
t.Fatalf("failed to check migration status: %v", err)
}
if !migrated {
t.Error("expected config to be marked as migrated")
}
// Create a new store instance to verify migration doesn't run again
s2 := Store{DBPath: filepath.Join(tmpDir, "db.sqlite")}
defer s2.Close()
// Delete the config file to ensure we're not reading from it
os.Remove(configPath)
// Verify data is still there
id2, err := s2.ID()
if err != nil {
t.Fatalf("failed to get ID from second store: %v", err)
}
if id2 != "test-device-id-12345" {
t.Errorf("expected persisted ID 'test-device-id-12345', got '%s'", id2)
}
}
func TestNoConfigToMigrate(t *testing.T) {
tmpDir := t.TempDir()
// Override the legacy config path for testing
oldLegacyConfigPath := legacyConfigPath
legacyConfigPath = filepath.Join(tmpDir, "config.json")
defer func() { legacyConfigPath = oldLegacyConfigPath }()
// Create store without any config.json
s := Store{DBPath: filepath.Join(tmpDir, "db.sqlite")}
defer s.Close()
// Should generate a new ID
id, err := s.ID()
if err != nil {
t.Fatalf("failed to get ID: %v", err)
}
if id == "" {
t.Error("expected auto-generated ID, got empty string")
}
// HasCompletedFirstRun should be false (default)
hasCompleted, err := s.HasCompletedFirstRun()
if err != nil {
t.Fatalf("failed to get has completed first run: %v", err)
}
if hasCompleted {
t.Error("expected has completed first run to be false by default")
}
// Migration should still be marked as complete
migrated, err := s.db.isConfigMigrated()
if err != nil {
t.Fatalf("failed to check migration status: %v", err)
}
if !migrated {
t.Error("expected config to be marked as migrated even with no config.json")
}
}
const (
v1Schema = `
CREATE TABLE IF NOT EXISTS settings (
id INTEGER PRIMARY KEY CHECK (id = 1),
device_id TEXT NOT NULL DEFAULT '',
has_completed_first_run BOOLEAN NOT NULL DEFAULT 0,
expose BOOLEAN NOT NULL DEFAULT 0,
browser BOOLEAN NOT NULL DEFAULT 0,
models TEXT NOT NULL DEFAULT '',
remote TEXT NOT NULL DEFAULT '',
agent BOOLEAN NOT NULL DEFAULT 0,
tools BOOLEAN NOT NULL DEFAULT 0,
working_dir TEXT NOT NULL DEFAULT '',
window_width INTEGER NOT NULL DEFAULT 0,
window_height INTEGER NOT NULL DEFAULT 0,
config_migrated BOOLEAN NOT NULL DEFAULT 0,
schema_version INTEGER NOT NULL DEFAULT 1
);
-- Insert default settings row if it doesn't exist
INSERT OR IGNORE INTO settings (id) VALUES (1);
CREATE TABLE IF NOT EXISTS chats (
id TEXT PRIMARY KEY,
title TEXT NOT NULL DEFAULT '',
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
);
CREATE TABLE IF NOT EXISTS 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,
FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS idx_messages_chat_id ON messages(chat_id);
CREATE TABLE IF NOT EXISTS tool_calls (
id INTEGER PRIMARY KEY AUTOINCREMENT,
message_id INTEGER NOT NULL,
type TEXT NOT NULL,
function_name TEXT NOT NULL,
function_arguments TEXT NOT NULL,
function_result TEXT,
FOREIGN KEY (message_id) REFERENCES messages(id) ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS idx_tool_calls_message_id ON tool_calls(message_id);
`
)
func TestMigrationFromEpoc(t *testing.T) {
tmpDir := t.TempDir()
s := Store{DBPath: filepath.Join(tmpDir, "db.sqlite")}
defer s.Close()
// Open database connection
conn, err := sql.Open("sqlite3", s.DBPath+"?_foreign_keys=on&_journal_mode=WAL")
if err != nil {
t.Fatal(err)
}
// Test the connection
if err := conn.Ping(); err != nil {
conn.Close()
t.Fatal(err)
}
s.db = &database{conn: conn}
t.Logf("DB created: %s", s.DBPath)
_, err = s.db.conn.Exec(v1Schema)
if err != nil {
t.Fatal(err)
}
version, err := s.db.getSchemaVersion()
if err != nil {
t.Fatalf("failed to get schema version: %v", err)
}
if version != 1 {
t.Fatalf("expected: %d\n got: %d", 1, version)
}
t.Logf("v1 schema created")
if err := s.db.migrate(); err != nil {
t.Fatal(err)
}
t.Logf("migrations completed")
version, err = s.db.getSchemaVersion()
if err != nil {
t.Fatalf("failed to get schema version: %v", err)
}
if version != currentSchemaVersion {
t.Fatalf("expected: %d\n got: %d", currentSchemaVersion, version)
}
}
app/store/schema.sql 0 → 100644
View file @ d3b4b997
-- This is the version 2 schema for the app database, the first released schema to users.
-- Do not modify this file. It is used to test that the database schema stays in a consistent state between schema migrations.
CREATE TABLE IF NOT EXISTS settings (
id INTEGER PRIMARY KEY CHECK (id = 1),
device_id TEXT NOT NULL DEFAULT '',
has_completed_first_run BOOLEAN NOT NULL DEFAULT 0,
expose BOOLEAN NOT NULL DEFAULT 0,
survey BOOLEAN NOT NULL DEFAULT TRUE,
browser BOOLEAN NOT NULL DEFAULT 0,
models TEXT NOT NULL DEFAULT '',
remote TEXT NOT NULL DEFAULT '',
agent BOOLEAN NOT NULL DEFAULT 0,
tools BOOLEAN NOT NULL DEFAULT 0,
working_dir TEXT NOT NULL DEFAULT '',
context_length INTEGER NOT NULL DEFAULT 4096,
window_width INTEGER NOT NULL DEFAULT 0,
window_height INTEGER NOT NULL DEFAULT 0,
config_migrated BOOLEAN NOT NULL DEFAULT 0,
schema_version INTEGER NOT NULL DEFAULT 2
);
-- Insert default settings row if it doesn't exist
INSERT OR IGNORE INTO settings (id) VALUES (1);
CREATE TABLE IF NOT EXISTS chats (
id TEXT PRIMARY KEY,
title TEXT NOT NULL DEFAULT '',
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
);
CREATE TABLE IF NOT EXISTS 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,
FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS idx_messages_chat_id ON messages(chat_id);
CREATE TABLE IF NOT EXISTS tool_calls (
id INTEGER PRIMARY KEY AUTOINCREMENT,
message_id INTEGER NOT NULL,
type TEXT NOT NULL,
function_name TEXT NOT NULL,
function_arguments TEXT NOT NULL,
function_result TEXT,
FOREIGN KEY (message_id) REFERENCES messages(id) ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS idx_tool_calls_message_id ON tool_calls(message_id);
app/store/schema_test.go 0 → 100644
View file @ d3b4b997
//go:build windows || darwin
package store
import (
"path/filepath"
"testing"
)
func TestSchemaVersioning(t *testing.T) {
tmpDir := t.TempDir()
// Override legacy config path to avoid migration logs
oldLegacyConfigPath := legacyConfigPath
legacyConfigPath = filepath.Join(tmpDir, "config.json")
defer func() { legacyConfigPath = oldLegacyConfigPath }()
t.Run("new database has correct schema version", func(t *testing.T) {
dbPath := filepath.Join(tmpDir, "new_db.sqlite")
db, err := newDatabase(dbPath)
if err != nil {
t.Fatalf("failed to create database: %v", err)
}
defer db.Close()
// Check schema version
version, err := db.getSchemaVersion()
if err != nil {
t.Fatalf("failed to get schema version: %v", err)
}
if version != currentSchemaVersion {
t.Errorf("expected schema version %d, got %d", currentSchemaVersion, version)
}
})
t.Run("can update schema version", func(t *testing.T) {
dbPath := filepath.Join(tmpDir, "update_db.sqlite")
db, err := newDatabase(dbPath)
if err != nil {
t.Fatalf("failed to create database: %v", err)
}
defer db.Close()
// Set a different version
testVersion := 42
if err := db.setSchemaVersion(testVersion); err != nil {
t.Fatalf("failed to set schema version: %v", err)
}
// Verify it was updated
version, err := db.getSchemaVersion()
if err != nil {
t.Fatalf("failed to get schema version: %v", err)
}
if version != testVersion {
t.Errorf("expected schema version %d, got %d", testVersion, version)
}
})
}
app/store/store.go
View file @ d3b4b997
//go:build windows || darwin
// Package store provides a simple JSON file store for the desktop application
// to save and load data such as ollama server configuration, messages,
// login information and more.
package store
import (
"encoding/json"
"errors"
"fmt"
"log/slog"
"os"
"path/filepath"
"runtime"
"sync"
"time"
"github.com/google/uuid"
"github.com/ollama/ollama/app/types/not"
)
type Store struct {
ID string `json:"id"`
FirstTimeRun bool `json:"first-time-run"`
type File struct {
Filename string `json:"filename"`
Data []byte `json:"data"`
}
var (
lock sync.Mutex
store Store
)
type User struct {
Name string `json:"name"`
Email string `json:"email"`
Plan string `json:"plan"`
CachedAt time.Time `json:"cachedAt"`
}
type Message struct {
Role string `json:"role"`
Content string `json:"content"`
Thinking string `json:"thinking"`
Stream bool `json:"stream"`
Model string `json:"model,omitempty"`
Attachments []File `json:"attachments,omitempty"`
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
ToolCall *ToolCall `json:"tool_call,omitempty"`
ToolName string `json:"tool_name,omitempty"`
ToolResult *json.RawMessage `json:"tool_result,omitempty"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
ThinkingTimeStart *time.Time `json:"thinkingTimeStart,omitempty" ts_type:"Date | undefined" ts_transform:"__VALUE__ && new Date(__VALUE__)"`
ThinkingTimeEnd *time.Time `json:"thinkingTimeEnd,omitempty" ts_type:"Date | undefined" ts_transform:"__VALUE__ && new Date(__VALUE__)"`
}
// MessageOptions contains optional parameters for creating a Message
type MessageOptions struct {
Model string
Attachments []File
Stream bool
Thinking string
ToolCalls []ToolCall
ToolCall *ToolCall
ToolResult *json.RawMessage
ThinkingTimeStart *time.Time
ThinkingTimeEnd *time.Time
}
func GetID() string {
lock.Lock()
defer lock.Unlock()
if store.ID == "" {
initStore()
// NewMessage creates a new Message with the given options
func NewMessage(role, content string, opts *MessageOptions) Message {
now := time.Now()
msg := Message{
Role: role,
Content: content,
CreatedAt: now,
UpdatedAt: now,
}
return store.ID
if opts != nil {
msg.Model = opts.Model
msg.Attachments = opts.Attachments
msg.Stream = opts.Stream
msg.Thinking = opts.Thinking
msg.ToolCalls = opts.ToolCalls
msg.ToolCall = opts.ToolCall
msg.ToolResult = opts.ToolResult
msg.ThinkingTimeStart = opts.ThinkingTimeStart
msg.ThinkingTimeEnd = opts.ThinkingTimeEnd
}
return msg
}
type ToolCall struct {
Type string `json:"type"`
Function ToolFunction `json:"function"`
}
type ToolFunction struct {
Name string `json:"name"`
Arguments string `json:"arguments"`
Result any `json:"result,omitempty"`
}
type Model struct {
Model string `json:"model"` // Model name
Digest string `json:"digest,omitempty"` // Model digest from the registry
ModifiedAt *time.Time `json:"modified_at,omitempty"` // When the model was last modified locally
}
type Chat struct {
ID string `json:"id"`
Messages []Message `json:"messages"`
Title string `json:"title"`
CreatedAt time.Time `json:"created_at"`
BrowserState json.RawMessage `json:"browser_state,omitempty" ts_type:"BrowserStateData"`
}
func GetFirstTimeRun() bool {
lock.Lock()
defer lock.Unlock()
if store.ID == "" {
initStore()
// NewChat creates a new Chat with the ID, with CreatedAt timestamp initialized
func NewChat(id string) *Chat {
return &Chat{
ID: id,
Messages: []Message{},
CreatedAt: time.Now(),
}
return store.FirstTimeRun
}
func SetFirstTimeRun(val bool) {
lock.Lock()
defer lock.Unlock()
if store.FirstTimeRun == val {
return
type Settings struct {
// Expose is a boolean that indicates if the ollama server should
// be exposed to the network
Expose bool
// Browser is a boolean that indicates if the ollama server should
// be exposed to browser windows (e.g. CORS set to allow all origins)
Browser bool
// Survey is a boolean that indicates if the user allows anonymous
// inference information to be shared with Ollama
Survey bool
// Models is a string that contains the models to load on startup
Models string
// TODO(parthsareen): temporary for experimentation
// Agent indicates if the app should use multi-turn tools to fulfill user requests
Agent bool
// Tools indicates if the app should use single-turn tools to fulfill user requests
Tools bool
// WorkingDir specifies the working directory for all agent operations
WorkingDir string
// ContextLength specifies the context length for the ollama server (using OLLAMA_CONTEXT_LENGTH)
ContextLength int
// AirplaneMode when true, turns off Ollama Turbo features and only uses local models
AirplaneMode bool
// TurboEnabled indicates if Ollama Turbo features are enabled
TurboEnabled bool
// Maps gpt-oss specific frontend name' BrowserToolEnabled' to db field 'websearch_enabled'
WebSearchEnabled bool
// ThinkEnabled indicates if thinking is enabled
ThinkEnabled bool
// ThinkLevel indicates the level of thinking to use for models that support multiple levels
ThinkLevel string
// SelectedModel stores the last model that the user selected
SelectedModel string
// SidebarOpen indicates if the chat sidebar is open
SidebarOpen bool
}
type Store struct {
// DBPath allows overriding the default database path (mainly for testing)
DBPath string
// dbMu protects database initialization only
dbMu sync.Mutex
db *database
}
var defaultDBPath = func() string {
switch runtime.GOOS {
case "windows":
return filepath.Join(os.Getenv("LOCALAPPDATA"), "Ollama", "db.sqlite")
case "darwin":
return filepath.Join(os.Getenv("HOME"), "Library", "Application Support", "Ollama", "db.sqlite")
default:
return filepath.Join(os.Getenv("HOME"), ".ollama", "db.sqlite")
}
}()
// legacyConfigPath is the path to the old config.json file
var legacyConfigPath = func() string {
switch runtime.GOOS {
case "windows":
return filepath.Join(os.Getenv("LOCALAPPDATA"), "Ollama", "config.json")
case "darwin":
return filepath.Join(os.Getenv("HOME"), "Library", "Application Support", "Ollama", "config.json")
default:
return filepath.Join(os.Getenv("HOME"), ".ollama", "config.json")
}
store.FirstTimeRun = val
writeStore(getStorePath())
}()
// legacyData represents the old config.json structure (only fields we need to migrate)
type legacyData struct {
ID string `json:"id"`
FirstTimeRun bool `json:"first-time-run"`
}
// lock must be held
func initStore() {
storeFile, err := os.Open(getStorePath())
if err == nil {
defer storeFile.Close()
err = json.NewDecoder(storeFile).Decode(&store)
func (s *Store) ensureDB() error {
// Fast path: check if db is already initialized
if s.db != nil {
return nil
}
// Slow path: initialize database with lock
s.dbMu.Lock()
defer s.dbMu.Unlock()
// Double-check after acquiring lock
if s.db != nil {
return nil
}
dbPath := s.DBPath
if dbPath == "" {
dbPath = defaultDBPath
}
// Ensure directory exists
if err := os.MkdirAll(filepath.Dir(dbPath), 0o755); err != nil {
return fmt.Errorf("create db directory: %w", err)
}
database, err := newDatabase(dbPath)
if err != nil {
return fmt.Errorf("open database: %w", err)
}
// Generate device ID if needed
id, err := database.getID()
if err != nil || id == "" {
// Generate new UUID for device
u, err := uuid.NewV7()
if err == nil {
slog.Debug(fmt.Sprintf("loaded existing store %s - ID: %s", getStorePath(), store.ID))
return
database.setID(u.String())
}
} else if !errors.Is(err, os.ErrNotExist) {
slog.Debug(fmt.Sprintf("unexpected error searching for store: %s", err))
}
slog.Debug("initializing new store")
store.ID = uuid.NewString()
writeStore(getStorePath())
s.db = database
// Check if we need to migrate from config.json
migrated, err := database.isConfigMigrated()
if err != nil || !migrated {
if err := s.migrateFromConfig(database); err != nil {
slog.Warn("failed to migrate from config.json", "error", err)
}
}
return nil
}
func writeStore(storeFilename string) {
ollamaDir := filepath.Dir(storeFilename)
_, err := os.Stat(ollamaDir)
if errors.Is(err, os.ErrNotExist) {
if err := os.MkdirAll(ollamaDir, 0o755); err != nil {
slog.Error(fmt.Sprintf("create ollama dir %s: %v", ollamaDir, err))
return
// migrateFromConfig attempts to migrate ID and FirstTimeRun from config.json
func (s *Store) migrateFromConfig(database *database) error {
configPath := legacyConfigPath
// Check if config.json exists
if _, err := os.Stat(configPath); os.IsNotExist(err) {
// No config to migrate, mark as migrated
return database.setConfigMigrated(true)
}
// Read the config file
b, err := os.ReadFile(configPath)
if err != nil {
return fmt.Errorf("read legacy config: %w", err)
}
var legacy legacyData
if err := json.Unmarshal(b, &legacy); err != nil {
// If we can't parse it, just mark as migrated and move on
slog.Warn("failed to parse legacy config.json", "error", err)
return database.setConfigMigrated(true)
}
// Migrate the ID if present
if legacy.ID != "" {
if err := database.setID(legacy.ID); err != nil {
return fmt.Errorf("migrate device ID: %w", err)
}
slog.Info("migrated device ID from config.json")
}
hasCompleted := legacy.FirstTimeRun // If old FirstTimeRun is true, it means first run was completed
if err := database.setHasCompletedFirstRun(hasCompleted); err != nil {
return fmt.Errorf("migrate first time run: %w", err)
}
slog.Info("migrated first run status from config.json", "hasCompleted", hasCompleted)
// Mark as migrated
if err := database.setConfigMigrated(true); err != nil {
return fmt.Errorf("mark config as migrated: %w", err)
}
slog.Info("successfully migrated settings from config.json")
return nil
}
func (s *Store) ID() (string, error) {
if err := s.ensureDB(); err != nil {
return "", err
}
payload, err := json.Marshal(store)
return s.db.getID()
}
func (s *Store) HasCompletedFirstRun() (bool, error) {
if err := s.ensureDB(); err != nil {
return false, err
}
return s.db.getHasCompletedFirstRun()
}
func (s *Store) SetHasCompletedFirstRun(hasCompleted bool) error {
if err := s.ensureDB(); err != nil {
return err
}
return s.db.setHasCompletedFirstRun(hasCompleted)
}
func (s *Store) Settings() (Settings, error) {
if err := s.ensureDB(); err != nil {
return Settings{}, fmt.Errorf("load settings: %w", err)
}
settings, err := s.db.getSettings()
if err != nil {
slog.Error(fmt.Sprintf("failed to marshal store: %s", err))
return
return Settings{}, err
}
// Set default models directory if not set
if settings.Models == "" {
dir := os.Getenv("OLLAMA_MODELS")
if dir != "" {
settings.Models = dir
} else {
home, err := os.UserHomeDir()
if err == nil {
settings.Models = filepath.Join(home, ".ollama", "models")
}
}
}
return settings, nil
}
func (s *Store) SetSettings(settings Settings) error {
if err := s.ensureDB(); err != nil {
return err
}
fp, err := os.OpenFile(storeFilename, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o755)
return s.db.setSettings(settings)
}
func (s *Store) Chats() ([]Chat, error) {
if err := s.ensureDB(); err != nil {
return nil, err
}
return s.db.getAllChats()
}
func (s *Store) Chat(id string) (*Chat, error) {
return s.ChatWithOptions(id, true)
}
func (s *Store) ChatWithOptions(id string, loadAttachmentData bool) (*Chat, error) {
if err := s.ensureDB(); err != nil {
return nil, err
}
chat, err := s.db.getChatWithOptions(id, loadAttachmentData)
if err != nil {
slog.Error(fmt.Sprintf("write store payload %s: %v", storeFilename, err))
return
return nil, fmt.Errorf("%w: chat %s", not.Found, id)
}
return chat, nil
}
func (s *Store) SetChat(chat Chat) error {
if err := s.ensureDB(); err != nil {
return err
}
defer fp.Close()
if n, err := fp.Write(payload); err != nil || n != len(payload) {
slog.Error(fmt.Sprintf("write store payload %s: %d vs %d -- %v", storeFilename, n, len(payload), err))
return
return s.db.saveChat(chat)
}
func (s *Store) DeleteChat(id string) error {
if err := s.ensureDB(); err != nil {
return err
}
// Delete from database
if err := s.db.deleteChat(id); err != nil {
return fmt.Errorf("%w: chat %s", not.Found, id)
}
// Also delete associated images
chatImgDir := filepath.Join(s.ImgDir(), id)
if err := os.RemoveAll(chatImgDir); err != nil {
// Log error but don't fail the deletion
slog.Warn("failed to delete chat images", "chat_id", id, "error", err)
}
return nil
}
func (s *Store) WindowSize() (int, int, error) {
if err := s.ensureDB(); err != nil {
return 0, 0, err
}
return s.db.getWindowSize()
}
func (s *Store) SetWindowSize(width, height int) error {
if err := s.ensureDB(); err != nil {
return err
}
return s.db.setWindowSize(width, height)
}
func (s *Store) UpdateLastMessage(chatID string, message Message) error {
if err := s.ensureDB(); err != nil {
return err
}
return s.db.updateLastMessage(chatID, message)
}
func (s *Store) AppendMessage(chatID string, message Message) error {
if err := s.ensureDB(); err != nil {
return err
}
return s.db.appendMessage(chatID, message)
}
func (s *Store) UpdateChatBrowserState(chatID string, state json.RawMessage) error {
if err := s.ensureDB(); err != nil {
return err
}
return s.db.updateChatBrowserState(chatID, state)
}
func (s *Store) User() (*User, error) {
if err := s.ensureDB(); err != nil {
return nil, err
}
return s.db.getUser()
}
func (s *Store) SetUser(user User) error {
if err := s.ensureDB(); err != nil {
return err
}
user.CachedAt = time.Now()
return s.db.setUser(user)
}
func (s *Store) ClearUser() error {
if err := s.ensureDB(); err != nil {
return err
}
return s.db.clearUser()
}
func (s *Store) Close() error {
s.dbMu.Lock()
defer s.dbMu.Unlock()
if s.db != nil {
return s.db.Close()
}
slog.Debug("Store contents: " + string(payload))
slog.Info(fmt.Sprintf("wrote store: %s", storeFilename))
return nil
}
app/store/store_darwin.go deleted 100644 → 0
View file @ a4770107
package store
import (
"os"
"path/filepath"
)
func getStorePath() string {
// TODO - system wide location?
home := os.Getenv("HOME")
return filepath.Join(home, "Library", "Application Support", "Ollama", "config.json")
}
app/store/store_linux.go deleted 100644 → 0
View file @ a4770107
package store
import (
"os"
"path/filepath"
)
func getStorePath() string {
if os.Geteuid() == 0 {
// TODO where should we store this on linux for system-wide operation?
return "/etc/ollama/config.json"
}
home := os.Getenv("HOME")
return filepath.Join(home, ".ollama", "config.json")
}
app/store/store_test.go 0 → 100644
View file @ d3b4b997
//go:build windows || darwin
package store
import (
"path/filepath"
"testing"
)
func TestStore(t *testing.T) {
s, cleanup := setupTestStore(t)
defer cleanup()
t.Run("default id", func(t *testing.T) {
// ID should be automatically generated
id, err := s.ID()
if err != nil {
t.Fatal(err)
}
if id == "" {
t.Error("expected non-empty ID")
}
// Verify ID is persisted
id2, err := s.ID()
if err != nil {
t.Fatal(err)
}
if id != id2 {
t.Errorf("expected ID %s, got %s", id, id2)
}
})
t.Run("has completed first run", func(t *testing.T) {
// Default should be false (hasn't completed first run yet)
hasCompleted, err := s.HasCompletedFirstRun()
if err != nil {
t.Fatal(err)
}
if hasCompleted {
t.Error("expected has completed first run to be false by default")
}
if err := s.SetHasCompletedFirstRun(true); err != nil {
t.Fatal(err)
}
hasCompleted, err = s.HasCompletedFirstRun()
if err != nil {
t.Fatal(err)
}
if !hasCompleted {
t.Error("expected has completed first run to be true")
}
})
t.Run("settings", func(t *testing.T) {
sc := Settings{
Expose: true,
Browser: true,
Survey: true,
Models: "/tmp/models",
Agent: true,
Tools: false,
WorkingDir: "/tmp/work",
}
if err := s.SetSettings(sc); err != nil {
t.Fatal(err)
}
loaded, err := s.Settings()
if err != nil {
t.Fatal(err)
}
// Compare fields individually since Models might get a default
if loaded.Expose != sc.Expose || loaded.Browser != sc.Browser ||
loaded.Agent != sc.Agent || loaded.Survey != sc.Survey ||
loaded.Tools != sc.Tools || loaded.WorkingDir != sc.WorkingDir {
t.Errorf("expected %v, got %v", sc, loaded)
}
})
t.Run("window size", func(t *testing.T) {
if err := s.SetWindowSize(1024, 768); err != nil {
t.Fatal(err)
}
width, height, err := s.WindowSize()
if err != nil {
t.Fatal(err)
}
if width != 1024 || height != 768 {
t.Errorf("expected 1024x768, got %dx%d", width, height)
}
})
t.Run("create and retrieve chat", func(t *testing.T) {
chat := NewChat("test-chat-1")
chat.Title = "Test Chat"
chat.Messages = append(chat.Messages, NewMessage("user", "Hello", nil))
chat.Messages = append(chat.Messages, NewMessage("assistant", "Hi there!", &MessageOptions{
Model: "llama4",
}))
if err := s.SetChat(*chat); err != nil {
t.Fatalf("failed to save chat: %v", err)
}
retrieved, err := s.Chat("test-chat-1")
if err != nil {
t.Fatalf("failed to retrieve chat: %v", err)
}
if retrieved.ID != chat.ID {
t.Errorf("expected ID %s, got %s", chat.ID, retrieved.ID)
}
if retrieved.Title != chat.Title {
t.Errorf("expected title %s, got %s", chat.Title, retrieved.Title)
}
if len(retrieved.Messages) != 2 {
t.Fatalf("expected 2 messages, got %d", len(retrieved.Messages))
}
if retrieved.Messages[0].Content != "Hello" {
t.Errorf("expected first message 'Hello', got %s", retrieved.Messages[0].Content)
}
if retrieved.Messages[1].Content != "Hi there!" {
t.Errorf("expected second message 'Hi there!', got %s", retrieved.Messages[1].Content)
}
})
t.Run("list chats", func(t *testing.T) {
chat2 := NewChat("test-chat-2")
chat2.Title = "Another Chat"
chat2.Messages = append(chat2.Messages, NewMessage("user", "Test", nil))
if err := s.SetChat(*chat2); err != nil {
t.Fatalf("failed to save chat: %v", err)
}
chats, err := s.Chats()
if err != nil {
t.Fatalf("failed to list chats: %v", err)
}
if len(chats) != 2 {
t.Fatalf("expected 2 chats, got %d", len(chats))
}
})
t.Run("delete chat", func(t *testing.T) {
if err := s.DeleteChat("test-chat-1"); err != nil {
t.Fatalf("failed to delete chat: %v", err)
}
// Verify it's gone
_, err := s.Chat("test-chat-1")
if err == nil {
t.Error("expected error retrieving deleted chat")
}
// Verify other chat still exists
chats, err := s.Chats()
if err != nil {
t.Fatalf("failed to list chats: %v", err)
}
if len(chats) != 1 {
t.Fatalf("expected 1 chat after deletion, got %d", len(chats))
}
})
}
// setupTestStore creates a temporary store for testing
func setupTestStore(t *testing.T) (*Store, func()) {
t.Helper()
tmpDir := t.TempDir()
// Override legacy config path to ensure no migration happens
oldLegacyConfigPath := legacyConfigPath
legacyConfigPath = filepath.Join(tmpDir, "config.json")
s := &Store{DBPath: filepath.Join(tmpDir, "db.sqlite")}
cleanup := func() {
s.Close()
legacyConfigPath = oldLegacyConfigPath
}
return s, cleanup
}
app/store/store_windows.go deleted 100644 → 0
View file @ a4770107
package store
import (
"os"
"path/filepath"
)
func getStorePath() string {
localAppData := os.Getenv("LOCALAPPDATA")
return filepath.Join(localAppData, "Ollama", "config.json")
}
app/store/testdata/schema.sql 0 → 100644
View file @ d3b4b997
-- This is the version 2 schema for the app database, the first released schema to users.
-- Do not modify this file. It is used to test that the database schema stays in a consistent state between schema migrations.
CREATE TABLE IF NOT EXISTS settings (
id INTEGER PRIMARY KEY CHECK (id = 1),
device_id TEXT NOT NULL DEFAULT '',
has_completed_first_run BOOLEAN NOT NULL DEFAULT 0,
expose BOOLEAN NOT NULL DEFAULT 0,
survey BOOLEAN NOT NULL DEFAULT TRUE,
browser BOOLEAN NOT NULL DEFAULT 0,
models TEXT NOT NULL DEFAULT '',
remote TEXT NOT NULL DEFAULT '',
agent BOOLEAN NOT NULL DEFAULT 0,
tools BOOLEAN NOT NULL DEFAULT 0,
working_dir TEXT NOT NULL DEFAULT '',
context_length INTEGER NOT NULL DEFAULT 4096,
window_width INTEGER NOT NULL DEFAULT 0,
window_height INTEGER NOT NULL DEFAULT 0,
config_migrated BOOLEAN NOT NULL DEFAULT 0,
schema_version INTEGER NOT NULL DEFAULT 2
);
-- Insert default settings row if it doesn't exist
INSERT OR IGNORE INTO settings (id) VALUES (1);
CREATE TABLE IF NOT EXISTS chats (
id TEXT PRIMARY KEY,
title TEXT NOT NULL DEFAULT '',
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
);
CREATE TABLE IF NOT EXISTS 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,
FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS idx_messages_chat_id ON messages(chat_id);
CREATE TABLE IF NOT EXISTS tool_calls (
id INTEGER PRIMARY KEY AUTOINCREMENT,
message_id INTEGER NOT NULL,
type TEXT NOT NULL,
function_name TEXT NOT NULL,
function_arguments TEXT NOT NULL,
function_result TEXT,
FOREIGN KEY (message_id) REFERENCES messages(id) ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS idx_tool_calls_message_id ON tool_calls(message_id);
app/tools/browser.go 0 → 100644
View file @ d3b4b997
//go:build windows || darwin
package tools
import (
"context"
"fmt"
"net/url"
"regexp"
"strings"
"sync"
"time"
"github.com/ollama/ollama/app/ui/responses"
)
type PageType string
const (
PageTypeSearchResults PageType = "initial_results"
PageTypeWebpage PageType = "webpage"
)
// DefaultViewTokens is the number of tokens to show to the model used when calling displayPage
const DefaultViewTokens = 1024
/*
The Browser tool provides web browsing capability for gpt-oss.
The model uses the tool by usually doing a search first and then choosing to either open a page,
find a term in a page, or do another search.
The tool optionally may open a URL directly - especially if one is passed in.
Each action is saved into an append-only page stack `responses.BrowserStateData` to keep
track of the history of the browsing session.
Each `Execute()` for a tool returns the full current state of the browser. ui.go manages the
browser state representation between the tool, ui, and db.
A new Browser object is created per request - the state is reconstructed by ui.go.
The initialization of the browser will receive a `responses.BrowserStateData` with the stitched history.
*/
// BrowserState manages the browsing session on a per-chat basis
type BrowserState struct {
mu sync.RWMutex
Data *responses.BrowserStateData
}
type Browser struct {
state *BrowserState
}
// State is only accessed in a single thread, as each chat has its own browser state
func (b *Browser) State() *responses.BrowserStateData {
b.state.mu.RLock()
defer b.state.mu.RUnlock()
return b.state.Data
}
func (b *Browser) savePage(page *responses.Page) {
b.state.Data.URLToPage[page.URL] = page
b.state.Data.PageStack = append(b.state.Data.PageStack, page.URL)
}
func (b *Browser) getPageFromStack(url string) (*responses.Page, error) {
page, ok := b.state.Data.URLToPage[url]
if !ok {
return nil, fmt.Errorf("page not found for url %s", url)
}
return page, nil
}
func NewBrowser(state *responses.BrowserStateData) *Browser {
if state == nil {
state = &responses.BrowserStateData{
PageStack: []string{},
ViewTokens: DefaultViewTokens,
URLToPage: make(map[string]*responses.Page),
}
}
b := &BrowserState{
Data: state,
}
return &Browser{
state: b,
}
}
type BrowserSearch struct {
Browser
webSearch *BrowserWebSearch
}
// NewBrowserSearch creates a new browser search instance
func NewBrowserSearch(bb *Browser) *BrowserSearch {
if bb == nil {
bb = &Browser{
state: &BrowserState{
Data: &responses.BrowserStateData{
PageStack: []string{},
ViewTokens: DefaultViewTokens,
URLToPage: make(map[string]*responses.Page),
},
},
}
}
return &BrowserSearch{
Browser: *bb,
webSearch: &BrowserWebSearch{},
}
}
func (b *BrowserSearch) Name() string {
return "browser.search"
}
func (b *BrowserSearch) Description() string {
return "Search the web for information"
}
func (b *BrowserSearch) Prompt() string {
return ""
}
func (b *BrowserSearch) Schema() map[string]any {
return map[string]any{}
}
func (b *BrowserSearch) Execute(ctx context.Context, args map[string]any) (any, string, error) {
query, ok := args["query"].(string)
if !ok {
return nil, "", fmt.Errorf("query parameter is required")
}
topn, ok := args["topn"].(int)
if !ok {
topn = 5
}
searchArgs := map[string]any{
"queries": []any{query},
"max_results": topn,
}
result, err := b.webSearch.Execute(ctx, searchArgs)
if err != nil {
return nil, "", fmt.Errorf("search error: %w", err)
}
searchResponse, ok := result.(*WebSearchResponse)
if !ok {
return nil, "", fmt.Errorf("invalid search results format")
}
// Build main search results page that contains all search results
searchResultsPage := b.buildSearchResultsPageCollection(query, searchResponse)
b.savePage(searchResultsPage)
cursor := len(b.state.Data.PageStack) - 1
// cache result for each page
for _, queryResults := range searchResponse.Results {
for i, result := range queryResults {
resultPage := b.buildSearchResultsPage(&result, i+1)
// save to global only, do not add to visited stack
b.state.Data.URLToPage[resultPage.URL] = resultPage
}
}
page := searchResultsPage
pageText, err := b.displayPage(page, cursor, 0, -1)
if err != nil {
return nil, "", fmt.Errorf("failed to display page: %w", err)
}
return b.state.Data, pageText, nil
}
func (b *Browser) buildSearchResultsPageCollection(query string, results *WebSearchResponse) *responses.Page {
page := &responses.Page{
URL: "search_results_" + query,
Title: query,
Links: make(map[int]string),
FetchedAt: time.Now(),
}
var textBuilder strings.Builder
linkIdx := 0
// Add the header lines to match format
textBuilder.WriteString("\n") // L0: empty
textBuilder.WriteString("URL: \n") // L1: URL: (empty for search)
textBuilder.WriteString("# Search Results\n") // L2: # Search Results
textBuilder.WriteString("\n") // L3: empty
for _, queryResults := range results.Results {
for _, result := range queryResults {
domain := result.URL
if u, err := url.Parse(result.URL); err == nil && u.Host != "" {
domain = u.Host
domain = strings.TrimPrefix(domain, "www.")
}
linkFormat := fmt.Sprintf("* 【%d†%s†%s】", linkIdx, result.Title, domain)
textBuilder.WriteString(linkFormat)
numChars := min(len(result.Content.FullText), 400)
snippet := strings.TrimSpace(result.Content.FullText[:numChars])
textBuilder.WriteString(snippet)
textBuilder.WriteString("\n")
page.Links[linkIdx] = result.URL
linkIdx++
}
}
page.Text = textBuilder.String()
page.Lines = wrapLines(page.Text, 80)
return page
}
func (b *Browser) buildSearchResultsPage(result *WebSearchResult, linkIdx int) *responses.Page {
page := &responses.Page{
URL: result.URL,
Title: result.Title,
Links: make(map[int]string),
FetchedAt: time.Now(),
}
var textBuilder strings.Builder
// Format the individual result page (only used when no full text is available)
linkFormat := fmt.Sprintf("【%d†%s】", linkIdx, result.Title)
textBuilder.WriteString(linkFormat)
textBuilder.WriteString("\n")
textBuilder.WriteString(fmt.Sprintf("URL: %s\n", result.URL))
numChars := min(len(result.Content.FullText), 300)
textBuilder.WriteString(result.Content.FullText[:numChars])
textBuilder.WriteString("\n\n")
// Only store link and snippet if we won't be processing full text later
// (full text processing will handle all links consistently)
if result.Content.FullText == "" {
page.Links[linkIdx] = result.URL
}
// Use full text if available, otherwise use snippet
if result.Content.FullText != "" {
// Prepend the URL line to the full text
page.Text = fmt.Sprintf("URL: %s\n%s", result.URL, result.Content.FullText)
// Process markdown links in the full text
processedText, processedLinks := processMarkdownLinks(page.Text)
page.Text = processedText
page.Links = processedLinks
} else {
page.Text = textBuilder.String()
}
page.Lines = wrapLines(page.Text, 80)
return page
}
// getEndLoc calculates the end location for viewport based on token limits
func (b *Browser) getEndLoc(loc, numLines, totalLines int, lines []string) int {
if numLines <= 0 {
// Auto-calculate based on viewTokens
txt := b.joinLinesWithNumbers(lines[loc:])
// If text is very short, no need to truncate (at least 1 char per token)
if len(txt) > b.state.Data.ViewTokens {
// Simple heuristic: approximate token counting
// Typical token is ~4 characters, but can be up to 128 chars
maxCharsPerToken := 128
// upper bound for text to analyze
upperBound := min((b.state.Data.ViewTokens+1)*maxCharsPerToken, len(txt))
textToAnalyze := txt[:upperBound]
// Simple approximation: count tokens as ~4 chars each
// This is less accurate than tiktoken but more performant
approxTokens := len(textToAnalyze) / 4
if approxTokens > b.state.Data.ViewTokens {
// Find the character position at viewTokens
endIdx := min(b.state.Data.ViewTokens*4, len(txt))
// Count newlines up to that position to get line count
numLines = strings.Count(txt[:endIdx], "\n") + 1
} else {
numLines = totalLines
}
} else {
numLines = totalLines
}
}
return min(loc+numLines, totalLines)
}
// joinLinesWithNumbers creates a string with line numbers, matching Python's join_lines
func (b *Browser) joinLinesWithNumbers(lines []string) string {
var builder strings.Builder
var hadZeroLine bool
for i, line := range lines {
if i == 0 {
builder.WriteString("L0:\n")
hadZeroLine = true
}
if hadZeroLine {
builder.WriteString(fmt.Sprintf("L%d: %s\n", i+1, line))
} else {
builder.WriteString(fmt.Sprintf("L%d: %s\n", i, line))
}
}
return builder.String()
}
// processMarkdownLinks finds all markdown links in the text and replaces them with the special format
// Returns the processed text and a map of link IDs to URLs
func processMarkdownLinks(text string) (string, map[int]string) {
links := make(map[int]string)
// Always start from 0 for consistent numbering across all pages
linkID := 0
// First, handle multi-line markdown links by joining them
// This regex finds markdown links that might be split across lines
multiLinePattern := regexp.MustCompile(`\[([^\]]+)\]\s*\n\s*\(([^)]+)\)`)
text = multiLinePattern.ReplaceAllStringFunc(text, func(match string) string {
// Replace newlines with spaces in the match
cleaned := strings.ReplaceAll(match, "\n", " ")
// Remove extra spaces
cleaned = regexp.MustCompile(`\s+`).ReplaceAllString(cleaned, " ")
return cleaned
})
// Now process all markdown links (including the cleaned multi-line ones)
linkPattern := regexp.MustCompile(`\[([^\]]+)\]\(([^)]+)\)`)
processedText := linkPattern.ReplaceAllStringFunc(text, func(match string) string {
matches := linkPattern.FindStringSubmatch(match)
if len(matches) != 3 {
return match
}
linkText := strings.TrimSpace(matches[1])
linkURL := strings.TrimSpace(matches[2])
// Extract domain from URL
domain := linkURL
if u, err := url.Parse(linkURL); err == nil && u.Host != "" {
domain = u.Host
// Remove www. prefix if present
domain = strings.TrimPrefix(domain, "www.")
}
// Create the formatted link
formatted := fmt.Sprintf("【%d†%s†%s】", linkID, linkText, domain)
// Store the link
links[linkID] = linkURL
linkID++
return formatted
})
return processedText, links
}
func wrapLines(text string, width int) []string {
if width <= 0 {
width = 80
}
lines := strings.Split(text, "\n")
var wrapped []string
for _, line := range lines {
if line == "" {
// Preserve empty lines
wrapped = append(wrapped, "")
} else if len(line) <= width {
wrapped = append(wrapped, line)
} else {
// Word wrapping while preserving whitespace structure
words := strings.Fields(line)
if len(words) == 0 {
// Line with only whitespace
wrapped = append(wrapped, line)
continue
}
currentLine := ""
for _, word := range words {
// Check if adding this word would exceed width
testLine := currentLine
if testLine != "" {
testLine += " "
}
testLine += word
if len(testLine) > width && currentLine != "" {
// Current line would be too long, wrap it
wrapped = append(wrapped, currentLine)
currentLine = word
} else {
// Add word to current line
if currentLine != "" {
currentLine += " "
}
currentLine += word
}
}
// Add any remaining content
if currentLine != "" {
wrapped = append(wrapped, currentLine)
}
}
}
return wrapped
}
// displayPage formats and returns the page display for the model
func (b *Browser) displayPage(page *responses.Page, cursor, loc, numLines int) (string, error) {
totalLines := len(page.Lines)
if loc >= totalLines {
return "", fmt.Errorf("invalid location: %d (max: %d)", loc, totalLines-1)
}
// get viewport end location
endLoc := b.getEndLoc(loc, numLines, totalLines, page.Lines)
var displayBuilder strings.Builder
displayBuilder.WriteString(fmt.Sprintf("[%d] %s", cursor, page.Title))
if page.URL != "" {
displayBuilder.WriteString(fmt.Sprintf("(%s)\n", page.URL))
} else {
displayBuilder.WriteString("\n")
}
displayBuilder.WriteString(fmt.Sprintf("**viewing lines [%d - %d] of %d**\n\n", loc, endLoc-1, totalLines-1))
// Content with line numbers
var hadZeroLine bool
for i := loc; i < endLoc; i++ {
if i == 0 {
displayBuilder.WriteString("L0:\n")
hadZeroLine = true
}
if hadZeroLine {
displayBuilder.WriteString(fmt.Sprintf("L%d: %s\n", i+1, page.Lines[i]))
} else {
displayBuilder.WriteString(fmt.Sprintf("L%d: %s\n", i, page.Lines[i]))
}
}
return displayBuilder.String(), nil
}
type BrowserOpen struct {
Browser
crawlPage *BrowserCrawler
}
func NewBrowserOpen(bb *Browser) *BrowserOpen {
if bb == nil {
bb = &Browser{
state: &BrowserState{
Data: &responses.BrowserStateData{
PageStack: []string{},
ViewTokens: DefaultViewTokens,
URLToPage: make(map[string]*responses.Page),
},
},
}
}
return &BrowserOpen{
Browser: *bb,
crawlPage: &BrowserCrawler{},
}
}
func (b *BrowserOpen) Name() string {
return "browser.open"
}
func (b *BrowserOpen) Description() string {
return "Open a link in the browser"
}
func (b *BrowserOpen) Prompt() string {
return ""
}
func (b *BrowserOpen) Schema() map[string]any {
return map[string]any{}
}
func (b *BrowserOpen) Execute(ctx context.Context, args map[string]any) (any, string, error) {
// Get cursor parameter first
cursor := -1
if c, ok := args["cursor"].(float64); ok {
cursor = int(c)
} else if c, ok := args["cursor"].(int); ok {
cursor = c
}
// Get loc parameter
loc := 0
if l, ok := args["loc"].(float64); ok {
loc = int(l)
} else if l, ok := args["loc"].(int); ok {
loc = l
}
// Get num_lines parameter
numLines := -1
if n, ok := args["num_lines"].(float64); ok {
numLines = int(n)
} else if n, ok := args["num_lines"].(int); ok {
numLines = n
}
// get page from cursor
var page *responses.Page
if cursor >= 0 {
if cursor >= len(b.state.Data.PageStack) {
return nil, "", fmt.Errorf("cursor %d is out of range (pageStack length: %d)", cursor, len(b.state.Data.PageStack))
}
var err error
page, err = b.getPageFromStack(b.state.Data.PageStack[cursor])
if err != nil {
return nil, "", fmt.Errorf("page not found for cursor %d: %w", cursor, err)
}
} else {
// get last page
if len(b.state.Data.PageStack) != 0 {
pageURL := b.state.Data.PageStack[len(b.state.Data.PageStack)-1]
var err error
page, err = b.getPageFromStack(pageURL)
if err != nil {
return nil, "", fmt.Errorf("page not found for cursor %d: %w", cursor, err)
}
}
}
// Try to get id as string (URL) first
if url, ok := args["id"].(string); ok {
// Check if we already have this page cached
if existingPage, ok := b.state.Data.URLToPage[url]; ok {
// Use cached page
b.savePage(existingPage)
// Always update cursor to point to the newly added page
cursor = len(b.state.Data.PageStack) - 1
pageText, err := b.displayPage(existingPage, cursor, loc, numLines)
if err != nil {
return nil, "", fmt.Errorf("failed to display page: %w", err)
}
return b.state.Data, pageText, nil
}
// Page not in cache, need to crawl it
if b.crawlPage == nil {
b.crawlPage = &BrowserCrawler{}
}
crawlResponse, err := b.crawlPage.Execute(ctx, map[string]any{
"urls": []any{url},
"latest": false,
})
if err != nil {
return nil, "", fmt.Errorf("failed to crawl URL %s: %w", url, err)
}
newPage, err := b.buildPageFromCrawlResult(url, crawlResponse)
if err != nil {
return nil, "", fmt.Errorf("failed to build page from crawl result: %w", err)
}
// Need to fall through if first search is directly an open command - no existing page
b.savePage(newPage)
// Always update cursor to point to the newly added page
cursor = len(b.state.Data.PageStack) - 1
pageText, err := b.displayPage(newPage, cursor, loc, numLines)
if err != nil {
return nil, "", fmt.Errorf("failed to display page: %w", err)
}
return b.state.Data, pageText, nil
}
// Try to get id as integer (link ID from current page)
if id, ok := args["id"].(float64); ok {
if page == nil {
return nil, "", fmt.Errorf("no current page to resolve link from")
}
idInt := int(id)
pageURL, ok := page.Links[idInt]
if !ok {
return nil, "", fmt.Errorf("invalid link id %d", idInt)
}
// Check if we have the linked page cached
newPage, ok := b.state.Data.URLToPage[pageURL]
if !ok {
if b.crawlPage == nil {
b.crawlPage = &BrowserCrawler{}
}
crawlResponse, err := b.crawlPage.Execute(ctx, map[string]any{
"urls": []any{pageURL},
"latest": false,
})
if err != nil {
return nil, "", fmt.Errorf("failed to crawl URL %s: %w", pageURL, err)
}
// Create new page from crawl result
newPage, err = b.buildPageFromCrawlResult(pageURL, crawlResponse)
if err != nil {
return nil, "", fmt.Errorf("failed to build page from crawl result: %w", err)
}
}
// Add to history stack regardless of cache status
b.savePage(newPage)
// Always update cursor to point to the newly added page
cursor = len(b.state.Data.PageStack) - 1
pageText, err := b.displayPage(newPage, cursor, loc, numLines)
if err != nil {
return nil, "", fmt.Errorf("failed to display page: %w", err)
}
return b.state.Data, pageText, nil
}
// If no id provided, just display current page
if page == nil {
return nil, "", fmt.Errorf("no current page to display")
}
// Only add to PageStack without updating URLToPage
b.state.Data.PageStack = append(b.state.Data.PageStack, page.URL)
cursor = len(b.state.Data.PageStack) - 1
pageText, err := b.displayPage(page, cursor, loc, numLines)
if err != nil {
return nil, "", fmt.Errorf("failed to display page: %w", err)
}
return b.state.Data, pageText, nil
}
// buildPageFromCrawlResult creates a Page from crawl API results
func (b *Browser) buildPageFromCrawlResult(requestedURL string, crawlResponse *CrawlResponse) (*responses.Page, error) {
// Initialize page with defaults
page := &responses.Page{
URL: requestedURL,
Title: requestedURL,
Text: "",
Links: make(map[int]string),
FetchedAt: time.Now(),
}
// Process crawl results - the API returns results grouped by URL
for url, urlResults := range crawlResponse.Results {
if len(urlResults) > 0 {
// Get the first result for this URL
result := urlResults[0]
// Extract content
if result.Content.FullText != "" {
page.Text = result.Content.FullText
}
// Extract title if available
if result.Title != "" {
page.Title = result.Title
}
// Update URL to the actual URL from results
page.URL = url
// Extract links if available from extras
for i, link := range result.Extras.Links {
if link.Href != "" {
page.Links[i] = link.Href
} else if link.URL != "" {
page.Links[i] = link.URL
}
}
// Only process the first URL's results
break
}
}
// If no text was extracted, set a default message
if page.Text == "" {
page.Text = "No content could be extracted from this page."
} else {
// Prepend the URL line to match Python implementation
page.Text = fmt.Sprintf("URL: %s\n%s", page.URL, page.Text)
}
// Process markdown links in the text
processedText, processedLinks := processMarkdownLinks(page.Text)
page.Text = processedText
page.Links = processedLinks
// Wrap lines for display
page.Lines = wrapLines(page.Text, 80)
return page, nil
}
type BrowserFind struct {
Browser
}
func NewBrowserFind(bb *Browser) *BrowserFind {
return &BrowserFind{
Browser: *bb,
}
}
func (b *BrowserFind) Name() string {
return "browser.find"
}
func (b *BrowserFind) Description() string {
return "Find a term in the browser"
}
func (b *BrowserFind) Prompt() string {
return ""
}
func (b *BrowserFind) Schema() map[string]any {
return map[string]any{}
}
func (b *BrowserFind) Execute(ctx context.Context, args map[string]any) (any, string, error) {
pattern, ok := args["pattern"].(string)
if !ok {
return nil, "", fmt.Errorf("pattern parameter is required")
}
// Get cursor parameter if provided, default to current page
cursor := -1
if c, ok := args["cursor"].(float64); ok {
cursor = int(c)
}
// Get the page to search in
var page *responses.Page
if cursor == -1 {
// Use current page
if len(b.state.Data.PageStack) == 0 {
return nil, "", fmt.Errorf("no pages to search in")
}
var err error
page, err = b.getPageFromStack(b.state.Data.PageStack[len(b.state.Data.PageStack)-1])
if err != nil {
return nil, "", fmt.Errorf("page not found for cursor %d: %w", cursor, err)
}
} else {
// Use specific cursor
if cursor < 0 || cursor >= len(b.state.Data.PageStack) {
return nil, "", fmt.Errorf("cursor %d is out of range [0-%d]", cursor, len(b.state.Data.PageStack)-1)
}
var err error
page, err = b.getPageFromStack(b.state.Data.PageStack[cursor])
if err != nil {
return nil, "", fmt.Errorf("page not found for cursor %d: %w", cursor, err)
}
}
if page == nil {
return nil, "", fmt.Errorf("page not found")
}
// Create find results page
findPage := b.buildFindResultsPage(pattern, page)
// Add the find results page to state
b.savePage(findPage)
newCursor := len(b.state.Data.PageStack) - 1
pageText, err := b.displayPage(findPage, newCursor, 0, -1)
if err != nil {
return nil, "", fmt.Errorf("failed to display page: %w", err)
}
return b.state.Data, pageText, nil
}
func (b *Browser) buildFindResultsPage(pattern string, page *responses.Page) *responses.Page {
findPage := &responses.Page{
Title: fmt.Sprintf("Find results for text: `%s` in `%s`", pattern, page.Title),
Links: make(map[int]string),
FetchedAt: time.Now(),
}
findPage.URL = fmt.Sprintf("find_results_%s", pattern)
var textBuilder strings.Builder
matchIdx := 0
maxResults := 50
numShowLines := 4
patternLower := strings.ToLower(pattern)
// Search through the page lines following the reference algorithm
var resultChunks []string
lineIdx := 0
for lineIdx < len(page.Lines) {
line := page.Lines[lineIdx]
lineLower := strings.ToLower(line)
if !strings.Contains(lineLower, patternLower) {
lineIdx++
continue
}
// Build snippet context
endLine := min(lineIdx+numShowLines, len(page.Lines))
var snippetBuilder strings.Builder
for j := lineIdx; j < endLine; j++ {
snippetBuilder.WriteString(page.Lines[j])
if j < endLine-1 {
snippetBuilder.WriteString("\n")
}
}
snippet := snippetBuilder.String()
// Format the match
linkFormat := fmt.Sprintf("【%d†match at L%d】", matchIdx, lineIdx)
resultChunk := fmt.Sprintf("%s\n%s", linkFormat, snippet)
resultChunks = append(resultChunks, resultChunk)
if len(resultChunks) >= maxResults {
break
}
matchIdx++
lineIdx += numShowLines
}
// Build final display text
if len(resultChunks) > 0 {
textBuilder.WriteString(strings.Join(resultChunks, "\n\n"))
}
if matchIdx == 0 {
findPage.Text = fmt.Sprintf("No `find` results for pattern: `%s`", pattern)
} else {
findPage.Text = textBuilder.String()
}
findPage.Lines = wrapLines(findPage.Text, 80)
return findPage
}
app/tools/browser_crawl.go 0 → 100644
View file @ d3b4b997
//go:build windows || darwin
package tools
import (
"context"
"encoding/json"
"fmt"
)
// CrawlContent represents the content of a crawled page
type CrawlContent struct {
Snippet string `json:"snippet"`
FullText string `json:"full_text"`
}
// CrawlExtras represents additional data from the crawl API
type CrawlExtras struct {
Links []CrawlLink `json:"links"`
}
// CrawlLink represents a link found on a crawled page
type CrawlLink struct {
URL string `json:"url"`
Href string `json:"href"`
Text string `json:"text"`
}
// CrawlResult represents a single crawl result
type CrawlResult struct {
Title string `json:"title"`
URL string `json:"url"`
Content CrawlContent `json:"content"`
Extras CrawlExtras `json:"extras"`
}
// CrawlResponse represents the complete response from the crawl API
type CrawlResponse struct {
Results map[string][]CrawlResult `json:"results"`
}
// BrowserCrawler tool for crawling web pages using ollama.com crawl API
type BrowserCrawler struct{}
func (g *BrowserCrawler) Name() string {
return "get_webpage"
}
func (g *BrowserCrawler) Description() string {
return "Crawl and extract text content from web pages"
}
func (g *BrowserCrawler) Prompt() string {
return `When you need to read content from web pages, use the get_webpage tool. Simply provide the URLs you want to read and I'll fetch their content for you.
For each URL, I'll extract the main text content in a readable format. If you need to discover links within those pages, set extract_links to true. If the user requires the latest information, set livecrawl to true.
Only use this tool when you need to access current web content. Make sure the URLs are valid and accessible. Do not use this tool for:
- Downloading files or media
- Accessing private/authenticated pages
- Scraping data at high volumes
Always check the returned content to ensure it's relevant before using it in your response.`
}
func (g *BrowserCrawler) Schema() map[string]any {
schemaBytes := []byte(`{
"type": "object",
"properties": {
"urls": {
"type": "array",
"items": {
"type": "string"
},
"description": "List of URLs to crawl and extract content from"
}
},
"required": ["urls"]
}`)
var schema map[string]any
if err := json.Unmarshal(schemaBytes, &schema); err != nil {
return nil
}
return schema
}
func (g *BrowserCrawler) Execute(ctx context.Context, args map[string]any) (*CrawlResponse, error) {
urlsRaw, ok := args["urls"].([]any)
if !ok {
return nil, fmt.Errorf("urls parameter is required and must be an array of strings")
}
urls := make([]string, 0, len(urlsRaw))
for _, u := range urlsRaw {
if urlStr, ok := u.(string); ok {
urls = append(urls, urlStr)
}
}
if len(urls) == 0 {
return nil, fmt.Errorf("at least one URL is required")
}
return g.performWebCrawl(ctx, urls)
}
// performWebCrawl handles the actual HTTP request to ollama.com crawl API
func (g *BrowserCrawler) performWebCrawl(ctx context.Context, urls []string) (*CrawlResponse, error) {
result := &CrawlResponse{Results: make(map[string][]CrawlResult, len(urls))}
for _, targetURL := range urls {
fetchResp, err := performWebFetch(ctx, targetURL)
if err != nil {
return nil, fmt.Errorf("web_fetch failed for %q: %w", targetURL, err)
}
links := make([]CrawlLink, 0, len(fetchResp.Links))
for _, link := range fetchResp.Links {
links = append(links, CrawlLink{URL: link, Href: link})
}
snippet := truncateString(fetchResp.Content, 400)
result.Results[targetURL] = []CrawlResult{{
Title: fetchResp.Title,
URL: targetURL,
Content: CrawlContent{
Snippet: snippet,
FullText: fetchResp.Content,
},
Extras: CrawlExtras{Links: links},
}}
}
return result, nil
}
app/tools/browser_test.go 0 → 100644
View file @ d3b4b997
//go:build windows || darwin
package tools
import (
"strings"
"testing"
"time"
"github.com/ollama/ollama/app/ui/responses"
)
func makeTestPage(url string) *responses.Page {
return &responses.Page{
URL: url,
Title: "Title " + url,
Text: "Body for " + url,
Lines: []string{"line1", "line2", "line3"},
Links: map[int]string{0: url},
FetchedAt: time.Now(),
}
}
func TestBrowser_Scroll_AppendsOnlyPageStack(t *testing.T) {
b := NewBrowser(&responses.BrowserStateData{PageStack: []string{}, ViewTokens: 1024, URLToPage: map[string]*responses.Page{}})
p1 := makeTestPage("https://example.com/1")
b.savePage(p1)
initialStackLen := len(b.state.Data.PageStack)
initialMapLen := len(b.state.Data.URLToPage)
bo := NewBrowserOpen(b)
// Scroll without id — should push only to PageStack
_, _, err := bo.Execute(t.Context(), map[string]any{"loc": float64(1), "num_lines": float64(1)})
if err != nil {
t.Fatalf("scroll execute failed: %v", err)
}
if got, want := len(b.state.Data.PageStack), initialStackLen+1; got != want {
t.Fatalf("page stack length = %d, want %d", got, want)
}
if got, want := len(b.state.Data.URLToPage), initialMapLen; got != want {
t.Fatalf("url_to_page length changed = %d, want %d", got, want)
}
}
func TestBrowserOpen_UseCacheByURL(t *testing.T) {
b := NewBrowser(&responses.BrowserStateData{PageStack: []string{}, ViewTokens: 1024, URLToPage: map[string]*responses.Page{}})
bo := NewBrowserOpen(b)
p := makeTestPage("https://example.com/cached")
b.state.Data.URLToPage[p.URL] = p
initialStackLen := len(b.state.Data.PageStack)
initialMapLen := len(b.state.Data.URLToPage)
_, _, err := bo.Execute(t.Context(), map[string]any{"id": p.URL})
if err != nil {
t.Fatalf("open cached execute failed: %v", err)
}
if got, want := len(b.state.Data.PageStack), initialStackLen+1; got != want {
t.Fatalf("page stack length = %d, want %d", got, want)
}
if got, want := len(b.state.Data.URLToPage), initialMapLen; got != want {
t.Fatalf("url_to_page length changed = %d, want %d", got, want)
}
}
func TestDisplayPage_InvalidLoc(t *testing.T) {
b := NewBrowser(&responses.BrowserStateData{PageStack: []string{}, ViewTokens: 1024, URLToPage: map[string]*responses.Page{}})
p := makeTestPage("https://example.com/x")
// ensure lines are set
p.Lines = []string{"a", "b"}
_, err := b.displayPage(p, 0, 10, -1)
if err == nil || !strings.Contains(err.Error(), "invalid location") {
t.Fatalf("expected invalid location error, got %v", err)
}
}
func TestBrowserOpen_LinkId_UsesCacheAndAppends(t *testing.T) {
b := NewBrowser(&responses.BrowserStateData{PageStack: []string{}, ViewTokens: 1024, URLToPage: map[string]*responses.Page{}})
// Seed a main page with a link id 0 to a linked URL
main := makeTestPage("https://example.com/main")
linked := makeTestPage("https://example.com/linked")
main.Links = map[int]string{0: linked.URL}
// Save the main page (adds to PageStack and URLToPage)
b.savePage(main)
// Pre-cache the linked page so open by id avoids network
b.state.Data.URLToPage[linked.URL] = linked
initialStackLen := len(b.state.Data.PageStack)
initialMapLen := len(b.state.Data.URLToPage)
bo := NewBrowserOpen(b)
_, _, err := bo.Execute(t.Context(), map[string]any{"id": float64(0)})
if err != nil {
t.Fatalf("open by link id failed: %v", err)
}
if got, want := len(b.state.Data.PageStack), initialStackLen+1; got != want {
t.Fatalf("page stack length = %d, want %d", got, want)
}
if got, want := len(b.state.Data.URLToPage), initialMapLen; got != want {
t.Fatalf("url_to_page length changed = %d, want %d", got, want)
}
if last := b.state.Data.PageStack[len(b.state.Data.PageStack)-1]; last != linked.URL {
t.Fatalf("last page in stack = %s, want %s", last, linked.URL)
}
}
func TestWrapLines_PreserveAndWidth(t *testing.T) {
long := strings.Repeat("word ", 50)
text := "Line1\n\n" + long + "\nLine3"
lines := wrapLines(text, 40)
// Ensure empty line preserved at index 1
if lines[1] != "" {
t.Fatalf("expected preserved empty line at index 1, got %q", lines[1])
}
// All lines should be <= 40 chars
for i, l := range lines {
if len(l) > 40 {
t.Fatalf("line %d exceeds width: %d > 40", i, len(l))
}
}
}
func TestDisplayPage_FormatHeaderAndLines(t *testing.T) {
b := NewBrowser(&responses.BrowserStateData{PageStack: []string{}, ViewTokens: 1024, URLToPage: map[string]*responses.Page{}})
p := &responses.Page{
URL: "https://example.com/x",
Title: "Example",
Lines: []string{"URL: https://example.com/x", "A", "B", "C"},
}
out, err := b.displayPage(p, 3, 0, 2)
if err != nil {
t.Fatalf("displayPage failed: %v", err)
}
if !strings.HasPrefix(out, "[3] Example(") {
t.Fatalf("header not formatted as expected: %q", out)
}
if !strings.Contains(out, "L0:\n") {
t.Fatalf("missing L0 label: %q", out)
}
if !strings.Contains(out, "L1: URL: https://example.com/x\n") || !strings.Contains(out, "L2: A\n") {
t.Fatalf("missing expected line numbers/content: %q", out)
}
}
app/tools/browser_websearch.go 0 → 100644
View file @ d3b4b997
//go:build windows || darwin
package tools
import (
"context"
"encoding/json"
"fmt"
"strconv"
"time"
)
// WebSearchContent represents the content of a search result
type WebSearchContent struct {
Snippet string `json:"snippet"`
FullText string `json:"full_text"`
}
// WebSearchMetadata represents metadata for a search result
type WebSearchMetadata struct {
PublishedDate *time.Time `json:"published_date,omitempty"`
}
// WebSearchResult represents a single search result
type WebSearchResult struct {
Title string `json:"title"`
URL string `json:"url"`
Content WebSearchContent `json:"content"`
Metadata WebSearchMetadata `json:"metadata"`
}
// WebSearchResponse represents the complete response from the websearch API
type WebSearchResponse struct {
Results map[string][]WebSearchResult `json:"results"`
}
// BrowserWebSearch tool for searching the web using ollama.com search API
type BrowserWebSearch struct{}
func (w *BrowserWebSearch) Name() string {
return "gpt_oss_web_search"
}
func (w *BrowserWebSearch) Description() string {
return "Search the web for real-time information using ollama.com search API."
}
func (w *BrowserWebSearch) Prompt() string {
return `Use the gpt_oss_web_search tool to search the web.
1. Come up with a list of search queries to get comprehensive information (typically 2-3 related queries work well)
2. Use the gpt_oss_web_search tool with multiple queries to get results organized by query
3. Use the search results to provide current up to date, accurate information
Today's date is ` + time.Now().Format("January 2, 2006") + `
Add "` + time.Now().Format("January 2, 2006") + `" for news queries and ` + strconv.Itoa(time.Now().Year()+1) + ` for other queries that need current information.`
}
func (w *BrowserWebSearch) Schema() map[string]any {
schemaBytes := []byte(`{
"type": "object",
"properties": {
"queries": {
"type": "array",
"items": {
"type": "string"
},
"description": "List of search queries to look up"
},
"max_results": {
"type": "integer",
"description": "Maximum number of results to return per query (default: 2) up to 5",
"default": 2
}
},
"required": ["queries"]
}`)
var schema map[string]any
if err := json.Unmarshal(schemaBytes, &schema); err != nil {
return nil
}
return schema
}
func (w *BrowserWebSearch) Execute(ctx context.Context, args map[string]any) (any, error) {
queriesRaw, ok := args["queries"].([]any)
if !ok {
return nil, fmt.Errorf("queries parameter is required and must be an array of strings")
}
queries := make([]string, 0, len(queriesRaw))
for _, q := range queriesRaw {
if query, ok := q.(string); ok {
queries = append(queries, query)
}
}
if len(queries) == 0 {
return nil, fmt.Errorf("at least one query is required")
}
maxResults := 5
if mr, ok := args["max_results"].(int); ok {
maxResults = mr
}
return w.performWebSearch(ctx, queries, maxResults)
}
// performWebSearch handles the actual HTTP request to ollama.com search API
func (w *BrowserWebSearch) performWebSearch(ctx context.Context, queries []string, maxResults int) (*WebSearchResponse, error) {
response := &WebSearchResponse{Results: make(map[string][]WebSearchResult, len(queries))}
for _, query := range queries {
searchResp, err := performWebSearch(ctx, query, maxResults)
if err != nil {
return nil, fmt.Errorf("web_search failed for %q: %w", query, err)
}
converted := make([]WebSearchResult, 0, len(searchResp.Results))
for _, item := range searchResp.Results {
converted = append(converted, WebSearchResult{
Title: item.Title,
URL: item.URL,
Content: WebSearchContent{
Snippet: truncateString(item.Content, 400),
FullText: item.Content,
},
Metadata: WebSearchMetadata{},
})
}
response.Results[query] = converted
}
return response, nil
}
func truncateString(input string, limit int) string {
if limit <= 0 || len(input) <= limit {
return input
}
return input[:limit]
}
app/tools/tools.go 0 → 100644
View file @ d3b4b997
//go:build windows || darwin
package tools
import (
"context"
"encoding/json"
"fmt"
)
// Tool defines the interface that all tools must implement
type Tool interface {
// Name returns the unique identifier for the tool
Name() string
// Description returns a human-readable description of what the tool does
Description() string
// Schema returns the JSON schema for the tool's parameters
Schema() map[string]any
// Execute runs the tool with the given arguments and returns result to store in db, and a string result for the model
Execute(ctx context.Context, args map[string]any) (any, string, error)
// Prompt returns a prompt for the tool
Prompt() string
}
// Registry manages the available tools and their execution
type Registry struct {
tools map[string]Tool
workingDir string // Working directory for all tool operations
}
// NewRegistry creates a new tool registry with no tools
func NewRegistry() *Registry {
return &Registry{
tools: make(map[string]Tool),
}
}
// Register adds a tool to the registry
func (r *Registry) Register(tool Tool) {
r.tools[tool.Name()] = tool
}
// Get retrieves a tool by name
func (r *Registry) Get(name string) (Tool, bool) {
tool, exists := r.tools[name]
return tool, exists
}
// List returns all available tools
func (r *Registry) List() []Tool {
tools := make([]Tool, 0, len(r.tools))
for _, tool := range r.tools {
tools = append(tools, tool)
}
return tools
}
// SetWorkingDir sets the working directory for all tool operations
func (r *Registry) SetWorkingDir(dir string) {
r.workingDir = dir
}
// Execute runs a tool with the given name and arguments
func (r *Registry) Execute(ctx context.Context, name string, args map[string]any) (any, string, error) {
tool, ok := r.tools[name]
if !ok {
return nil, "", fmt.Errorf("unknown tool: %s", name)
}
result, text, err := tool.Execute(ctx, args)
if err != nil {
return nil, "", err
}
return result, text, nil
}
// ToolCall represents a request to execute a tool
type ToolCall struct {
ID string `json:"id"`
Type string `json:"type"`
Function ToolFunction `json:"function"`
}
// ToolFunction represents the function call details
type ToolFunction struct {
Name string `json:"name"`
Arguments json.RawMessage `json:"arguments"`
}
// ToolResult represents the result of a tool execution
type ToolResult struct {
ToolCallID string `json:"tool_call_id"`
Content any `json:"content"`
Error string `json:"error,omitempty"`
}
// ToolSchemas returns all tools as schema maps suitable for API calls
func (r *Registry) AvailableTools() []map[string]any {
schemas := make([]map[string]any, 0, len(r.tools))
for _, tool := range r.tools {
schema := map[string]any{
"name": tool.Name(),
"description": tool.Description(),
"schema": tool.Schema(),
}
schemas = append(schemas, schema)
}
return schemas
}
// ToolNames returns a list of all tool names
func (r *Registry) ToolNames() []string {
names := make([]string, 0, len(r.tools))
for name := range r.tools {
names = append(names, name)
}
return names
}
app/tools/web_fetch.go 0 → 100644
View file @ d3b4b997
//go:build windows || darwin
package tools
import (
"bytes"
"context"
"encoding/json"
"fmt"
"net/http"
"net/url"
"strconv"
"strings"
"time"
"github.com/ollama/ollama/auth"
)
type WebFetch struct{}
type FetchRequest struct {
URL string `json:"url"`
}
type FetchResponse struct {
Title string `json:"title"`
Content string `json:"content"`
Links []string `json:"links"`
}
func (w *WebFetch) Name() string {
return "web_fetch"
}
func (w *WebFetch) Description() string {
return "Crawl and extract text content from web pages"
}
func (g *WebFetch) Schema() map[string]any {
schemaBytes := []byte(`{
"type": "object",
"properties": {
"url": {
"type": "string",
"description": "URL to crawl and extract content from"
}
},
"required": ["url"]
}`)
var schema map[string]any
if err := json.Unmarshal(schemaBytes, &schema); err != nil {
return nil
}
return schema
}
func (w *WebFetch) Prompt() string {
return ""
}
func (w *WebFetch) Execute(ctx context.Context, args map[string]any) (any, string, error) {
urlRaw, ok := args["url"]
if !ok {
return nil, "", fmt.Errorf("url parameter is required")
}
urlStr, ok := urlRaw.(string)
if !ok || strings.TrimSpace(urlStr) == "" {
return nil, "", fmt.Errorf("url must be a non-empty string")
}
result, err := performWebFetch(ctx, urlStr)
if err != nil {
return nil, "", err
}
return result, "", nil
}
func performWebFetch(ctx context.Context, targetURL string) (*FetchResponse, error) {
reqBody := FetchRequest{URL: targetURL}
jsonBody, err := json.Marshal(reqBody)
if err != nil {
return nil, fmt.Errorf("failed to marshal request body: %w", err)
}
crawlURL, err := url.Parse("https://ollama.com/api/web_fetch")
if err != nil {
return nil, fmt.Errorf("failed to parse fetch URL: %w", err)
}
query := crawlURL.Query()
query.Add("ts", strconv.FormatInt(time.Now().Unix(), 10))
crawlURL.RawQuery = query.Encode()
data := fmt.Appendf(nil, "%s,%s", http.MethodPost, crawlURL.RequestURI())
signature, err := auth.Sign(ctx, data)
if err != nil {
return nil, fmt.Errorf("failed to sign request: %w", err)
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, crawlURL.String(), bytes.NewBuffer(jsonBody))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
if signature != "" {
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", signature))
}
client := &http.Client{Timeout: 30 * time.Second}
resp, err := client.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to execute fetch request: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("fetch API error (status %d)", resp.StatusCode)
}
var result FetchResponse
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
return nil, fmt.Errorf("failed to decode response: %w", err)
}
return &result, nil
}
app/tools/web_search.go 0 → 100644
View file @ d3b4b997
//go:build windows || darwin
package tools
import (
"bytes"
"context"
"encoding/json"
"fmt"
"net/http"
"net/url"
"strconv"
"strings"
"time"
"github.com/ollama/ollama/auth"
)
type WebSearch struct{}
type SearchRequest struct {
Query string `json:"query"`
MaxResults int `json:"max_results,omitempty"`
}
type SearchResult struct {
Title string `json:"title"`
URL string `json:"url"`
Content string `json:"content"`
}
type SearchResponse struct {
Results []SearchResult `json:"results"`
}
func (w *WebSearch) Name() string {
return "web_search"
}
func (w *WebSearch) Description() string {
return "Search the web for real-time information using ollama.com web search API."
}
func (w *WebSearch) Prompt() string {
return ""
}
func (g *WebSearch) Schema() map[string]any {
schemaBytes := []byte(`{
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "The search query to execute"
},
"max_results": {
"type": "integer",
"description": "Maximum number of search results to return",
"default": 3
}
},
"required": ["query"]
}`)
var schema map[string]any
if err := json.Unmarshal(schemaBytes, &schema); err != nil {
return nil
}
return schema
}
func (w *WebSearch) Execute(ctx context.Context, args map[string]any) (any, string, error) {
rawQuery, ok := args["query"]
if !ok {
return nil, "", fmt.Errorf("query parameter is required")
}
queryStr, ok := rawQuery.(string)
if !ok || strings.TrimSpace(queryStr) == "" {
return nil, "", fmt.Errorf("query must be a non-empty string")
}
maxResults := 5
if v, ok := args["max_results"].(float64); ok && int(v) > 0 {
maxResults = int(v)
}
result, err := performWebSearch(ctx, queryStr, maxResults)
if err != nil {
return nil, "", err
}
return result, "", nil
}
func performWebSearch(ctx context.Context, query string, maxResults int) (*SearchResponse, error) {
reqBody := SearchRequest{Query: query, MaxResults: maxResults}
jsonBody, err := json.Marshal(reqBody)
if err != nil {
return nil, fmt.Errorf("failed to marshal request body: %w", err)
}
searchURL, err := url.Parse("https://ollama.com/api/web_search")
if err != nil {
return nil, fmt.Errorf("failed to parse search URL: %w", err)
}
q := searchURL.Query()
q.Add("ts", strconv.FormatInt(time.Now().Unix(), 10))
searchURL.RawQuery = q.Encode()
data := fmt.Appendf(nil, "%s,%s", http.MethodPost, searchURL.RequestURI())
signature, err := auth.Sign(ctx, data)
if err != nil {
return nil, fmt.Errorf("failed to sign request: %w", err)
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, searchURL.String(), bytes.NewBuffer(jsonBody))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
if signature != "" {
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", signature))
}
client := &http.Client{Timeout: 10 * time.Second}
resp, err := client.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to execute search request: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("search API error (status %d)", resp.StatusCode)
}
var result SearchResponse
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
return nil, fmt.Errorf("failed to decode response: %w", err)
}
return &result, nil
}
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment