server/chatStream.js

import { LIGHTNING_BASE } from "./config.js";

const SYSTEM_PROMPT =
  "CRITICAL RULE: Every response MUST use HTML <span data-color=\"{COLOR NAME}\"> tags to color main points and headings. " +
  "COLORS MUST HAVE MEANING AND CONSISTENCY ACROSS THE ENTIRE CONVERSATION. " +
  "You may ONLY use the following semantic color names: green, pink, blue, red, orange, yellow, purple, teal, gold, coral. " +
  "Never output text formatted with explicit black or white colors. " +
  "Use a variety of colors throughout every response to distinguish headings, sections, and key terms. " +
  "Keep code blocks plain, but color headings and important points in surrounding text. " +
  "Do not over-color responses. Use color intentionally and sparingly. " +
  "CRITICAL RULE: MARKDOWN FORMATTING SUCH AS #, ##, ###, **, * MUST BE PLACED OUTSIDE <span> tags. " +
  "You are a helpful, friendly AI assistant. Use tools when appropriate to help the user. " +
  "When generating media, do not include URLs — it is displayed automatically. " +
  "You can render SVG images by outputting SVG code in a code block tagged exactly as:\n```svg\n<svg>...</svg>\n```\n" +
  "Never use single backslashes. You may use emojis where appropriate. " +
  "Use markdown for everything other than coloring your text. Use tables, lists, and other markdown elements.";

export async function streamChat(ws, {
  sessionId,
  model,
  history,
  userMessage,
  tools,
  accessToken,
  clientId,
  onToken,
  onDone,
  onError,
  onToolCall,
  onNewAsset,
  abortSignal,
}) {
  const headers = {
    "Content-Type": "application/json",
    "Accept": "text/event-stream",
  };
  if (accessToken) headers["Authorization"] = `Bearer ${accessToken}`;
  if (clientId) headers["X-Client-ID"] = clientId;

  const messages = [
    { role: "system", content: SYSTEM_PROMPT },
    ...history.map(normalizeMessage),
    { role: "user", content: userMessage },
  ];

  const enabledTools = buildToolList(tools);

  try {
    const response = await fetch(`${LIGHTNING_BASE}/gen/chat/completions`, {
      method: "POST",
      headers,
      body: JSON.stringify({
        model: model || "lightning",
        messages,
        tools: enabledTools.length > 0 ? enabledTools : undefined,
        stream: true,
      }),
      signal: abortSignal,
    });

    if (!response.ok) {
      const err = await response.text();
      onError(err);
      return;
    }

    const reader = response.body.getReader();
    const decoder = new TextDecoder();
    let assistantText = "";
    const toolCallBuffer = new Map();
    let buffer = "";

    while (true) {
      const { done, value } = await reader.read();
      if (done) break;

      buffer += decoder.decode(value, { stream: true });
      const lines = buffer.split("\n");
      buffer = lines.pop() || "";

      for (const line of lines) {
        if (!line.startsWith("data: ")) continue;
        const data = line.slice(6).trim();
        if (data === "[DONE]") continue;

        let chunk;
        try { chunk = JSON.parse(data); } catch { continue; }

        const delta = chunk.choices?.[0]?.delta;
        if (!delta) continue;

        if (delta.content) {
          assistantText += delta.content;
          onToken(delta.content);
        }

        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.function?.arguments) entry.arguments += call.function.arguments;
            toolCallBuffer.set(call.index, entry);
          }
        }
      }
    }

    // Process tool calls
    const toolCalls = [...toolCallBuffer.values()].map(t => ({
      id: t.id || `call_${crypto.randomUUID()}`,
      type: "function",
      function: { name: t.name, arguments: t.arguments },
    }));

    if (toolCalls.length > 0) {
      const toolResults = await processToolCalls(ws, toolCalls, tools, accessToken, clientId, abortSignal, onToolCall, onNewAsset);

      // Follow-up response after tool calls
      const followUpMessages = [
        { role: "system", content: SYSTEM_PROMPT },
        ...history.map(normalizeMessage),
        { role: "user", content: userMessage },
        { role: "assistant", content: assistantText || "", tool_calls: toolCalls },
        ...toolResults,
      ];

      const followUp = await fetch(`${LIGHTNING_BASE}/gen/chat/completions`, {
        method: "POST",
        headers: { ...headers, "Accept": "text/event-stream" },
        body: JSON.stringify({
          model: model || "lightning",
          messages: followUpMessages,
          stream: true,
        }),
        signal: abortSignal,
      });

      if (followUp.ok) {
        const fuReader = followUp.body.getReader();
        let fuBuffer = "";
        while (true) {
          const { done, value } = await fuReader.read();
          if (done) break;
          fuBuffer += decoder.decode(value, { stream: true });
          const fuLines = fuBuffer.split("\n");
          fuBuffer = fuLines.pop() || "";
          for (const line of fuLines) {
            if (!line.startsWith("data: ")) continue;
            const data = line.slice(6).trim();
            if (data === "[DONE]") continue;
            let chunk;
            try { chunk = JSON.parse(data); } catch { continue; }
            const delta = chunk.choices?.[0]?.delta;
            if (delta?.content) {
              assistantText += delta.content;
              onToken(delta.content);
            }
          }
        }
      }
    }

    onDone(assistantText, toolCalls);
  } catch (err) {
    if (err.name === "AbortError") {
      onDone(null, null, true); // aborted
    } else {
      onError(String(err));
    }
  }
}

function normalizeMessage(msg) {
  if (msg.role === "tool") {
    return { role: "tool", tool_call_id: msg.tool_call_id, content: String(msg.content) };
  }
  if (msg.role === "assistant" && msg.tool_calls) {
    return { role: "assistant", content: "", tool_calls: msg.tool_calls };
  }
  return { role: msg.role, content: msg.content };
}

function buildToolList(tools) {
  if (!tools) return [];
  const list = [];
  if (tools.webSearch) {
    list.push({
      type: "function",
      function: {
        name: "ollama_search",
        description: "Search the web for current information",
        parameters: {
          type: "object",
          properties: { query: { type: "string", description: "Search query" } },
          required: ["query"],
        },
      },
    });
    list.push({
      type: "function",
      function: {
        name: "read_web_page",
        description: "Read the content of a web page by URL",
        parameters: {
          type: "object",
          properties: { url: { type: "string", description: "URL to fetch" } },
          required: ["url"],
        },
      },
    });
  }
  if (tools.imageGen) {
    list.push({
      type: "function",
      function: {
        name: "generate_image",
        description: "Generate an image from a prompt",
        parameters: {
          type: "object",
          properties: {
            prompt: { type: "string" },
            mode: { type: "string", enum: ["auto", "fantasy", "realistic"] },
            image_urls: { type: "array", items: { type: "string" } },
          },
          required: ["prompt"],
        },
      },
    });
  }
  if (tools.videoGen) {
    list.push({
      type: "function",
      function: {
        name: "generate_video",
        description: "Generate a video from a prompt",
        parameters: {
          type: "object",
          properties: {
            prompt: { type: "string" },
            ratio: { type: "string", enum: ["3:2", "2:3", "1:1"] },
            mode: { type: "string", enum: ["normal", "fun"] },
            duration: { type: "number" },
            image_urls: { type: "array", items: { type: "string" } },
          },
          required: ["prompt"],
        },
      },
    });
  }
  if (tools.audioGen) {
    list.push({
      type: "function",
      function: {
        name: "generate_audio",
        description: "Generate music or sound effects from a prompt",
        parameters: {
          type: "object",
          properties: { prompt: { type: "string" } },
          required: ["prompt"],
        },
      },
    });
  }
  return list;
}

async function processToolCalls(ws, toolCalls, tools, accessToken, clientId, abortSignal, onToolCall, onNewAsset) {
  const toolResults = [];
  const authHeaders = {};
  if (accessToken) authHeaders["Authorization"] = `Bearer ${accessToken}`;
  if (clientId) authHeaders["X-Client-ID"] = clientId;

  for (const call of toolCalls) {
    let args;
    try { args = JSON.parse(call.function.arguments || "{}"); } catch { args = {}; }

    onToolCall({ id: call.id, name: call.function.name, state: "pending", args });

    let result = "Tool completed.";

    try {
      if (call.function.name === "ollama_search") {
        const { Client } = await import("@gradio/client");
        const client = await Client.connect("incognitolm/Web-Search");
        const r = await client.predict("/perform_search", { query: args.query });
        result = JSON.stringify(r.data);
      }

      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),
          });
        }
      }

      else if (call.function.name === "generate_image") {
        const body = { prompt: args.prompt };
        if (args.mode) body.mode = args.mode;
        if (args.image_urls?.length) body.image_urls = args.image_urls;

        const res = await fetch(`${LIGHTNING_BASE}/gen/image`, {
          method: "POST",
          headers: { "Content-Type": "application/json", ...authHeaders },
          body: JSON.stringify(body),
          signal: abortSignal,
        });
        if (res.ok) {
          const buf = await res.arrayBuffer();
          const ct = res.headers.get("content-type") || "image/png";
          const b64 = Buffer.from(buf).toString("base64");
          const dataUrl = `data:${ct};base64,${b64}`;
          onNewAsset({ role: "image", content: dataUrl });
          result = "Image generated successfully and shown to the user.";
        } else {
          result = `Image generation failed: ${res.status}`;
        }
      }

      else if (call.function.name === "generate_video") {
        const body = { prompt: args.prompt };
        if (args.ratio) body.ratio = args.ratio;
        if (args.mode) body.mode = args.mode;
        if (args.duration) body.duration = args.duration;
        if (args.image_urls?.length) body.image_urls = args.image_urls;

        const res = await fetch(`${LIGHTNING_BASE}/gen/video`, {
          method: "POST",
          headers: { "Content-Type": "application/json", ...authHeaders },
          body: JSON.stringify(body),
          signal: abortSignal,
        });
        if (res.ok) {
          const buf = await res.arrayBuffer();
          const b64 = Buffer.from(buf).toString("base64");
          const dataUrl = `data:video/mp4;base64,${b64}`;
          onNewAsset({ role: "video", content: dataUrl });
          result = "Video generated successfully and shown to the user.";
        } else {
          result = `Video generation failed: ${res.status}`;
        }
      }

      else if (call.function.name === "generate_audio") {
        const res = await fetch(`${LIGHTNING_BASE}/gen/sfx`, {
          method: "POST",
          headers: { "Content-Type": "application/json", ...authHeaders },
          body: JSON.stringify({ prompt: args.prompt }),
          signal: abortSignal,
        });
        if (res.ok) {
          const buf = await res.arrayBuffer();
          const b64 = Buffer.from(buf).toString("base64");
          const dataUrl = `data:audio/mpeg;base64,${b64}`;
          onNewAsset({ role: "audio", content: dataUrl });
          result = "Audio generated successfully and shown to the user.";
        } else {
          result = `Audio generation failed: ${res.status}`;
        }
      }
    } catch (err) {
      console.log(`Tool error: ${String(err)}`);
      result = `Tool error: ${String(err)}`;
    }

    onToolCall({ id: call.id, name: call.function.name, state: "resolved", result });

    toolResults.push({
      role: "tool",
      tool_call_id: call.id,
      content: typeof result === "string" ? result : JSON.stringify(result),
    });
  }

  return toolResults;
}