From 9914c043cbb0306ca7e775e156750be243128711 Mon Sep 17 00:00:00 2001
From: Zoe <62722391+juls0730@users.noreply.github.com>
Date: Mon, 27 Apr 2026 12:05:47 -0500
Subject: [PATCH] feat: add web search tool with reranking support

Integrates SearXNG web search as a tool available during chat when the
agent has search enabled. Supports optional reranking of results via a
configurable reranking model. Replaces Python subprocess evaluation with
@pydantic/monty WASM runtime. Introduces stable tool call ID mapping to
avoid exposing provider-native IDs to the database.
---
 app/components/ChatInput.vue                  |  71 +++-
 app/components/SearchSelector.vue             | 140 +++++++
 app/composables/useAgents.ts                  |   8 +-
 app/composables/useChat.ts                    |  52 +--
 app/composables/useModels.ts                  |   1 +
 server/api/topic/[topicId]/chat/index.post.ts | 354 +++++++++++++-----
 6 files changed, 490 insertions(+), 136 deletions(-)
 create mode 100644 app/components/SearchSelector.vue

diff --git a/app/components/ChatInput.vue b/app/components/ChatInput.vue
index 5113d4b..f961f49 100644
--- a/app/components/ChatInput.vue
+++ b/app/components/ChatInput.vue
@@ -6,6 +6,7 @@ import type { Agent } from '~/composables/useAgents';
 import type FileSelector from './FileSelector.vue';
 
 const { allModels } = await useModels();
+const { updateAgent, patchAgentLocally } = await useAgents();
 
 const inputHeight: Ref<string> = ref('auto');
 const fileSelectorRef = ref<InstanceType<typeof FileSelector> | null>(null);
@@ -39,6 +40,35 @@ const props = defineProps<{
     providers?: ProviderWithModels[];
 }>();
 
+const searchConfig = ref({
+    enabled: props.agent?.config?.search?.enabled ?? false,
+    maxResults: props.agent?.config?.search?.maxResults ?? 10,
+    rerank: props.agent?.config?.search?.rerank ?? false,
+});
+
+watch(() => props.agent?.config?.search, (val) => {
+    searchConfig.value = {
+        enabled: val?.enabled ?? false,
+        maxResults: val?.maxResults ?? 10,
+        rerank: val?.rerank ?? false,
+    };
+}, { deep: true });
+
+const saveSearchConfig = async () => {
+    if (!props.agent) return;
+    const currentConfig = props.agent.config ?? {};
+    const newConfig = {
+        ...currentConfig,
+        search: {
+            enabled: searchConfig.value.enabled,
+            maxResults: searchConfig.value.maxResults,
+            rerank: searchConfig.value.rerank,
+        },
+    };
+    patchAgentLocally(props.agent.id, { config: newConfig });
+    updateAgent(props.agent.id, { config: newConfig });
+};
+
 const selectedModel = ref<ModelWithProvider | null>(null);
 
 const handlePaste = async (event: ClipboardEvent) => {
@@ -165,25 +195,21 @@ const handleWindowKeyDown = async (event: KeyboardEvent) => {
     }
 };
 
-watch(textAreaValue, async () => {
+const resizeTextArea = () => {
     const textarea = inputRef.value;
     if (!textarea) return;
 
     inputHeight.value = 'auto';
-    await nextTick();
+    nextTick().then(() => {
+        const lineHeight = 24;
+        const maxLines = 10;
+        const maxHeight = maxLines * lineHeight;
+        const height = Math.min(textarea.scrollHeight, maxHeight);
+        inputHeight.value = `${height}px`;
+    });
+};
 
-    const lineHeight = 24;
-    const maxLines = 10;
-    const maxHeight = maxLines * lineHeight;
-
-    const newHeight = textarea.scrollHeight;
-
-    if (newHeight > maxHeight) {
-        inputHeight.value = `${maxHeight}px`;
-    } else {
-        inputHeight.value = `${newHeight}px`;
-    }
-}, { immediate: true });
+watch(textAreaValue, resizeTextArea, { immediate: true });
 
 let hasCommandKey = false;
 if (import.meta.server) {
@@ -197,14 +223,26 @@ onBeforeMount(() => {
     tempInput = (document.getElementById('chat') as HTMLInputElement)?.value ?? '';
 });
 
+let resizeRafId: number | undefined;
+const handleWindowResize = () => {
+    if (resizeRafId) return;
+    resizeRafId = requestAnimationFrame(() => {
+        resizeRafId = undefined;
+        resizeTextArea();
+    });
+}
+
 onMounted(() => {
     textAreaValue.value = tempInput;
     document.addEventListener('keydown', handleWindowKeyDown);
+    window.addEventListener('resize', handleWindowResize);
     inputRef.value?.addEventListener('paste', handlePaste);
 });
 
 onUnmounted(() => {
     document.removeEventListener('keydown', handleWindowKeyDown);
+    window.removeEventListener('resize', handleWindowResize);
+    if (resizeRafId !== undefined) cancelAnimationFrame(resizeRafId);
     inputRef.value?.removeEventListener('paste', handlePaste);
 });
 </script>
@@ -232,6 +270,11 @@ onUnmounted(() => {
                 <div class="flex flex-1 gap-1 min-w-0">
                     <ModelSelector v-if="providers !== undefined" :add-hotkey="true" v-model="selectedModel"
                         :providers="providers" />
+                    <SearchSelector v-if="selectedModel?.capabilities.includes('tools')" :enabled="searchConfig.enabled"
+                        :max-results="searchConfig.maxResults" :rerank="searchConfig.rerank"
+                        @update:enabled="(v: boolean) => { searchConfig.enabled = v; saveSearchConfig() }"
+                        @update:max-results="(v: number) => { searchConfig.maxResults = v; saveSearchConfig() }"
+                        @update:rerank="(v: boolean) => { searchConfig.rerank = v; saveSearchConfig() }" />
                     <FileSelector ref="fileSelectorRef" :selected-model="selectedModel" v-model="files" />
                 </div>
                 <button aria-label="Send message" @click="handleSubmit"
diff --git a/app/components/SearchSelector.vue b/app/components/SearchSelector.vue
new file mode 100644
index 0000000..e4a5d03
--- /dev/null
+++ b/app/components/SearchSelector.vue
@@ -0,0 +1,140 @@
+<script setup lang="ts">
+import { DialogType } from '~/composables/useDialog';
+
+const { openDialog } = await useDialog();
+
+const props = defineProps<{
+    enabled: boolean;
+    maxResults: number;
+    rerank: boolean;
+}>();
+
+const emit = defineEmits<{
+    'update:enabled': [value: boolean];
+    'update:maxResults': [value: number];
+    'update:rerank': [value: boolean];
+}>();
+
+const adjustMaxResults = (delta: number) => {
+    const next = Math.min(50, Math.max(1, props.maxResults + delta));
+    emit('update:maxResults', next);
+};
+
+const { settings } = await useUserSettings();
+
+const systemAssistantsRerank = computed(() => {
+    return settings.value.systemAssistants?.rerank ?? null;
+});
+
+const isRerankConfigured = computed(() => {
+    const sa = systemAssistantsRerank.value;
+    return sa?.enabled === true && sa?.modelId != null && sa.modelId !== '';
+});
+
+const openSettings = () => {
+    openDialog(DialogType.Settings, undefined, { page: 'systemAssistants' });
+};
+</script>
+
+<template>
+    <Dropdown dropdownClass="text-sm" placement="top">
+        <template #default="{ toggle, setRef }">
+            <button :ref="setRef" @click="toggle"
+                class="flex items-center justify-center h-8.5 w-8.5 @hover:bg-[var(--color-hover)] rounded-lg transition-colors duration-200 ease-[cubic-bezier(0.5,_1,_0.89,_1)]">
+                <span class="pointer-events-none i-tabler-world text-5 transition-colors duration-200"
+                    :class="enabled ? 'text-[var(--color-accent)]' : 'text-[var(--text-secondary)]'"></span>
+            </button>
+        </template>
+
+        <template #dropdown="{ close }">
+            <div class="flex p-1 gap-4">
+                <!-- Left: mode toggle group -->
+                <div class="flex flex-col gap-1">
+                    <button @click="emit('update:enabled', false)"
+                        class="flex items-start gap-3 rounded-xl px-3 py-2.5 text-left transition-colors duration-150"
+                        :class="!enabled
+                            ? 'bg-[var(--color-active)]'
+                            : '@hover:bg-[var(--color-hover)]'">
+                        <span class="i-tabler-world-off text-5 mt-0.5 shrink-0"
+                            :class="!enabled ? 'text-[var(--text-primary)]' : 'text-[var(--text-dim)]'"></span>
+                        <div class="flex flex-col gap-0.5">
+                            <span class="text-sm font-medium"
+                                :class="!enabled ? 'text-[var(--text-primary)]' : 'text-[var(--text-secondary)]'">
+                                Off
+                            </span>
+                            <span class="text-xs leading-snug"
+                                :class="!enabled ? 'text-[var(--text-secondary)]' : 'text-[var(--text-dim)]'">
+                                Disable web access
+                            </span>
+                        </div>
+                    </button>
+
+                    <button @click="emit('update:enabled', true)"
+                        class="flex items-start gap-3 rounded-xl px-3 py-2.5 text-left transition-colors duration-150"
+                        :class="enabled
+                            ? 'bg-[var(--color-active)]'
+                            : '@hover:bg-[var(--color-hover)]'">
+                        <span class="i-tabler-world text-5 mt-0.5 shrink-0"
+                            :class="enabled ? 'text-[var(--color-accent)]' : 'text-[var(--text-dim)]'"></span>
+                        <div class="flex flex-col gap-0.5">
+                            <span class="text-sm font-medium"
+                                :class="enabled ? 'text-[var(--text-primary)]' : 'text-[var(--text-secondary)]'">
+                                Auto
+                            </span>
+                            <span class="text-xs leading-snug"
+                                :class="enabled ? 'text-[var(--text-secondary)]' : 'text-[var(--text-dim)]'">
+                                Search the web automatically when needed
+                            </span>
+                        </div>
+                    </button>
+                </div>
+
+                <!-- Right: settings (only when enabled) -->
+                <div v-if="enabled" class="flex flex-col gap-3 border-l border-[var(--color-border)] pl-4">
+                    <!-- Max results stepper -->
+                    <div class="flex flex-col gap-1.5">
+                        <span class="text-xs font-medium text-[var(--text-tertiary)] uppercase tracking-wider">
+                            Max results
+                        </span>
+                        <div class="flex items-center gap-2">
+                            <button @click="adjustMaxResults(-1)"
+                                class="h-7 w-7 flex items-center justify-center rounded-lg @hover:bg-[var(--color-hover)] text-[var(--text-secondary)] @hover:text-[var(--text-primary)] transition-colors duration-150">
+                                <span class="text-sm i-mynaui-minus"></span>
+                            </button>
+                            <span class="w-8 text-center text-sm font-medium tabular-nums text-[var(--text-primary)]">
+                                {{ maxResults }}
+                            </span>
+                            <button @click="adjustMaxResults(1)"
+                                class="h-7 w-7 flex items-center justify-center rounded-lg @hover:bg-[var(--color-hover)] text-[var(--text-secondary)] @hover:text-[var(--text-primary)] transition-colors duration-150">
+                                <span class="text-sm i-mynaui-plus"></span>
+                            </button>
+                        </div>
+                    </div>
+
+                    <!-- Rerank toggle -->
+                    <div class="flex flex-col gap-1.5">
+                        <span class="text-xs font-medium text-[var(--color-tertiary)] uppercase tracking-wider">
+                            Rerank
+                        </span>
+                        <div v-if="!isRerankConfigured"
+                            class="flex items-center gap-2 px-3 py-1.5 rounded-lg bg-[var(--bg-container)] border border-[var(--color-border)]">
+                            <button @click="() => { openSettings(); close() }"
+                                class="flex items-center gap-2 text-xs text-[var(--text-dim)] hover:text-[var(--text-primary)] transition-colors duration-150">
+                                <span class="i-mynaui-sparkles text-3.5"></span>
+                                <span>Setup reranking</span>
+                                <span class="i-mynaui-arrow-right text-3 ml-auto"></span>
+                            </button>
+                        </div>
+                        <label v-else
+                            class="flex items-center gap-2 px-3 py-1.5 rounded-lg bg-[var(--bg-container)] border border-[var(--color-border)] cursor-pointer @hover:border-[var(--color-hover)] transition-colors duration-150">
+                            <input type="checkbox" :checked="props.rerank"
+                                @change="emit('update:rerank', ($event.target as HTMLInputElement).checked)"
+                                class="w-4 h-4 rounded border-[var(--color-border)] bg-transparent text-[var(--color-accent)] focus:ring-[var(--color-accent)] focus:ring-offset-0 cursor-pointer" />
+                            <span class="text-xs text-[var(--text-secondary)]">Rerank results</span>
+                        </label>
+                    </div>
+                </div>
+            </div>
+        </template>
+    </Dropdown>
+</template>
diff --git a/app/composables/useAgents.ts b/app/composables/useAgents.ts
index d9a5618..0185c3c 100644
--- a/app/composables/useAgents.ts
+++ b/app/composables/useAgents.ts
@@ -123,9 +123,9 @@ export const useAgents = async () => {
                     a.id === id ? agent : a
                 );
             },
-            async onResponse() {
-                await refresh();
-            }
+            // async onResponse() {
+            //     await refresh();
+            // }
         });
     }
 
@@ -166,7 +166,7 @@ export const useAgents = async () => {
             body: topic,
             onRequest() {
                 agents.value = agents.value.map(a =>
-                    a.id === agentId ? { ...a, topics: [{ ...topic, createdAt: new Date() }, ...a.topics] } : a
+                    a.id === agentId ? { ...a, topics: [{ ...topic, createdAt: new Date() }, ...a.topics] } as AgentWithTopics : a
                 );
             },
             onResponseError() {
diff --git a/app/composables/useChat.ts b/app/composables/useChat.ts
index c28228c..91028de 100644
--- a/app/composables/useChat.ts
+++ b/app/composables/useChat.ts
@@ -52,7 +52,6 @@ export const useChat = async (topicId: MaybeRef<string>, connect = true) => {
             if (!msg) continue;
 
             for (const [partId, content] of parts) {
-                console.log("content", content);
                 const part = msg.parts?.find(p => p.id === partId);
                 if (part) {
                     part.content += content;
@@ -135,7 +134,6 @@ export const useChat = async (topicId: MaybeRef<string>, connect = true) => {
                 const messageBuffer = textDeltaBuffer.get(payload.messageId)!;
                 const existing = messageBuffer.get(payload.partId) || '';
                 messageBuffer.set(payload.partId, existing + payload.content);
-                console.log("messageBuffer", messageBuffer, existing + payload.content);
 
                 scheduleFlush();
                 break;
@@ -303,8 +301,6 @@ export const useChat = async (topicId: MaybeRef<string>, connect = true) => {
         baseMessage: BaseMessage,
         onRequest?: () => void,
     ): Promise<Result<void, ChatErrorType>> => {
-        console.log("sendMessage", baseMessage);
-
         const { user } = useAuth();
         if (!user.value) return Err(ChatErrorType.NoUser);
 
@@ -316,8 +312,6 @@ export const useChat = async (topicId: MaybeRef<string>, connect = true) => {
                 fileIds: baseMessage.fileIds,
             }
 
-            console.log("sendMessage", message, baseMessage.content, baseMessage.fileIds);
-
             await $fetch(`/api/topic/${unref(topicId)}/message`, {
                 method: 'POST',
                 body: {
@@ -355,19 +349,46 @@ export const useChat = async (topicId: MaybeRef<string>, connect = true) => {
         }
     };
 
-    const startGeneration = async (model: ModelWithProvider) => {
-        console.log(model.provider);
+    const startGeneration = async (model: ModelWithProvider, parentMessageId?: string, agentIdOverride?: string) => {
         const providerApiKeyRes = await getProviderAPIKey(model.provider);
         if (providerApiKeyRes.ok === false) {
             return providerApiKeyRes;
         }
         const providerApiKey = providerApiKeyRes.data;
 
+        const { agents } = await useAgents();
+
+        const agent = agents.value.find(agent => agent.id === (agentIdOverride || topic.value?.agentId));
+        if (!agent) {
+            return Err(ChatErrorType.NoAgent);
+        }
+
+        const { settings } = await useUserSettings();
+
+        let rerank = undefined;
+
+        if (agent.config?.search?.rerank && settings.value.systemAssistants.rerank.enabled) {
+            const rerankModel = await useModels().then(m => m.allModels.value.find(m => m.id === settings.value.systemAssistants.rerank.modelId));
+            if (rerankModel) {
+                const rerankProviderApiKeyRes = await getProviderAPIKey(rerankModel.provider);
+                if (rerankProviderApiKeyRes.ok === false) {
+                    return rerankProviderApiKeyRes;
+                }
+                const rerankProviderApiKey = rerankProviderApiKeyRes.data;
+                rerank = {
+                    modelId: rerankModel.id,
+                    providerApiKey: rerankProviderApiKey,
+                };
+            }
+        }
+
         await $fetch(`/api/topic/${unref(topicId)}/chat`, {
             method: 'POST',
             body: {
                 modelId: model.id,
                 providerApiKey,
+                rerank,
+                parentMessageId,
             },
         });
     }
@@ -399,12 +420,6 @@ export const useChat = async (topicId: MaybeRef<string>, connect = true) => {
         }
         let targetMessageIndex = topicMessages.indexOf(targetMessage);
 
-        const providerApiKeyRes = await getProviderAPIKey(model.provider);
-        if (providerApiKeyRes.ok === false) {
-            return providerApiKeyRes;
-        }
-        const providerApiKey = providerApiKeyRes.data;
-
         let parentMessageId = undefined;
         let focusedMessages: MessageEntity[] | undefined;
         if (targetMessage.role === 'user') {
@@ -429,14 +444,7 @@ export const useChat = async (topicId: MaybeRef<string>, connect = true) => {
             focusedMessages = topicMessages;
         }
 
-        await $fetch(`/api/topic/${unref(topicId)}/chat`, {
-            method: 'POST',
-            body: {
-                modelId: model.id,
-                providerApiKey,
-                parentMessageId,
-            },
-        });
+        startGeneration(model, parentMessageId);
 
         return Ok(undefined);
     }
diff --git a/app/composables/useModels.ts b/app/composables/useModels.ts
index 9b74add..01a4f8e 100644
--- a/app/composables/useModels.ts
+++ b/app/composables/useModels.ts
@@ -6,6 +6,7 @@ export type Provider = typeof schema.providers.$inferSelect;
 
 export type ProviderWithModels = Provider & {
     models: Model[];
+    defaultBaseUrl?: string;
 };
 
 export type ModelWithProvider = Model & {
diff --git a/server/api/topic/[topicId]/chat/index.post.ts b/server/api/topic/[topicId]/chat/index.post.ts
index ef0d52e..ffd0271 100644
--- a/server/api/topic/[topicId]/chat/index.post.ts
+++ b/server/api/topic/[topicId]/chat/index.post.ts
@@ -3,15 +3,16 @@ import * as z from 'zod';
 import { type MessageEntity } from '~/composables/useChat';
 import { promises as fs } from 'fs';
 import { glob } from 'glob';
-import { type ModelMessage, streamText, type StreamTextTransform, tool } from "ai";
+import { isLoopFinished, type ModelMessage, rerank, type RerankingModel, streamText, type StreamTextTransform, Tool, tool } from "ai";
 import { generations, messageParts, messages, toolCalls, ToolCallType } from "~~/drizzle/schema";
 import { topicEvents } from "~~/server/utils/events";
 import { nanoid } from "nanoid";
-import { spawn } from "child_process";
+import { Monty, MontyRuntimeError, MontySyntaxError, MontyTypingError } from '@pydantic/monty';
 import { type Model } from "~/composables/useModels";
 import { eq } from "drizzle-orm";
 import { buildFocusedMessageTree, buildMessageTree, marshallMessages } from "~~/utils/message";
 import path from "path";
+import { isRerankingProvider } from "~~/server/utils/ai-provider";
 
 export default defineEventHandler(async (event) => {
     await protectRoute(event);
@@ -22,6 +23,10 @@ export default defineEventHandler(async (event) => {
     const result = await readValidatedBody(event, z.object({
         parentMessageId: z.string().optional(),
         modelId: z.string(),
+        rerank: z.object({
+            modelId: z.string(),
+            providerApiKey: z.string().optional(),
+        }).optional(),
         args: z.record(z.string(), z.any()).optional(),
         providerApiKey: z.string().optional(),
     }).safeParse);
@@ -33,7 +38,7 @@ export default defineEventHandler(async (event) => {
         });
     }
 
-    const { parentMessageId, modelId, providerApiKey, args } = result.data;
+    const { parentMessageId, modelId, rerank: rerankConfig, providerApiKey, args } = result.data;
 
     const topic = await db.query.topics.findFirst({
         where: {
@@ -109,6 +114,91 @@ export default defineEventHandler(async (event) => {
         }
     }
 
+    const agentSearchConfig = topic.agent.config?.search;
+    let searchParam: false | { config: SearchTheWebConfig } = false;
+
+    if (agentSearchConfig?.enabled) {
+        searchParam = { config: { rerank: false, maxResults: agentSearchConfig.maxResults ?? 10 } };
+
+        if (agentSearchConfig.rerank && rerankConfig) {
+            const rerankModel = await db.query.models.findFirst({
+                where: {
+                    id: rerankConfig.modelId,
+                    userId,
+                },
+                with: {
+                    provider: true,
+                }
+            });
+
+            if (rerankModel === undefined) {
+                throw createError({
+                    statusCode: 400,
+                    statusMessage: 'Invalid rerank model',
+                    data: {
+                        code: 'INVALID_RERANK_MODEL',
+                        ok: false,
+                    }
+                });
+            }
+
+            const rerankProviderDetails = await getProviderDetails(rerankModel.provider, rerankConfig.providerApiKey, rerankModel);
+            if (!rerankProviderDetails.ok) {
+                switch (rerankProviderDetails.error) {
+                    case GatewayFetchError.NoProviderApiKey: {
+                        throw createError({
+                            statusCode: 400,
+                            statusMessage: `${rerankModel.provider.type} provider requires an API key`,
+                            data: {
+                                code: 'NO_RERANK_PROVIDER_API_KEY',
+                                ok: false,
+                            }
+                        });
+                    }
+                    case GatewayFetchError.NoProviderBaseUrl: {
+                        throw createError({
+                            statusCode: 400,
+                            statusMessage: 'Invalid provider URL',
+                            data: {
+                                code: 'BAD_RERANK_PROVIDER_URL',
+                                ok: false,
+                            }
+                        });
+                    }
+                }
+            }
+
+            const { gateway: rerankGateway } = rerankProviderDetails.data;
+            if (rerankGateway === null) {
+                throw createError({
+                    statusCode: 500,
+                    statusMessage: 'Invalid gateway',
+                    data: {
+                        code: 'INVALID_RERANK_GATEWAY',
+                        ok: false,
+                    }
+                });
+            }
+
+            if (isRerankingProvider(rerankGateway.gateway)) {
+                searchParam.config = {
+                    rerank: true,
+                    maxResults: agentSearchConfig.maxResults ?? 10,
+                    model: rerankGateway.gateway.reranking(rerankModel.externalId)
+                };
+            } else {
+                throw createError({
+                    statusCode: 500,
+                    statusMessage: 'Invalid rerank model',
+                    data: {
+                        code: 'INVALID_RERANK_MODEL',
+                        ok: false,
+                    }
+                });
+            }
+        }
+    }
+
     let agentmessage = await db.transaction(async tx => {
         const generationId = nanoid();
 
@@ -234,6 +324,7 @@ export default defineEventHandler(async (event) => {
         generateResponse(
             agentmessage as MessageEntity,
             { gateway: gateway.gateway, model, parameters: args },
+            searchParam,
             agentmessage.generationId!,
             userId,
             topicId,
@@ -255,54 +346,83 @@ const INTERNAL_ERROR = 'An internal error occurred';
 // todo message takes in variadics like console.log
 const todo = (...args: any[]) => {
     console.error('TODO', ...args);
-    throw new Error('TODO');
+};
+
+const formatPartId = (partType: string, existingId: string) => {
+    return `veridian__part-${partType}-${existingId}-${nanoid()}`;
+};
+
+const formatToolCallId = (nativeId: string) => {
+    return `veridian__tool-${nativeId}-${nanoid()}`;
 };
 
 const evalPython = async (code: string) => {
-    // This is the wrapper logic from above, minified or stored as a string
-    // Or you can save the wrapper script to a file and call that.
-    const wrapper = `
-import ast
-import sys
-code = sys.stdin.read()
-tree = ast.parse(code)
-last_node = tree.body[-1] if tree.body else None
-namespace = {}
-if len(tree.body) > 1:
-    exec(compile(ast.Module(tree.body[:-1], []), "<ast>", "exec"), namespace)
-if isinstance(last_node, ast.Expr):
-    res = eval(compile(ast.Expression(last_node.value), "<ast>", "eval"), namespace)
-    if res is not None: print(res)
-elif last_node:
-    exec(compile(ast.Module([last_node], []), "<ast>", "exec"), namespace)
-`.trim();
+    try {
+        let stdout = ''
+        const printCallback = (_: string, text: string) => {
+            stdout += text
+        }
+        const m = new Monty(code)
+        await m.run({ printCallback })
+        return stdout
+    } catch (error) {
+        if (error instanceof MontySyntaxError) {
+            console.log('Syntax error:', error.message)
+        } else if (error instanceof MontyRuntimeError) {
+            console.log('Runtime error:', error.message)
+            console.log('Traceback:', error.traceback())
+        } else if (error instanceof MontyTypingError) {
+            console.log('Type error:', error.displayDiagnostics())
+        }
+    }
+};
 
-    return new Promise<string>((resolve, reject) => {
-        const child = spawn('python3', ['-c', wrapper]);
+interface SearchTheWebRerankedConfig {
+    rerank: true;
+    maxResults: number;
+    model: RerankingModel;
+}
 
-        let output = '';
-        let errorOutput = '';
+interface SearchTheWebUnrankedConfig {
+    rerank?: false;
+    maxResults: number;
+}
 
-        child.stdout.on('data', (data) => {
-            output += data.toString();
-        });
+type SearchTheWebConfig = SearchTheWebRerankedConfig | SearchTheWebUnrankedConfig;
 
-        child.stderr.on('data', (data) => {
-            errorOutput += data.toString();
-        });
-
-        child.on('close', (exitCode) => {
-            if (exitCode !== 0) {
-                reject(errorOutput || `Exit code ${exitCode}`);
-            } else {
-                resolve(output.trim());
+export const searchTheWeb = (config: SearchTheWebConfig) => {
+    return async (query: string) => {
+        const results = await $fetch<any>(`${process.env.SEARXNG_URL}/search`, {
+            query: {
+                q: query,
+                format: 'json',
             }
         });
 
-        // Send the agent's code to the wrapper via stdin
-        child.stdin.write(code);
-        child.stdin.end();
-    });
+        const sites = results.results.map((item: any) => ({
+            title: item.title,
+            link: item.url,
+            snippet: item.content,
+        }));
+
+        if (config.rerank) {
+            const { ranking } = await rerank({
+                model: config.model,
+                query,
+                documents: sites.map(site => site.snippet),
+            });
+
+            const ranked_sites = [];
+
+            for (let i = 0; i < ranking.length; i++) {
+                ranked_sites.push(sites[ranking[i]!.originalIndex]);
+            }
+
+            return ranked_sites.slice(0, config.maxResults || 10);
+        }
+
+        return sites.slice(0, config.maxResults || 10);
+    };
 };
 
 // TODO: obviously come up with a better way for the user to define their own tools
@@ -462,6 +582,9 @@ async function generateResponse(
         model: Model,
         parameters?: Record<string, any>,
     },
+    search: false | {
+        config: SearchTheWebConfig,
+    },
     generationId: string,
     userId: string,
     topicId: string,
@@ -477,9 +600,10 @@ async function generateResponse(
     let ttft = undefined;
     const activeParts = new Map<string, { id: string; accumulatedContent: string; providerOptions?: any }>();
     const activeToolCalls = new Set<string>();
+    const nativeToDbToolCallId = new Map<string, string>();
 
     // TODO: somehow let the user turn on and off tools
-    const tools = {
+    const tools: Record<string, Tool> = {
         // writeFile: tool({
         //     inputSchema: z.object({
         //         path: z.string(),
@@ -504,7 +628,25 @@ async function generateResponse(
         bash: bashTool,
     };
 
-    console.log({ messages });
+    if (search) {
+        tools.search = tool({
+            description: 'Searches the web',
+            inputSchema: z.object({
+                query: z.string(),
+            }),
+            outputSchema: z.array(
+                z.object({
+                    title: z.string(),
+                    link: z.string(),
+                    snippet: z.string(),
+                    engine: z.string(),
+                })
+            ),
+            execute: async ({ query }) => {
+                return await searchTheWeb(search.config)(query);
+            },
+        });
+    }
 
     const response = streamText({
         model: model.gateway(model.model.externalId),
@@ -518,8 +660,7 @@ async function generateResponse(
             }
         },
         experimental_transform: streamTransoforms,
-        // a little trick that makes it so that the stream doesnt stop because of tool calls, and will continue an unbounded amount of time and steps
-        stopWhen: [],
+        stopWhen: isLoopFinished(),
         tools: model.model.capabilities.includes('tools') ? tools : undefined,
         onError: async (error: any) => {
             // TODO: the docs say "The stream processing will pause until the callback promise is resolved." Suggesting that this error might not be fatal?
@@ -646,8 +787,10 @@ async function generateResponse(
                 case 'reasoning-start': {
                     const type = token.type.split('-')[0] as 'text' | 'reasoning';
                     const key = `${type}-${curStepIdx}`;
+                    const nativeId = nanoid();
 
                     const [part] = await db.insert(messageParts).values({
+                        id: formatPartId(type, nativeId),
                         userId,
                         topicId,
                         messageId: message.id,
@@ -772,9 +915,11 @@ async function generateResponse(
                     key = `tool-call-${curStepIdx}`;
 
                     const toolCallId = token.id;
+                    const dbToolCallId = formatToolCallId(toolCallId);
+                    nativeToDbToolCallId.set(toolCallId, dbToolCallId);
 
                     const [toolCall] = await db.insert(toolCalls).values({
-                        id: toolCallId,
+                        id: dbToolCallId,
                         userId: userId,
                         toolName: token.toolName,
                         status: 'pending',
@@ -789,10 +934,11 @@ async function generateResponse(
                     }
 
                     const [part] = await db.insert(messageParts).values({
+                        id: formatPartId('tool-call', dbToolCallId),
                         userId,
                         topicId,
                         messageId: message.id,
-                        toolCallId,
+                        toolCallId: dbToolCallId,
                         type: 'tool-call',
                         content: null,
                         finished: false,
@@ -816,7 +962,7 @@ async function generateResponse(
                         }
                     })
 
-                    activeToolCalls.add(toolCallId);
+                    activeToolCalls.add(dbToolCallId);
 
                     activeParts.set(key, { id: part.id, accumulatedContent: '' });
                 } break;
@@ -838,7 +984,9 @@ async function generateResponse(
                             break;
                     }
 
-                    if (activeToolCalls.has(token.toolCallId)) {
+                    const dbToolCallIdFromMap = nativeToDbToolCallId.get(token.toolCallId);
+                    if (dbToolCallIdFromMap && activeToolCalls.has(dbToolCallIdFromMap)) {
+                        const dbToolCallId = dbToolCallIdFromMap;
                         await db.update(toolCalls)
                             .set({
                                 status: 'pending',
@@ -847,13 +995,13 @@ async function generateResponse(
                                     value: inputValue,
                                 },
                             })
-                            .where(eq(toolCalls.id, token.toolCallId));
+                            .where(eq(toolCalls.id, dbToolCallId));
 
                         await topicEvents.emit(topicId, {
                             type: 'tool-call-delta',
                             payload: {
                                 messageId: message.id,
-                                toolCallId: token.toolCallId,
+                                toolCallId: dbToolCallId,
                                 toolName: token.toolName,
                                 input: {
                                     type: inputType,
@@ -865,9 +1013,11 @@ async function generateResponse(
                         key = `tool-call-${curStepIdx}`;
 
                         const toolCallId = token.toolCallId;
+                        const dbToolCallId = formatToolCallId(toolCallId);
+                        nativeToDbToolCallId.set(toolCallId, dbToolCallId);
 
                         const [toolCall] = await db.insert(toolCalls).values({
-                            id: toolCallId,
+                            id: dbToolCallId,
                             userId: userId,
                             toolName: token.toolName,
                             status: 'pending',
@@ -885,10 +1035,11 @@ async function generateResponse(
                         }
 
                         const [part] = await db.insert(messageParts).values({
+                            id: formatPartId('tool-call', dbToolCallId),
                             userId,
                             topicId,
                             messageId: message.id,
-                            toolCallId: token.toolCallId,
+                            toolCallId: dbToolCallId,
                             type: 'tool-call',
                             content: null,
                             finished: false,
@@ -911,7 +1062,7 @@ async function generateResponse(
                             }
                         })
 
-                        activeToolCalls.add(toolCallId);
+                        activeToolCalls.add(dbToolCallId);
 
                         activeParts.set(key, { id: part.id, accumulatedContent: '' });
                     }
@@ -919,6 +1070,7 @@ async function generateResponse(
                 case 'tool-result': {
                     let outputType: ToolCallType = ToolCallType.Text;
                     let outputValue: string = '';
+                    const dbToolCallId = nativeToDbToolCallId.get(token.toolCallId);
 
                     switch (typeof token.output) {
                         case 'string':
@@ -931,52 +1083,56 @@ async function generateResponse(
                             break;
                         default:
                             console.error('Unknown output type', token.output);
-                            await db.update(toolCalls).set({
-                                status: 'failed',
-                                error: { type: ToolCallType.Text, value: 'Tool returned invalid output' }
-                            }).where(eq(toolCalls.id, token.toolCallId));
-
-                            await topicEvents.emit(topicId, {
-                                type: 'tool-call-delta',
-                                payload: {
-                                    messageId: message.id,
-                                    toolCallId: token.toolCallId,
-                                    toolName: token.toolName,
-                                    output: {
-                                        type: outputType,
-                                        value: outputValue,
-                                    },
+                            if (dbToolCallId) {
+                                await db.update(toolCalls).set({
                                     status: 'failed',
-                                }
-                            })
+                                    error: { type: ToolCallType.Text, value: 'Tool returned invalid output' }
+                                }).where(eq(toolCalls.id, dbToolCallId));
 
-                            activeToolCalls.delete(token.toolCallId);
+                                await topicEvents.emit(topicId, {
+                                    type: 'tool-call-delta',
+                                    payload: {
+                                        messageId: message.id,
+                                        toolCallId: dbToolCallId,
+                                        toolName: token.toolName,
+                                        output: {
+                                            type: outputType,
+                                            value: outputValue,
+                                        },
+                                        status: 'failed',
+                                    }
+                                })
+
+                                activeToolCalls.delete(dbToolCallId);
+                            }
                             break;
                     }
 
-                    await db.update(toolCalls).set({
-                        status: 'completed',
-                        output: {
-                            type: outputType,
-                            value: outputValue,
-                        },
-                    }).where(eq(toolCalls.id, token.toolCallId));
-
-                    await topicEvents.emit(topicId, {
-                        type: 'tool-call-delta',
-                        payload: {
-                            messageId: message.id,
-                            toolCallId: token.toolCallId,
-                            toolName: token.toolName,
+                    if (dbToolCallId) {
+                        await db.update(toolCalls).set({
+                            status: 'completed',
                             output: {
                                 type: outputType,
                                 value: outputValue,
                             },
-                            status: 'completed',
-                        }
-                    })
+                        }).where(eq(toolCalls.id, dbToolCallId));
 
-                    activeToolCalls.delete(token.toolCallId);
+                        await topicEvents.emit(topicId, {
+                            type: 'tool-call-delta',
+                            payload: {
+                                messageId: message.id,
+                                toolCallId: dbToolCallId,
+                                toolName: token.toolName,
+                                output: {
+                                    type: outputType,
+                                    value: outputValue,
+                                },
+                                status: 'completed',
+                            }
+                        })
+
+                        activeToolCalls.delete(dbToolCallId);
+                    }
                     break;
                 }
 
@@ -1002,20 +1158,21 @@ async function generateResponse(
                             break;
                     }
 
-                    if (activeToolCalls.has(token.toolCallId)) {
+                    const existingDbToolCallId = nativeToDbToolCallId.get(token.toolCallId);
+                    if (existingDbToolCallId && activeToolCalls.has(existingDbToolCallId)) {
                         await db.update(toolCalls).set({
                             status: 'failed',
                             error: {
                                 type: outputType as ToolCallType,
                                 value: outputValue as string,
                             }
-                        }).where(eq(toolCalls.id, token.toolCallId));
+                        }).where(eq(toolCalls.id, existingDbToolCallId));
 
                         await topicEvents.emit(topicId, {
                             type: 'tool-call-delta',
                             payload: {
                                 messageId: message.id,
-                                toolCallId: token.toolCallId,
+                                toolCallId: existingDbToolCallId,
                                 toolName: token.toolName,
                                 error: {
                                     type: outputType as ToolCallType,
@@ -1027,9 +1184,11 @@ async function generateResponse(
                         key = `tool-call-${curStepIdx}`;
 
                         const toolCallId = token.toolCallId;
+                        const dbToolCallId = formatToolCallId(toolCallId);
+                        nativeToDbToolCallId.set(toolCallId, dbToolCallId);
 
                         const [toolCall] = await db.insert(toolCalls).values({
-                            id: toolCallId,
+                            id: dbToolCallId,
                             userId: userId,
                             toolName: token.toolName,
                             status: 'failed',
@@ -1047,10 +1206,11 @@ async function generateResponse(
                         }
 
                         const [part] = await db.insert(messageParts).values({
+                            id: formatPartId('tool-call', dbToolCallId),
                             userId,
                             topicId,
                             messageId: message.id,
-                            toolCallId: token.toolCallId,
+                            toolCallId: dbToolCallId,
                             type: 'tool-call',
                             content: null,
                             finished: false,
@@ -1073,7 +1233,9 @@ async function generateResponse(
                         activeParts.set(key, { id: part.id, accumulatedContent: '' });
                     }
 
-                    activeToolCalls.delete(token.toolCallId);
+                    if (existingDbToolCallId) {
+                        activeToolCalls.delete(existingDbToolCallId);
+                    }
                     break;
                 }