refactor: create generic useChat hook and simplify architecture

- Add useChat hook that receives tools as parameter and handles UI + ML state
- Create unified Zustand store (useChatStore) replacing 2 stores
- Simplify worker to only handle model loading and inference (~90 lines)
- Extract function calling logic to lib/functionCalling.ts
- Move tool definitions to config/tools.ts
- Add shared types in types/chat.ts

package.json

@@ -7,7 +7,7 @@
     "dev": "vite",
     "build": "tsc -b && vite build",
     "preview": "vite preview",
-    "typecheck": "tsc --noEmit",
+    "typecheck": "tsc -b",
     "lint": "biome lint ./src",
     "format": "biome format --write ./src",
     "check": "biome check --write ./src"

src/components/AgenticInterface.tsx

@@ -1,75 +1,40 @@
-import { useState } from "react";
-import { useFunctionCalling } from "@/hooks/useFunctionCalling";
-import { useColorStore } from "@/store/useColorStore";
+import { useCallback } from "react";
+import { colorTools } from "@/config/tools";
+import { useChat } from "@/hooks/useChat";
+import { useChatStore } from "@/store/useChatStore";
+import type { FunctionCallResult } from "@/types/chat";
 import { ChatSidebar } from "./ChatSidebar";
 import { ColorControlForm } from "./ColorControlForm";
 import { ColorVisualizer } from "./ColorVisualizer";
-import type { Message } from "./MessageBubble";
-
-const INITIAL_MESSAGES: Message[] = [
-  {
-    id: 1,
-    text: "Hi! I can change the square color for you. Ask me to set a color or ask what color it currently is!",
-    sender: "bot",
-  },
-];
 
 export function AgenticInterface() {
-  const { squareColor, setSquareColor } = useColorStore();
-  const { processMessage, continueWithToolResult, isProcessing, loadingStatus } =
-    useFunctionCalling();
-  const [messages, setMessages] = useState<Message[]>(INITIAL_MESSAGES);
-
-  const handleSendMessage = async (text: string) => {
-    const userMessage: Message = { id: Date.now(), text, sender: "user" };
-    setMessages((prev) => [...prev, userMessage]);
-
-    let botResponse: string;
-    try {
-      const result = await processMessage(text);
-
-      if (result.functionCall?.functionName === "setSquareColor") {
-        const color = result.functionCall.args.color as string | undefined;
-        if (color) {
-          setSquareColor(color);
-          // Get response from model after executing the function
-          botResponse = await continueWithToolResult("set_square_color", {
-            success: true,
-            color,
-          });
-        } else {
-          botResponse =
-            "I understood you want to change the color, but I couldn't determine which color. Please try again with a specific color.";
-        }
-      } else if (result.functionCall?.functionName === "getSquareColor") {
-        // Get response from model after executing the function
-        botResponse = await continueWithToolResult("get_square_color", {
-          color: squareColor,
-        });
-      } else if (result.textResponse) {
-        botResponse = result.textResponse;
-      } else {
-        botResponse =
-          "Sorry, I had trouble understanding that. Could you please try again?";
+  const { squareColor, setSquareColor } = useChatStore();
+
+  const handleFunctionCall = useCallback(
+    async (fc: FunctionCallResult) => {
+      if (fc?.functionName === "setSquareColor") {
+        const color = fc.args.color as string;
+        setSquareColor(color);
+        return { success: true, color };
+      }
+      if (fc?.functionName === "getSquareColor") {
+        return { color: squareColor };
       }
-    } catch {
-      botResponse =
-        "Sorry, there was an error processing your request. Please try again.";
-    }
+      return { error: "Unknown function" };
+    },
+    [squareColor, setSquareColor],
+  );
 
-    const botMessage: Message = {
-      id: Date.now() + 1,
-      text: botResponse,
-      sender: "bot",
-    };
-    setMessages((prev) => [...prev, botMessage]);
-  };
+  const { messages, sendMessage, isProcessing, loadingStatus } = useChat({
+    tools: colorTools,
+    onFunctionCall: handleFunctionCall,
+  });
 
   return (
     <div className="flex flex-col md:flex-row h-screen w-screen overflow-hidden bg-background text-foreground font-sans selection:bg-primary/30">
       <ChatSidebar
         messages={messages}
-        onSendMessage={handleSendMessage}
+        onSendMessage={sendMessage}
         isLoading={isProcessing || loadingStatus.isLoading}
         modelReady={loadingStatus.isModelReady}
         progress={loadingStatus.downloadProgress}

src/components/MessageBubble.tsx

@@ -1,15 +1,11 @@
 import { IconRobot, IconUser } from "@tabler/icons-react";
+import type { UIMessage } from "@/types/chat";
 
-type MessageSender = "user" | "bot";
-
-type Message = {
-  id: number;
-  text: string;
-  sender: MessageSender;
-};
+// Re-export for backwards compatibility
+export type Message = UIMessage;
 
 type MessageBubbleProps = {
-  message: Message;
+  message: UIMessage;
 };
 
 export function MessageBubble({ message }: MessageBubbleProps) {
@@ -41,5 +37,3 @@ export function MessageBubble({ message }: MessageBubbleProps) {
     </div>
   );
 }
-
-export type { Message, MessageSender };

src/config/tools.ts

@@ -0,0 +1,34 @@
+import type { ToolDefinition } from "@/types/chat";
+
+export const colorTools: ToolDefinition[] = [
+  {
+    type: "function",
+    function: {
+      name: "set_square_color",
+      description: "Sets the color of the square displayed on the screen.",
+      parameters: {
+        type: "object",
+        properties: {
+          color: {
+            type: "string",
+            description: "The color to set, e.g. red, blue, green",
+          },
+        },
+        required: ["color"],
+      },
+    },
+  },
+  {
+    type: "function",
+    function: {
+      name: "get_square_color",
+      description:
+        "Returns the current color of the square. Use this when the user asks 'what color is the square' or 'tell me the color'.",
+      parameters: {
+        type: "object",
+        properties: {},
+        required: [],
+      },
+    },
+  },
+];

src/hooks/{useFunctionCalling.ts → useChat.ts}

@@ -1,6 +1,20 @@
 import * as Comlink from "comlink";
 import { useCallback, useEffect, useRef, useState } from "react";
-import type { FunctionCallResult, LoadingProgress, WorkerAPI } from "@/worker";
+import { useChatStore } from "@/store/useChatStore";
+import {
+  parseFunctionCall,
+  cleanupModelResponse,
+} from "@/lib/functionCalling";
+import { toSnakeCase } from "@/lib/utils";
+import type {
+  ToolDefinition,
+  ChatMessage,
+  FunctionCallResult,
+} from "@/types/chat";
+import type { WorkerAPI, LoadingProgress } from "@/worker";
+
+const DEVELOPER_PROMPT =
+  "You are a model that can do function calling with the following functions";
 
 export type LoadingStatus = {
   isLoading: boolean;
@@ -10,7 +24,12 @@ export type LoadingStatus = {
   error?: string;
 };
 
-export function useFunctionCalling() {
+type UseChatOptions = {
+  tools: ToolDefinition[];
+  onFunctionCall?: (fc: FunctionCallResult) => Promise<unknown>;
+};
+
+export function useChat({ tools, onFunctionCall }: UseChatOptions) {
   const workerRef = useRef<Worker | null>(null);
   const apiRef = useRef<Comlink.Remote<WorkerAPI> | null>(null);
   const initializationPromiseRef = useRef<Promise<void> | null>(null);
@@ -22,6 +41,9 @@ export function useFunctionCalling() {
   });
   const [isProcessing, setIsProcessing] = useState(false);
 
+  const { messages, addMessage, setConversation } = useChatStore();
+
+  // Initialize worker
   useEffect(() => {
     if (!("gpu" in navigator)) {
       setLoadingStatus((prev) => ({
@@ -45,12 +67,12 @@ export function useFunctionCalling() {
     };
   }, []);
 
+  // Initialize model
   const initModel = useCallback(async () => {
     const api = apiRef.current;
     if (!api) return;
 
     setLoadingStatus((prev) => ({ ...prev, isLoading: true }));
-    console.log("[useFunctionCalling] Starting model initialization...");
 
     const fileProgress = new Map<string, { loaded: number; total: number }>();
     let lastUpdate = 0;
@@ -65,7 +87,6 @@ export function useFunctionCalling() {
             });
           }
 
-          // Throttle updates to avoid "Maximum update depth exceeded"
           const now = Date.now();
           if (now - lastUpdate < 100) return;
           lastUpdate = now;
@@ -80,7 +101,6 @@ export function useFunctionCalling() {
           const currentProgress =
             grandTotal > 0 ? (totalLoaded / grandTotal) * 100 : 0;
 
-          console.log("[useFunctionCalling] Progress:", progress);
           setLoadingStatus((prev) => ({
             ...prev,
             currentFile: progress.file ?? prev.currentFile,
@@ -88,14 +108,13 @@ export function useFunctionCalling() {
           }));
         }),
       );
-      console.log("[useFunctionCalling] Model loaded successfully!");
+
       setLoadingStatus({
         isLoading: false,
         isModelReady: true,
         downloadProgress: 100,
       });
     } catch (error) {
-      console.error("[useFunctionCalling] Failed to load model:", error);
       setLoadingStatus((prev) => ({
         ...prev,
         isLoading: false,
@@ -105,13 +124,56 @@ export function useFunctionCalling() {
     }
   }, []);
 
-  const processMessage = useCallback(
-    async (
-      text: string,
-    ): Promise<{ functionCall: FunctionCallResult; textResponse?: string }> => {
+  // Continue conversation with tool result
+  const continueWithResult = useCallback(
+    async (result: unknown): Promise<string> => {
+      const api = apiRef.current;
+
+      // Get fresh state directly from store to avoid stale closure
+      const { conversationMessages, lastFunctionCall, clearConversation } =
+        useChatStore.getState();
+
+      if (!api || !lastFunctionCall || conversationMessages.length === 0) {
+        return "Error: no conversation to continue";
+      }
+
+      const snakeCaseName = toSnakeCase(lastFunctionCall.functionName);
+
+      const fullMessages: ChatMessage[] = [
+        ...conversationMessages,
+        {
+          role: "assistant",
+          tool_calls: [
+            {
+              type: "function",
+              function: {
+                name: snakeCaseName,
+                arguments: lastFunctionCall.args,
+              },
+            },
+          ],
+        },
+        {
+          role: "tool",
+          content: [{ name: snakeCaseName, response: result }],
+        },
+      ];
+
+      const rawOutput = await api.generate(fullMessages, tools);
+      clearConversation();
+
+      return cleanupModelResponse(rawOutput || "");
+    },
+    [tools],
+  );
+
+  // Send message
+  const sendMessage = useCallback(
+    async (text: string) => {
       const api = apiRef.current;
-      if (!api) return { functionCall: null };
+      if (!api) return;
 
+      // Ensure model is initialized
       if (!initializationPromiseRef.current) {
         initializationPromiseRef.current = initModel().catch((err) => {
           initializationPromiseRef.current = null;
@@ -122,36 +184,68 @@ export function useFunctionCalling() {
       try {
         await initializationPromiseRef.current;
       } catch {
-        return {
-          functionCall: null,
-          textResponse: "Failed to initialize model. Please try again.",
-        };
-      }
-
-      setIsProcessing(true);
-      try {
-        return await api.processMessage(text);
-      } finally {
-        setIsProcessing(false);
+        addMessage({
+          id: Date.now(),
+          text: "Failed to initialize model. Please try again.",
+          sender: "bot",
+        });
+        return;
       }
-    },
-    [initModel],
-  );
 
-  const continueWithToolResult = useCallback(
-    async (functionName: string, functionResult: unknown): Promise<string> => {
-      const api = apiRef.current;
-      if (!api) return "Error: Worker not available";
+      // Add user message to UI
+      addMessage({ id: Date.now(), text, sender: "user" });
 
       setIsProcessing(true);
+
       try {
-        return await api.continueWithToolResult(functionName, functionResult);
+        const mlMessages: ChatMessage[] = [
+          { role: "developer", content: DEVELOPER_PROMPT },
+          { role: "user", content: text },
+        ];
+
+        const rawOutput = await api.generate(mlMessages, tools);
+        const functionCall = parseFunctionCall(rawOutput || "");
+
+        if (functionCall && onFunctionCall) {
+          // Save conversation state for continuation
+          setConversation(mlMessages, functionCall);
+
+          // Execute function and get result
+          const result = await onFunctionCall(functionCall);
+
+          // Continue conversation with result
+          const response = await continueWithResult(result);
+          addMessage({ id: Date.now(), text: response, sender: "bot" });
+        } else if (functionCall) {
+          // Function call but no handler
+          addMessage({
+            id: Date.now(),
+            text: "Function call detected but no handler provided.",
+            sender: "bot",
+          });
+        } else {
+          // No function call, just text response
+          const response = cleanupModelResponse(rawOutput || "");
+          addMessage({ id: Date.now(), text: response, sender: "bot" });
+        }
       } finally {
         setIsProcessing(false);
       }
     },
-    [],
+    [
+      tools,
+      onFunctionCall,
+      addMessage,
+      setConversation,
+      continueWithResult,
+      initModel,
+    ],
   );
 
-  return { processMessage, continueWithToolResult, isProcessing, loadingStatus };
+  return {
+    messages,
+    sendMessage,
+    isProcessing,
+    loadingStatus,
+  };
 }

src/lib/functionCalling.ts

@@ -0,0 +1,53 @@
+import type { FunctionCallResult } from "@/types/chat";
+import { toCamelCase } from "./utils";
+
+/**
+ * Parses FunctionGemma output format:
+ * <start_function_call>call:func_name{param:<escape>value<escape>}<end_function_call>
+ */
+export function parseFunctionCall(output: string): FunctionCallResult {
+  const match = output.match(
+    /<start_function_call>call:(\w+)\{([^}]*)\}<end_function_call>/,
+  );
+
+  if (!match) return null;
+
+  const funcName = match[1];
+  const argsStr = match[2];
+
+  // Convert function name from snake_case to camelCase
+  const functionName = toCamelCase(funcName);
+
+  // Parse all arguments generically
+  // Format: key:<escape>value<escape>
+  const args: Record<string, unknown> = {};
+  const argRegex = /(\w+):<escape>([^<]*)<escape>/g;
+
+  for (const argMatch of argsStr.matchAll(argRegex)) {
+    const key = toCamelCase(argMatch[1]);
+    const value = argMatch[2].trim();
+    args[key] = value;
+  }
+
+  return { functionName, args };
+}
+
+/**
+ * Cleans up model response by removing special tokens
+ */
+export function cleanupModelResponse(text: string): string {
+  return (
+    text
+      .replace(/<start_function_call>[\s\S]*?<end_function_call>/g, "")
+      .replace(/<start_function_response>[\s\S]*?<end_function_response>/g, "")
+      .replace(/<end_of_turn>/g, "")
+      .replace(/<start_of_turn>/g, "")
+      .replace(/<\|im_end\|>/g, "")
+      .replace(/<\|im_start\|>/g, "")
+      .replace(/<\|endoftext\|>/g, "")
+      .replace(/<tool_call>[\s\S]*?<\/tool_call>/g, "")
+      .replace(/^model\s*/i, "")
+      .replace(/^assistant\s*/i, "")
+      .trim() || "Done!"
+  );
+}

src/store/useChatStore.ts

@@ -0,0 +1,50 @@
+import { create } from "zustand";
+import type {
+  ChatMessage,
+  FunctionCallResult,
+  UIMessage,
+} from "@/types/chat";
+
+const INITIAL_MESSAGE: UIMessage = {
+  id: 1,
+  text: "Hi! I can change the square color for you. Ask me to set a color or ask what color it currently is!",
+  sender: "bot",
+};
+
+interface ChatState {
+  // UI Messages (displayed in chat)
+  messages: UIMessage[];
+  addMessage: (message: UIMessage) => void;
+
+  // ML Conversation (for multi-turn function calling)
+  conversationMessages: ChatMessage[];
+  lastFunctionCall: FunctionCallResult;
+  setConversation: (
+    messages: ChatMessage[],
+    functionCall: FunctionCallResult,
+  ) => void;
+  clearConversation: () => void;
+
+  // App State
+  squareColor: string;
+  setSquareColor: (color: string) => void;
+}
+
+export const useChatStore = create<ChatState>((set) => ({
+  // UI Messages
+  messages: [INITIAL_MESSAGE],
+  addMessage: (message) =>
+    set((state) => ({ messages: [...state.messages, message] })),
+
+  // ML Conversation
+  conversationMessages: [],
+  lastFunctionCall: null,
+  setConversation: (conversationMessages, lastFunctionCall) =>
+    set({ conversationMessages, lastFunctionCall }),
+  clearConversation: () =>
+    set({ conversationMessages: [], lastFunctionCall: null }),
+
+  // App State
+  squareColor: "rebeccapurple",
+  setSquareColor: (squareColor) => set({ squareColor }),
+}));

src/store/useColorStore.ts

@@ -1,11 +0,0 @@
-import { create } from "zustand";
-
-interface ColorState {
-  squareColor: string;
-  setSquareColor: (color: string) => void;
-}
-
-export const useColorStore = create<ColorState>((set) => ({
-  squareColor: "rebeccapurple",
-  setSquareColor: (color: string) => set({ squareColor: color }),
-}));

src/types/chat.ts

@@ -0,0 +1,36 @@
+export type ToolCall = {
+  type: "function";
+  function: {
+    name: string;
+    arguments: Record<string, unknown>;
+  };
+};
+
+export type ToolDefinition = {
+  type: "function";
+  function: {
+    name: string;
+    description: string;
+    parameters: {
+      type: "object";
+      properties: Record<string, unknown>;
+      required: string[];
+    };
+  };
+};
+
+export type ChatMessage =
+  | { role: "developer" | "user"; content: string }
+  | { role: "assistant"; content?: string; tool_calls?: ToolCall[] }
+  | { role: "tool"; content: Array<{ name: string; response: unknown }> };
+
+export type FunctionCallResult = {
+  functionName: string;
+  args: Record<string, unknown>;
+} | null;
+
+export type UIMessage = {
+  id: number;
+  text: string;
+  sender: "user" | "bot";
+};

src/worker.ts

@@ -5,8 +5,7 @@ import {
   type PreTrainedTokenizer,
 } from "@huggingface/transformers";
 import * as Comlink from "comlink";
-import { createStore } from "zustand/vanilla";
-import { toCamelCase, toSnakeCase } from "./lib/utils";
+import type { ChatMessage, ToolDefinition } from "./types/chat";
 
 const MODEL_ID = "onnx-community/functiongemma-270m-it-ONNX";
 
@@ -25,42 +24,6 @@ export type LoadingProgress = {
 
 type ProgressCallback = (progress: LoadingProgress) => void;
 
-// Types for FunctionGemma multi-turn conversation
-type ToolCall = {
-  type: "function";
-  function: {
-    name: string;
-    arguments: Record<string, unknown>;
-  };
-};
-
-type ChatMessage =
-  | { role: "developer" | "user"; content: string }
-  | { role: "assistant"; content?: string; tool_calls?: ToolCall[] }
-  | { role: "tool"; content: Array<{ name: string; response: unknown }> };
-
-export type FunctionCallResult = {
-  functionName: string;
-  args: Record<string, unknown>;
-} | null;
-
-type ConversationState = {
-  messages: ChatMessage[];
-  lastFunctionCall: FunctionCallResult;
-  setMessages: (messages: ChatMessage[]) => void;
-  setLastFunctionCall: (fc: FunctionCallResult) => void;
-  clear: () => void;
-};
-
-// Store to maintain conversation state between processMessage and continueWithToolResult
-const conversationStore = createStore<ConversationState>((set) => ({
-  messages: [],
-  lastFunctionCall: null,
-  setMessages: (messages) => set({ messages }),
-  setLastFunctionCall: (fc) => set({ lastFunctionCall: fc }),
-  clear: () => set({ messages: [], lastFunctionCall: null }),
-}));
-
 function getModel(onProgress?: ProgressCallback) {
   modelPromise ??= (async () => {
     console.log("[Worker] Loading model...");
@@ -80,71 +43,6 @@ function getModel(onProgress?: ProgressCallback) {
   return modelPromise;
 }
 
-const tools = [
-  {
-    type: "function",
-    function: {
-      name: "set_square_color",
-      description: "Sets the color of the square displayed on the screen.",
-      parameters: {
-        type: "object",
-        properties: {
-          color: {
-            type: "string",
-            description: "The color to set, e.g. red, blue, green",
-          },
-        },
-        required: ["color"],
-      },
-    },
-  },
-  {
-    type: "function",
-    function: {
-      name: "get_square_color",
-      description:
-        "Returns the current color of the square. Use this when the user asks 'what color is the square' or 'tell me the color'.",
-      parameters: {
-        type: "object",
-        properties: {},
-        required: [],
-      },
-    },
-  },
-];
-
-// Exact prompt from functiongemma documentation
-const developerPrompt =
-  "You are a model that can do function calling with the following functions";
-
-// Parses functiongemma output: <start_function_call>call:func_name{param:<escape>value<escape>}<end_function_call>
-function parseFunctionCall(output: string): FunctionCallResult {
-  const match = output.match(
-    /<start_function_call>call:(\w+)\{([^}]*)\}<end_function_call>/,
-  );
-
-  if (!match) return null;
-
-  const funcName = match[1];
-  const argsStr = match[2];
-
-  // Convert function name from snake_case to camelCase
-  const functionName = toCamelCase(funcName);
-
-  // Parse all arguments generically
-  // Format: key:<escape>value<escape>
-  const args: Record<string, unknown> = {};
-  const argRegex = /(\w+):<escape>([^<]*)<escape>/g;
-
-  for (const argMatch of argsStr.matchAll(argRegex)) {
-    const key = toCamelCase(argMatch[1]);
-    const value = argMatch[2].trim();
-    args[key] = value;
-  }
-
-  return { functionName, args };
-}
-
 const workerAPI = {
   async initModel(progressCallback: ProgressCallback): Promise<void> {
     console.log("[Worker] initModel called");
@@ -152,161 +50,17 @@ const workerAPI = {
     console.log("[Worker] initModel completed");
   },
 
-  async processMessage(text: string): Promise<{
-    functionCall: FunctionCallResult;
-    textResponse?: string;
-  }> {
-    console.log("[Worker] ========== NEW REQUEST ==========");
-    console.log("[Worker] User input:", text);
-
-    if (!text.trim()) return { functionCall: null };
-
-    const { model, tokenizer } = await getModel();
-
-    const messages = [
-      { role: "developer", content: developerPrompt },
-      { role: "user", content: text },
-    ];
-
-    console.log("[Worker] Messages:", JSON.stringify(messages, null, 2));
-
-    const inputs = tokenizer.apply_chat_template(messages, {
-      tools,
-      tokenize: true,
-      add_generation_prompt: true,
-      return_dict: true,
-    }) as { input_ids: { dims: number[]; data: unknown } };
-
-    console.log("[Worker] Input token count:", inputs.input_ids.dims[1]);
-    console.log("[Worker] Input structure:", Object.keys(inputs));
-    console.log(
-      "[Worker] input_ids type:",
-      typeof inputs.input_ids,
-      inputs.input_ids.constructor?.name,
-    );
-
-    try {
-      const inputIds = Array.from(inputs.input_ids.data as Iterable<number>);
-      const decodedInput = tokenizer.decode(inputIds, {
-        skip_special_tokens: false,
-      });
-      console.log("[Worker] Decoded input (what model sees):", decodedInput);
-    } catch (e) {
-      console.log("[Worker] Could not decode input:", e);
-    }
-
-    // biome-ignore lint/complexity/noBannedTypes: model.generate type is not properly typed in transformers.js
-    const output = await (model.generate as Function)({
-      ...inputs,
-      max_new_tokens: 512,
-    });
-
-    const inputLength = inputs.input_ids.dims[1];
-    const outputLength = output.dims ? output.dims[1] : output.length;
-    console.log("[Worker] Output token count:", outputLength);
-    console.log("[Worker] New tokens generated:", outputLength - inputLength);
-
-    const decoded = tokenizer.decode(output.slice(0, [inputLength, null]), {
-      skip_special_tokens: false,
-    });
-
-    console.log("[Worker] Raw model output:", decoded);
-    const functionCall = parseFunctionCall(decoded as string);
-    console.log("[Worker] Parsed function call:", functionCall);
-
-    // If function was called, save conversation state for continueWithToolResult
-    if (functionCall) {
-      console.log(
-        "[Worker] ✅ Function call successful:",
-        functionCall.functionName,
-      );
-
-      // Save conversation state for the next turn
-      const { setMessages, setLastFunctionCall } = conversationStore.getState();
-      setMessages([
-        { role: "developer", content: developerPrompt },
-        { role: "user", content: text },
-      ]);
-      setLastFunctionCall(functionCall);
-
-      return { functionCall };
-    }
-
-    // If model generated an error format, return null to trigger fallback
-    if ((decoded as string).includes("<start_function_call>error:")) {
-      console.log(
-        "[Worker] ⚠️ Model generated error format, triggering fallback",
-      );
-      return { functionCall: null };
-    }
-
-    console.log("[Worker] ❌ No function call detected, returning fallback");
-
-    const textResponse = (decoded as string)
-      .replace(/<tool_call>[\s\S]*?<\/tool_call>/g, "")
-      .replace(/<\|im_end\|>/g, "")
-      .replace(/<\|im_start\|>/g, "")
-      .replace(/<\|endoftext\|>/g, "")
-      .replace(/<start_function_call>[\s\S]*?<end_function_call>/g, "")
-      .replace(/<start_function_response>/g, "")
-      .trim();
-
-    return { functionCall: null, textResponse: textResponse || undefined };
-  },
-
-  async continueWithToolResult(
-    functionName: string,
-    functionResult: unknown,
+  async generate(
+    messages: ChatMessage[],
+    tools: ToolDefinition[],
   ): Promise<string> {
-    console.log("[Worker] ========== CONTINUE WITH TOOL RESULT ==========");
-    console.log("[Worker] Function:", functionName);
-    console.log("[Worker] Result:", functionResult);
-
-    const { messages, lastFunctionCall, clear } = conversationStore.getState();
-
-    if (messages.length === 0 || !lastFunctionCall) {
-      console.log("[Worker] ⚠️ No conversation state found");
-      return "I apologize, but I lost track of our conversation. Could you please try again?";
-    }
+    console.log("[Worker] ========== GENERATE ==========");
+    console.log("[Worker] Messages:", JSON.stringify(messages, null, 2));
 
     const { model, tokenizer } = await getModel();
 
-    // Build the full conversation with tool result
-    // Convert camelCase back to snake_case for the model
-    const snakeCaseName = toSnakeCase(lastFunctionCall.functionName);
-
-    // Build messages in the format expected by the tokenizer
-    // Note: Using 'as unknown' to bypass strict typing since the tokenizer
-    // accepts a more flexible format than what TypeScript infers
-    const fullMessages = [
-      ...messages,
-      // Assistant turn with tool_calls
-      {
-        role: "assistant",
-        tool_calls: [
-          {
-            type: "function",
-            function: {
-              name: snakeCaseName,
-              arguments: lastFunctionCall.args,
-            },
-          },
-        ],
-      },
-      // Tool response
-      {
-        role: "tool",
-        content: [{ name: snakeCaseName, response: functionResult }],
-      },
-    ] as unknown[];
-
-    console.log(
-      "[Worker] Full messages for final response:",
-      JSON.stringify(fullMessages, null, 2),
-    );
-
     // biome-ignore lint/suspicious/noExplicitAny: tokenizer accepts flexible message format
-    const inputs = tokenizer.apply_chat_template(fullMessages as any, {
+    const inputs = tokenizer.apply_chat_template(messages as any, {
       tools,
       tokenize: true,
       add_generation_prompt: true,
@@ -315,23 +69,10 @@ const workerAPI = {
 
     console.log("[Worker] Input token count:", inputs.input_ids.dims[1]);
 
-    try {
-      const inputIds = Array.from(inputs.input_ids.data as Iterable<number>);
-      const decodedInput = tokenizer.decode(inputIds, {
-        skip_special_tokens: false,
-      });
-      console.log(
-        "[Worker] Decoded input for final response:",
-        decodedInput,
-      );
-    } catch (e) {
-      console.log("[Worker] Could not decode input:", e);
-    }
-
     // biome-ignore lint/complexity/noBannedTypes: model.generate type is not properly typed in transformers.js
     const output = await (model.generate as Function)({
       ...inputs,
-      max_new_tokens: 256,
+      max_new_tokens: 512,
     });
 
     const inputLength = inputs.input_ids.dims[1];
@@ -339,27 +80,9 @@ const workerAPI = {
       skip_special_tokens: false,
     });
 
-    console.log("[Worker] Raw final response:", decoded);
-
-    // Clear conversation state
-    clear();
-
-    // Clean up the response
-    const cleanResponse = (decoded as string)
-      .replace(/<start_function_call>[\s\S]*?<end_function_call>/g, "")
-      .replace(/<start_function_response>[\s\S]*?<end_function_response>/g, "")
-      .replace(/<\|im_end\|>/g, "")
-      .replace(/<\|im_start\|>/g, "")
-      .replace(/<\|endoftext\|>/g, "")
-      .replace(/<end_of_turn>/g, "")
-      .replace(/<start_of_turn>/g, "")
-      .replace(/^model\s*/i, "")
-      .replace(/^assistant\s*/i, "")
-      .trim();
-
-    console.log("[Worker] ✅ Final response:", cleanResponse);
+    console.log("[Worker] Raw output:", decoded);
 
-    return cleanResponse || "Done!";
+    return decoded as string;
   },
 };