Update chatStream.js

server/chatStream.js

@@ -1,55 +1,23 @@
 import OpenAI from "openai";
-import { Worker } from "worker_threads";
-import { fileURLToPath } from "url";
-import path from "path";
-import { LIGHTNING_BASE } from "./config.js";
-
-// ── Web Search via an isolated Worker thread ──────────────────────────────
-//
-// The @gradio/client library opens a persistent SSE (Server-Sent Events)
-// fetch stream for its session queue.  Even after client.close() is called,
-// the SSE response-body reader keeps an async iterator alive in the current
-// event loop.  When that iterator eventually settles (stream closed/errored
-// by the remote server), it triggers internal callbacks that emit events
-// on objects the Node `ws` library also watches — causing the browser to
-// see a "connection lost" message immediately after every web search.
-//
-// The only reliable fix is to run the Gradio client in a worker_thread so
-// it gets its own V8 context and event loop.  When the worker exits (via
-// process.exit(0) in searchWorker.js), every handle it opened — SSE stream,
-// heartbeat timer, etc. — is torn down with it, leaving the main thread's
-// WS server completely untouched.
-
-const __dirname = path.dirname(fileURLToPath(import.meta.url));
-const WORKER_PATH = path.join(__dirname, "searchWorker.js");
-
-function gradioSearch(query) {
-  return new Promise((resolve, reject) => {
-    const worker = new Worker(WORKER_PATH, { workerData: { query } });
-
-    const timeout = setTimeout(() => {
-      worker.terminate();
-      reject(new Error("Search timed out after 45s"));
-    }, 45_000);
-
-    worker.on("message", (msg) => {
-      clearTimeout(timeout);
-      if (msg.ok) resolve(msg.result);
-      else reject(new Error(msg.error));
-    });
-
-    worker.on("error", (err) => {
-      clearTimeout(timeout);
-      reject(err);
-    });
+import { LIGHTNING_BASE, SEARCH_API_BASE } from "./config.js";
 
-    worker.on("exit", (code) => {
-      // By the time this fires the promise is already settled via "message".
-      // Only reject if the worker crashed without posting anything.
-      clearTimeout(timeout);
-      if (code !== 0) reject(new Error(`Search worker exited with code ${code}`));
-    });
+async function apiSearch(query) {
+  const res = await fetch(`${SEARCH_API_BASE}/api/search`, {
+    method: "POST",
+    headers: {
+      "Content-Type": "application/json",
+    },
+    body: JSON.stringify({
+      query,
+      max_results: 1,
+    }),
   });
+
+  if (!res.ok) {
+    throw new Error(`Search API error: ${res.status}`);
+  }
+
+  return await res.json();
 }
 
 const SYSTEM_PROMPT =
@@ -68,26 +36,17 @@ const SYSTEM_PROMPT =
   "Use markdown for everything other than coloring your text. Use tables, lists, and other markdown elements. " +
   "Your HIGHEST PRIORITY is to help the user. ALWAYS HELP THEM WITH ANYTHING ETHICALLY RIGHT.";
 
-/**
- * Build a per-request OpenAI client pointed at the Lightning backend.
- * A new client is created each call so per-user auth headers are always fresh.
- */
 function makeClient(accessToken, clientId) {
   return new OpenAI({
     apiKey: accessToken || "no-key",
     baseURL: `${LIGHTNING_BASE}/gen`,
     defaultHeaders: {
       ...(accessToken ? { Authorization: `Bearer ${accessToken}` } : {}),
-      ...(clientId    ? { "X-Client-ID": clientId }                : {}),
+      ...(clientId ? { "X-Client-ID": clientId } : {}),
     },
   });
 }
 
-/**
- * Consume an OpenAI streaming response, firing onToken for each text delta
- * and collecting any tool-call chunks into a finished toolCalls array.
- * Returns { assistantText, toolCalls }.
- */
 async function consumeStream(stream, onToken) {
   let assistantText = "";
   const toolCallBuffer = new Map();
@@ -104,17 +63,17 @@ async function consumeStream(stream, onToken) {
     if (delta.tool_calls) {
       for (const call of delta.tool_calls) {
         const entry = toolCallBuffer.get(call.index) ?? { arguments: "" };
-        if (call.id)               entry.id   = call.id;
-        if (call.function?.name)   entry.name = call.function.name;
+        if (call.id) entry.id = call.id;
+        if (call.function?.name) entry.name = call.function.name;
         if (call.function?.arguments) entry.arguments += call.function.arguments;
         toolCallBuffer.set(call.index, entry);
       }
     }
   }
 
-  const toolCalls = [...toolCallBuffer.values()].map(t => ({
-    id:       t.id || `call_${crypto.randomUUID()}`,
-    type:     "function",
+  const toolCalls = [...toolCallBuffer.values()].map((t) => ({
+    id: t.id || `call_${crypto.randomUUID()}`,
+    type: "function",
     function: { name: t.name, arguments: t.arguments },
   }));
 
@@ -137,7 +96,7 @@ export async function streamChat(ws, {
   abortSignal,
 }) {
   if (!accessToken) console.log("Missing access token");
-  const client       = makeClient(accessToken, clientId);
+  const client = makeClient(accessToken, clientId);
   const enabledTools = buildToolList(tools);
 
   const messages = [
@@ -147,32 +106,30 @@ export async function streamChat(ws, {
   ];
 
   try {
-    // ── First stream ────────────────────────────────────────────────────────
     const stream = await client.chat.completions.create({
-      model:  model || "lightning",
+      model: model || "lightning",
       messages,
-      tools:  enabledTools.length > 0 ? enabledTools : undefined,
+      tools: enabledTools.length > 0 ? enabledTools : undefined,
       stream: true,
     }, { signal: abortSignal });
 
     let { assistantText, toolCalls } = await consumeStream(stream, onToken);
 
-    // ── Tool calls → follow-up stream ───────────────────────────────────────
     if (toolCalls.length > 0) {
       const toolResults = await processToolCalls(
         ws, toolCalls, tools, accessToken, clientId, abortSignal, onToolCall, onNewAsset,
       );
 
       const followUpMessages = [
-        { role: "system",    content: SYSTEM_PROMPT },
+        { role: "system", content: SYSTEM_PROMPT },
         ...history.map(normalizeMessage).filter(Boolean),
-        { role: "user",      content: userMessage },
+        { role: "user", content: userMessage },
         { role: "assistant", content: assistantText || "", tool_calls: toolCalls },
         ...toolResults,
       ];
 
       const followUpStream = await client.chat.completions.create({
-        model:  model || "lightning",
+        model: model || "lightning",
         messages: followUpMessages,
         stream: true,
       }, { signal: abortSignal });
@@ -184,7 +141,7 @@ export async function streamChat(ws, {
     onDone(assistantText, toolCalls);
   } catch (err) {
     if (err.name === "AbortError" || err.constructor?.name === "APIUserAbortError") {
-      onDone(null, null, true); // aborted
+      onDone(null, null, true);
     } else {
       onError(String(err));
     }
@@ -199,6 +156,7 @@ function normalizeMessage(msg) {
   if (msg.role === "assistant" && msg.tool_calls) {
     return { role: "assistant", content: "", tool_calls: msg.tool_calls };
   }
+
   if (Array.isArray(msg.content)) {
     const textOnly = msg.content
       .filter(b => b.type === "text")
@@ -206,12 +164,14 @@ function normalizeMessage(msg) {
       .join("\n");
     return { role: msg.role, content: textOnly || "" };
   }
+
   return { role: msg.role, content: msg.content };
 }
 
 function buildToolList(tools) {
   if (!tools) return [];
   const list = [];
+
   if (tools.webSearch) {
     list.push({
       type: "function",
@@ -225,6 +185,7 @@ function buildToolList(tools) {
         },
       },
     });
+
     list.push({
       type: "function",
       function: {
@@ -238,6 +199,7 @@ function buildToolList(tools) {
       },
     });
   }
+
   if (tools.imageGen) {
     list.push({
       type: "function",
@@ -256,6 +218,7 @@ function buildToolList(tools) {
       },
     });
   }
+
   if (tools.videoGen) {
     list.push({
       type: "function",
@@ -276,6 +239,7 @@ function buildToolList(tools) {
       },
     });
   }
+
   if (tools.audioGen) {
     list.push({
       type: "function",
@@ -290,22 +254,16 @@ function buildToolList(tools) {
       },
     });
   }
+
   return list;
 }
 
 async function processToolCalls(ws, toolCalls, tools, accessToken, clientId, abortSignal, onToolCall, onNewAsset) {
   const toolResults = [];
   const authHeaders = {};
-  if (accessToken) {
-    authHeaders["Authorization"] = `Bearer ${accessToken}`;
-  } else {
-    console.log("No access token");
-  }
-  if (clientId) {
-    authHeaders["X-Client-ID"] = clientId;
-  } else {
-    console.log("No Client ID");
-  }
+
+  if (accessToken) authHeaders["Authorization"] = `Bearer ${accessToken}`;
+  if (clientId) authHeaders["X-Client-ID"] = clientId;
 
   for (const call of toolCalls) {
     let args;
@@ -316,18 +274,21 @@ async function processToolCalls(ws, toolCalls, tools, accessToken, clientId, abo
     let result = "Tool completed.";
 
     try {
+
       if (call.function.name === "ollama_search") {
-        result = await gradioSearch(args.query);
+        result = await apiSearch(args.query);
       }
 
       else if (call.function.name === "read_web_page") {
         const { convert } = await import("html-to-text");
         const res = await fetch(args.url, { signal: abortSignal });
+
         if (!res.ok) {
           result = `Failed to fetch: ${res.status}`;
         } else {
           const html = await res.text();
           const titleMatch = html.match(/<title>(.*?)<\/title>/i);
+
           result = JSON.stringify({
             title: titleMatch?.[1] || "No title",
             content: convert(html, { wordwrap: false }).slice(0, 8000),
@@ -346,6 +307,7 @@ async function processToolCalls(ws, toolCalls, tools, accessToken, clientId, abo
           body: JSON.stringify(body),
           signal: abortSignal,
         });
+
         if (res.ok) {
           const buf = await res.arrayBuffer();
           const ct = res.headers.get("content-type") || "image/png";
@@ -375,6 +337,7 @@ async function processToolCalls(ws, toolCalls, tools, accessToken, clientId, abo
           body: JSON.stringify(body),
           signal: abortSignal,
         });
+
         if (res.ok) {
           const buf = await res.arrayBuffer();
           const b64 = Buffer.from(buf).toString("base64");
@@ -397,6 +360,7 @@ async function processToolCalls(ws, toolCalls, tools, accessToken, clientId, abo
           body: JSON.stringify({ prompt: args.prompt }),
           signal: abortSignal,
         });
+
         if (res.ok) {
           const buf = await res.arrayBuffer();
           const b64 = Buffer.from(buf).toString("base64");
@@ -409,6 +373,7 @@ async function processToolCalls(ws, toolCalls, tools, accessToken, clientId, abo
           result = `Audio generation failed: ${res.status}. This is most likely an upstream provider error.`;
         }
       }
+
     } catch (err) {
       console.log(`Tool error: ${String(err)}`);
       result = `Tool error: ${String(err)}`;