feat: implement worker for LLM model loading and response generation with message handling

src/App.tsx

@@ -3,6 +3,7 @@ import React, {
   useEffect,
   useCallback,
   useRef,
+  useMemo,
 } from "react";
 import { openDB, type IDBPDatabase } from "idb";
 import {
@@ -128,6 +129,59 @@ const App: React.FC = () => {
     connectAll: connectAllMCPServers,
   } = useMCP();
 
+  // Memoize tool schemas to avoid recalculating on every render
+  const toolSchemas = useMemo(() => {
+    return tools
+      .filter((tool) => tool.enabled)
+      .map((tool) => generateSchemaFromCode(tool.code));
+  }, [tools]);
+
+  // Memoize example prompts to prevent flickering
+  const examplePrompts = useMemo(() => {
+    const enabledTools = tools.filter((tool) => tool.enabled);
+
+    // Group tools by server (MCP tools have mcpServerId in their code)
+    const toolsByServer = enabledTools.reduce((acc, tool) => {
+      const mcpServerMatch = tool.code?.match(/mcpServerId: "([^"]+)"/);
+      const serverId = mcpServerMatch ? mcpServerMatch[1] : 'local';
+      if (!acc[serverId]) acc[serverId] = [];
+      acc[serverId].push(tool);
+      return acc;
+    }, {} as Record<string, typeof enabledTools>);
+
+    // Pick one tool from each server (up to 3 servers)
+    const serverIds = Object.keys(toolsByServer).slice(0, 3);
+    const selectedTools = serverIds.map(serverId => {
+      const serverTools = toolsByServer[serverId];
+      return serverTools[Math.floor(Math.random() * serverTools.length)];
+    });
+
+    return selectedTools.map((tool) => {
+      const schema = generateSchemaFromCode(tool.code);
+      const description = schema.description || tool.name;
+
+      // Create a cleaner natural language prompt
+      let displayText = description;
+      if (description !== tool.name) {
+        // If there's a description, make it conversational
+        displayText = description.charAt(0).toUpperCase() + description.slice(1);
+        if (!displayText.endsWith('?') && !displayText.endsWith('.')) {
+          displayText += '?';
+        }
+      } else {
+        // Fallback to tool name in a readable format
+        displayText = tool.name.replace(/_/g, ' ');
+        displayText = displayText.charAt(0).toUpperCase() + displayText.slice(1);
+      }
+
+      return {
+        icon: "🛠️",
+        displayText,
+        messageText: displayText,
+      };
+    });
+  }, [tools]);
+
   const loadTools = useCallback(async (): Promise<void> => {
     const db = await getDB();
     const allTools: Tool[] = await db.getAll(STORE_NAME);
@@ -405,10 +459,6 @@ const App: React.FC = () => {
     setIsGenerating(true);
 
     try {
-      const toolSchemas = tools
-        .filter((tool) => tool.enabled)
-        .map((tool) => generateSchemaFromCode(tool.code));
-
       while (true) {
         const messagesForGeneration = [
           { role: "system" as const, content: systemPrompt },
@@ -574,10 +624,6 @@ const App: React.FC = () => {
     setIsGenerating(true);
 
     try {
-      const toolSchemas = tools
-        .filter((tool) => tool.enabled)
-        .map((tool) => generateSchemaFromCode(tool.code));
-
       while (true) {
         const messagesForGeneration = [
           { role: "system" as const, content: systemPrompt },
@@ -716,50 +762,7 @@ const App: React.FC = () => {
             >
               {messages.length === 0 && isReady ? (
                 <ExamplePrompts
-                  examples={(() => {
-                    const enabledTools = tools.filter((tool) => tool.enabled);
-
-                    // Group tools by server (MCP tools have mcpServerId in their code)
-                    const toolsByServer = enabledTools.reduce((acc, tool) => {
-                      const mcpServerMatch = tool.code?.match(/mcpServerId: "([^"]+)"/);
-                      const serverId = mcpServerMatch ? mcpServerMatch[1] : 'local';
-                      if (!acc[serverId]) acc[serverId] = [];
-                      acc[serverId].push(tool);
-                      return acc;
-                    }, {} as Record<string, typeof enabledTools>);
-
-                    // Pick one tool from each server (up to 3 servers)
-                    const serverIds = Object.keys(toolsByServer).slice(0, 3);
-                    const selectedTools = serverIds.map(serverId => {
-                      const serverTools = toolsByServer[serverId];
-                      return serverTools[Math.floor(Math.random() * serverTools.length)];
-                    });
-
-                    return selectedTools.map((tool) => {
-                      const schema = generateSchemaFromCode(tool.code);
-                      const description = schema.description || tool.name;
-
-                      // Create a cleaner natural language prompt
-                      let displayText = description;
-                      if (description !== tool.name) {
-                        // If there's a description, make it conversational
-                        displayText = description.charAt(0).toUpperCase() + description.slice(1);
-                        if (!displayText.endsWith('?') && !displayText.endsWith('.')) {
-                          displayText += '?';
-                        }
-                      } else {
-                        // Fallback to tool name in a readable format
-                        displayText = tool.name.replace(/_/g, ' ');
-                        displayText = displayText.charAt(0).toUpperCase() + displayText.slice(1);
-                      }
-
-                      return {
-                        icon: "🛠️",
-                        displayText,
-                        messageText: displayText,
-                      };
-                    });
-                  })()}
+                  examples={examplePrompts}
                   onExampleClick={handleExampleClick}
                 />
               ) : (

src/components/ExamplePrompts.tsx

@@ -30,12 +30,18 @@ const ExamplePrompts: React.FC<ExamplePromptsProps> = ({
         <p className="text-sm text-gray-500">Click one to get started</p>
       </div>
 
-      <div className="grid grid-cols-1 sm:grid-cols-2 lg:grid-cols-3 gap-3 max-w-4xl w-full px-4">
+      <div className={`grid gap-3 max-w-4xl w-full px-4 ${
+        dynamicExamples.length === 1
+          ? 'grid-cols-1 justify-items-center'
+          : 'grid-cols-1 sm:grid-cols-2 lg:grid-cols-3'
+      }`}>
         {dynamicExamples.map((example, index) => (
           <button
             key={index}
             onClick={() => onExampleClick(example.messageText)}
-            className="flex items-start gap-3 p-4 bg-gray-700 hover:bg-gray-600 rounded-lg transition-colors text-left group cursor-pointer"
+            className={`flex items-start gap-3 p-4 bg-gray-700 hover:bg-gray-600 rounded-lg transition-colors text-left group cursor-pointer ${
+              dynamicExamples.length === 1 ? 'max-w-md' : ''
+            }`}
           >
             <span className="text-xl flex-shrink-0 group-hover:scale-110 transition-transform">
               {example.icon}

src/hooks/useLLM.ts

@@ -1,9 +1,4 @@
 import { useState, useEffect, useRef, useCallback } from "react";
-import {
-  AutoModelForCausalLM,
-  AutoTokenizer,
-  TextStreamer,
-} from "@huggingface/transformers";
 
 interface LLMState {
   isLoading: boolean;
@@ -14,18 +9,6 @@ interface LLMState {
   numTokens: number;
 }
 
-interface LLMInstance {
-  model: any;
-  tokenizer: any;
-}
-
-let moduleCache: {
-  [modelId: string]: {
-    instance: LLMInstance | null;
-    loadingPromise: Promise<LLMInstance> | null;
-  };
-} = {};
-
 export const useLLM = (modelId?: string) => {
   const [state, setState] = useState<LLMState>({
     isLoading: false,
@@ -36,54 +19,92 @@ export const useLLM = (modelId?: string) => {
     numTokens: 0,
   });
 
-  const instanceRef = useRef<LLMInstance | null>(null);
-  const loadingPromiseRef = useRef<Promise<LLMInstance> | null>(null);
+  const workerRef = useRef<Worker | null>(null);
+  const onTokenCallbackRef = useRef<((token: string) => void) | null>(null);
+  const resolveGenerationRef = useRef<((text: string) => void) | null>(null);
+  const rejectGenerationRef = useRef<((error: Error) => void) | null>(null);
 
-  const abortControllerRef = useRef<AbortController | null>(null);
-  const pastKeyValuesRef = useRef<any>(null);
-  const generationAbortControllerRef = useRef<AbortController | null>(null);
-
-  const loadModel = useCallback(async () => {
-    if (!modelId) {
-      throw new Error("Model ID is required");
-    }
-
-    const MODEL_ID = modelId;
-
-    if (!moduleCache[modelId]) {
-      moduleCache[modelId] = {
-        instance: null,
-        loadingPromise: null,
-      };
-    }
-
-    const cache = moduleCache[modelId];
-
-    const existingInstance = instanceRef.current || cache.instance;
-    if (existingInstance) {
-      instanceRef.current = existingInstance;
-      cache.instance = existingInstance;
-      setState((prev) => ({ ...prev, isReady: true, isLoading: false }));
-      return existingInstance;
-    }
+  // Initialize worker
+  useEffect(() => {
+    const worker = new Worker(
+      new URL("../workers/llm.worker.ts", import.meta.url),
+      { type: "module" }
+    );
+
+    workerRef.current = worker;
+
+    // Handle messages from worker
+    worker.onmessage = (e) => {
+      const message = e.data;
+
+      switch (message.type) {
+        case "progress":
+          setState((prev) => ({
+            ...prev,
+            progress: message.progress,
+            isLoading: true,
+          }));
+          break;
+
+        case "ready":
+          setState((prev) => ({
+            ...prev,
+            isLoading: false,
+            isReady: true,
+            progress: 100,
+          }));
+          break;
+
+        case "update":
+          setState((prev) => ({
+            ...prev,
+            tokensPerSecond: message.tokensPerSecond,
+            numTokens: message.numTokens,
+          }));
+          if (onTokenCallbackRef.current) {
+            onTokenCallbackRef.current(message.token);
+          }
+          break;
 
-    const existingPromise = loadingPromiseRef.current || cache.loadingPromise;
-    if (existingPromise) {
-      try {
-        const instance = await existingPromise;
-        instanceRef.current = instance;
-        cache.instance = instance;
-        setState((prev) => ({ ...prev, isReady: true, isLoading: false }));
-        return instance;
-      } catch (error) {
-        setState((prev) => ({
-          ...prev,
-          isLoading: false,
-          error:
-            error instanceof Error ? error.message : "Failed to load model",
-        }));
-        throw error;
+        case "complete":
+          if (resolveGenerationRef.current) {
+            resolveGenerationRef.current(message.text);
+            resolveGenerationRef.current = null;
+            rejectGenerationRef.current = null;
+          }
+          break;
+
+        case "error":
+          setState((prev) => ({
+            ...prev,
+            isLoading: false,
+            error: message.error,
+          }));
+          if (rejectGenerationRef.current) {
+            rejectGenerationRef.current(new Error(message.error));
+            resolveGenerationRef.current = null;
+            rejectGenerationRef.current = null;
+          }
+          break;
       }
+    };
+
+    worker.onerror = (error) => {
+      setState((prev) => ({
+        ...prev,
+        isLoading: false,
+        error: error.message,
+      }));
+    };
+
+    return () => {
+      worker.terminate();
+    };
+  }, []);
+
+  const loadModel = useCallback(async () => {
+    if (!modelId || !workerRef.current) {
+      throw new Error("Model ID or worker not available");
     }
 
     setState((prev) => ({
@@ -93,76 +114,10 @@ export const useLLM = (modelId?: string) => {
       progress: 0,
     }));
 
-    abortControllerRef.current = new AbortController();
-
-    const loadingPromise = (async () => {
-      try {
-        const progressCallback = (progress: any) => {
-          // Only update progress for weights
-          if (
-            progress.status === "progress" &&
-            progress.file.endsWith(".onnx_data")
-          ) {
-            const percentage = Math.round(
-              (progress.loaded / progress.total) * 100
-            );
-            setState((prev) => ({ ...prev, progress: percentage }));
-          }
-        };
-
-        const tokenizer = await AutoTokenizer.from_pretrained(MODEL_ID, {
-          progress_callback: progressCallback,
-        });
-
-        const model = await AutoModelForCausalLM.from_pretrained(MODEL_ID, {
-          dtype: "q4f16",
-          device: "webgpu",
-          progress_callback: progressCallback,
-        });
-
-        // Pre-warm the model with a dummy input for shader compilation
-        console.log("Pre-warming model...");
-        const dummyInput = tokenizer("Hello", {
-          return_tensors: "pt",
-          padding: false,
-          truncation: false,
-        });
-        await model.generate({
-          ...dummyInput,
-          max_new_tokens: 1,
-          do_sample: false,
-        });
-        console.log("Model pre-warmed");
-
-        const instance = { model, tokenizer };
-        instanceRef.current = instance;
-        cache.instance = instance;
-        loadingPromiseRef.current = null;
-        cache.loadingPromise = null;
-
-        setState((prev) => ({
-          ...prev,
-          isLoading: false,
-          isReady: true,
-          progress: 100,
-        }));
-        return instance;
-      } catch (error) {
-        loadingPromiseRef.current = null;
-        cache.loadingPromise = null;
-        setState((prev) => ({
-          ...prev,
-          isLoading: false,
-          error:
-            error instanceof Error ? error.message : "Failed to load model",
-        }));
-        throw error;
-      }
-    })();
-
-    loadingPromiseRef.current = loadingPromise;
-    cache.loadingPromise = loadingPromise;
-    return loadingPromise;
+    workerRef.current.postMessage({
+      type: "load",
+      modelId,
+    });
   }, [modelId]);
 
   const generateResponse = useCallback(
@@ -171,104 +126,44 @@ export const useLLM = (modelId?: string) => {
       tools: Array<any>,
       onToken?: (token: string) => void
     ): Promise<string> => {
-      const instance = instanceRef.current;
-      if (!instance) {
-        throw new Error("Model not loaded. Call loadModel() first.");
+      if (!workerRef.current) {
+        throw new Error("Worker not initialized");
       }
 
-      const { model, tokenizer } = instance;
-
-      // Create abort controller for this generation
-      generationAbortControllerRef.current = new AbortController();
-
-      // Apply chat template with tools
-      const input = tokenizer.apply_chat_template(messages, {
-        tools,
-        add_generation_prompt: true,
-        return_dict: true,
-      });
-
-      // Track tokens and timing
-      const startTime = performance.now();
-      let tokenCount = 0;
-
-      const streamer = onToken
-        ? new TextStreamer(tokenizer, {
-            skip_prompt: true,
-            skip_special_tokens: false,
-            callback_function: (token: string) => {
-              tokenCount++;
-              const elapsed = (performance.now() - startTime) / 1000;
-              const tps = tokenCount / elapsed;
-              setState((prev) => ({
-                ...prev,
-                tokensPerSecond: tps,
-                numTokens: tokenCount,
-              }));
-              onToken(token);
-            },
-          })
-        : undefined;
+      return new Promise((resolve, reject) => {
+        onTokenCallbackRef.current = onToken || null;
+        resolveGenerationRef.current = resolve;
+        rejectGenerationRef.current = reject;
 
-      try {
-        // Generate the response
-        const { sequences, past_key_values } = await model.generate({
-          ...input,
-          past_key_values: pastKeyValuesRef.current,
-          max_new_tokens: 1024,
-          do_sample: false,
-          streamer,
-          return_dict_in_generate: true,
+        workerRef.current!.postMessage({
+          type: "generate",
+          messages,
+          tools,
         });
-        pastKeyValuesRef.current = past_key_values;
-
-        // Decode the generated text with special tokens preserved (except end tokens) for tool call detection
-        const response = tokenizer
-          .batch_decode(sequences.slice(null, [input.input_ids.dims[1], null]), {
-            skip_special_tokens: false,
-          })[0]
-          .replace(/<\|im_end\|>$/, "")
-          .replace(/<\|end_of_text\|>$/, "");
-
-        return response;
-      } finally {
-        generationAbortControllerRef.current = null;
-      }
+      });
     },
     []
   );
 
   const interruptGeneration = useCallback(() => {
-    if (generationAbortControllerRef.current) {
-      generationAbortControllerRef.current.abort();
+    if (workerRef.current) {
+      workerRef.current.postMessage({ type: "interrupt" });
     }
   }, []);
 
   const clearPastKeyValues = useCallback(() => {
-    pastKeyValuesRef.current = null;
+    if (workerRef.current) {
+      workerRef.current.postMessage({ type: "reset" });
+    }
   }, []);
 
   const cleanup = useCallback(() => {
-    if (abortControllerRef.current) {
-      abortControllerRef.current.abort();
+    if (workerRef.current) {
+      workerRef.current.terminate();
+      workerRef.current = null;
     }
   }, []);
 
-  useEffect(() => {
-    return cleanup;
-  }, [cleanup]);
-
-  useEffect(() => {
-    if (modelId && moduleCache[modelId]) {
-      const existingInstance =
-        instanceRef.current || moduleCache[modelId].instance;
-      if (existingInstance) {
-        instanceRef.current = existingInstance;
-        setState((prev) => ({ ...prev, isReady: true }));
-      }
-    }
-  }, [modelId]);
-
   return {
     ...state,
     loadModel,

src/workers/llm.worker.ts

@@ -0,0 +1,253 @@
+import {
+  AutoModelForCausalLM,
+  AutoTokenizer,
+  TextStreamer,
+} from "@huggingface/transformers";
+
+// Worker state
+let model: any = null;
+let tokenizer: any = null;
+let pastKeyValues: any = null;
+let isGenerating = false;
+
+// Cache for loaded models
+const modelCache: {
+  [modelId: string]: {
+    model: any;
+    tokenizer: any;
+  };
+} = {};
+
+// Message types from main thread
+interface LoadMessage {
+  type: "load";
+  modelId: string;
+}
+
+interface GenerateMessage {
+  type: "generate";
+  messages: Array<{ role: string; content: string }>;
+  tools: Array<any>;
+}
+
+interface InterruptMessage {
+  type: "interrupt";
+}
+
+interface ResetMessage {
+  type: "reset";
+}
+
+type WorkerMessage = LoadMessage | GenerateMessage | InterruptMessage | ResetMessage;
+
+// Message types to main thread
+interface ProgressMessage {
+  type: "progress";
+  progress: number;
+  file?: string;
+}
+
+interface ReadyMessage {
+  type: "ready";
+}
+
+interface UpdateMessage {
+  type: "update";
+  token: string;
+  tokensPerSecond: number;
+  numTokens: number;
+}
+
+interface CompleteMessage {
+  type: "complete";
+  text: string;
+}
+
+interface ErrorMessage {
+  type: "error";
+  error: string;
+}
+
+type WorkerResponse = ProgressMessage | ReadyMessage | UpdateMessage | CompleteMessage | ErrorMessage;
+
+function postMessage(message: WorkerResponse) {
+  self.postMessage(message);
+}
+
+// Load model
+async function loadModel(modelId: string) {
+  try {
+    // Check cache first
+    if (modelCache[modelId]) {
+      model = modelCache[modelId].model;
+      tokenizer = modelCache[modelId].tokenizer;
+      postMessage({ type: "ready" });
+      return;
+    }
+
+    const progressCallback = (progress: any) => {
+      if (
+        progress.status === "progress" &&
+        progress.file.endsWith(".onnx_data")
+      ) {
+        const percentage = Math.round(
+          (progress.loaded / progress.total) * 100
+        );
+        postMessage({
+          type: "progress",
+          progress: percentage,
+          file: progress.file,
+        });
+      }
+    };
+
+    // Load tokenizer
+    tokenizer = await AutoTokenizer.from_pretrained(modelId, {
+      progress_callback: progressCallback,
+    });
+
+    // Load model
+    model = await AutoModelForCausalLM.from_pretrained(modelId, {
+      dtype: "q4f16",
+      device: "webgpu",
+      progress_callback: progressCallback,
+    });
+
+    // Pre-warm the model with a dummy input for shader compilation
+    const dummyInput = tokenizer("Hello", {
+      return_tensors: "pt",
+      padding: false,
+      truncation: false,
+    });
+    await model.generate({
+      ...dummyInput,
+      max_new_tokens: 1,
+      do_sample: false,
+    });
+
+    // Cache the loaded model
+    modelCache[modelId] = { model, tokenizer };
+
+    postMessage({ type: "ready" });
+  } catch (error) {
+    postMessage({
+      type: "error",
+      error: error instanceof Error ? error.message : "Failed to load model",
+    });
+  }
+}
+
+// Generate response
+async function generate(
+  messages: Array<{ role: string; content: string }>,
+  tools: Array<any>
+) {
+  if (!model || !tokenizer) {
+    postMessage({ type: "error", error: "Model not loaded" });
+    return;
+  }
+
+  try {
+    isGenerating = true;
+
+    // Apply chat template with tools
+    const input = tokenizer.apply_chat_template(messages, {
+      tools,
+      add_generation_prompt: true,
+      return_dict: true,
+    });
+
+    // Track tokens and timing
+    const startTime = performance.now();
+    let tokenCount = 0;
+
+    const streamer = new TextStreamer(tokenizer, {
+      skip_prompt: true,
+      skip_special_tokens: false,
+      callback_function: (token: string) => {
+        if (!isGenerating) return; // Check if interrupted
+
+        tokenCount++;
+        const elapsed = (performance.now() - startTime) / 1000;
+        const tps = tokenCount / elapsed;
+
+        postMessage({
+          type: "update",
+          token,
+          tokensPerSecond: tps,
+          numTokens: tokenCount,
+        });
+      },
+    });
+
+    // Generate the response
+    const { sequences, past_key_values } = await model.generate({
+      ...input,
+      past_key_values: pastKeyValues,
+      max_new_tokens: 1024,
+      do_sample: false,
+      streamer,
+      return_dict_in_generate: true,
+    });
+
+    pastKeyValues = past_key_values;
+
+    // Decode the generated text
+    const response = tokenizer
+      .batch_decode(sequences.slice(null, [input.input_ids.dims[1], null]), {
+        skip_special_tokens: false,
+      })[0]
+      .replace(/<\|im_end\|>$/, "")
+      .replace(/<\|end_of_text\|>$/, "");
+
+    if (isGenerating) {
+      postMessage({ type: "complete", text: response });
+    }
+
+    isGenerating = false;
+  } catch (error) {
+    isGenerating = false;
+    postMessage({
+      type: "error",
+      error: error instanceof Error ? error.message : "Generation failed",
+    });
+  }
+}
+
+// Interrupt generation
+function interrupt() {
+  isGenerating = false;
+  // Send a completion message with empty text to resolve the promise
+  postMessage({ type: "complete", text: "" });
+}
+
+// Reset past key values
+function reset() {
+  pastKeyValues = null;
+}
+
+// Handle messages from main thread
+self.onmessage = async (e: MessageEvent<WorkerMessage>) => {
+  const message = e.data;
+
+  switch (message.type) {
+    case "load":
+      await loadModel(message.modelId);
+      break;
+
+    case "generate":
+      await generate(message.messages, message.tools);
+      break;
+
+    case "interrupt":
+      interrupt();
+      break;
+
+    case "reset":
+      reset();
+      break;
+  }
+};
+
+// Export for TypeScript
+export {};