chore: tool calling (wip)

Author: pi0

demo/llm.ts

@@ -1,6 +1,45 @@
 import type { MLCEngine } from "@mlc-ai/web-llm";
 
 // https://github.com/mlc-ai/web-llm
+// Tool use via structural tags
+// Ref: https://github.com/mlc-ai/web-llm/blob/main/examples/structural-tag-tool-use
+
+type ToolInvocation = { name: string; arguments: Record<string, unknown> };
+
+const tools = [
+  {
+    name: "get_current_time",
+    description: "Return the current date and time in a given timezone.",
+    schema: {
+      type: "object",
+      properties: {
+        timezone: {
+          type: "string",
+          description: "IANA timezone name, defaults to UTC",
+        },
+      },
+      required: [] as string[],
+    },
+  },
+];
+
+const toolResponseFormat = {
+  type: "structural_tag",
+  structural_tag: {
+    type: "structural_tag",
+    format: {
+      type: "triggered_tags",
+      triggers: ["<tool_call>"],
+      tags: tools.map((t) => ({
+        begin: `<tool_call>\n{"name": "${t.name}", "arguments": `,
+        content: { type: "json_schema", json_schema: t.schema },
+        end: "}\n</tool_call>",
+      })),
+      at_least_one: false,
+      stop_after_first: false,
+    },
+  },
+};
 
 let engine: MLCEngine | null = null;
 let chatHistory: { role: string; content: string }[] = [];
@@ -18,10 +57,17 @@ export function isReady() {
 }
 
 export function setAgent(name: string) {
+  const toolList = tools
+    .map((t) => `- ${t.name}: ${t.description}`)
+    .join("\n");
   chatHistory = [
     {
       role: "system",
-      content: `You are ${name}, a helpful desktop assistant from Windows 98. Keep responses very short and fun (1 sentence max).`,
+      content: [
+        `You are ${name}, a helpful desktop assistant from Windows 98. Keep responses very short and fun (1 sentence max).`,
+        `You have tools available. To use one, emit a <tool_call> block.`,
+        toolList,
+      ].join("\n"),
     },
   ];
 }
@@ -69,7 +115,9 @@ chatBtn.addEventListener("click", async () => {
 
 let onReplyStream: ((stream: AsyncIterable<string>) => void) | null = null;
 
-export function onAgentReplyStream(cb: (stream: AsyncIterable<string>) => void) {
+export function onAgentReplyStream(
+  cb: (stream: AsyncIterable<string>) => void,
+) {
   onReplyStream = cb;
 }
 
@@ -85,34 +133,44 @@ async function sendChat() {
   chatHistory.push({ role: "user", content: text });
 
   try {
-    const chunks = await engine.chat.completions.create({
+    // First pass (non-streaming) to detect tool calls via structural tags
+    const firstReply = await engine.chat.completions.create({
       messages: chatHistory as any,
-      stream: true,
+      stream: false,
+      max_tokens: 512,
+      response_format: toolResponseFormat as any,
     });
 
-    let reply = "";
-    const msgEl = appendChatMsg("Agent", "", "assistant");
-
-    async function* deltaStream() {
-      for await (const chunk of chunks) {
-        const delta = chunk.choices[0]?.delta?.content || "";
-        if (!delta) continue;
-        reply += delta;
-        msgEl.querySelector(".chat-msg-text").textContent = reply;
-        chatMessages.scrollTop = chatMessages.scrollHeight;
-        yield delta;
-      }
-    }
-
-    const stream = deltaStream();
-    if (onReplyStream) {
-      onReplyStream(stream);
+    console.log("First pass reply:", firstReply);
+
+    const content = firstReply.choices[0]?.message?.content || "";
+    const calls = parseToolCalls(content);
+
+    if (calls.length > 0) {
+      // Execute tools and follow up with a streamed response
+      chatHistory.push({ role: "assistant", content });
+      const results = calls.map((c) => ({
+        tool: c.name,
+        result: executeTool(c),
+      }));
+      chatHistory.push({
+        role: "user",
+        content: `[Tool results]: ${JSON.stringify(results)}`,
+      });
+      const msgEl = appendChatMsg("Agent", "", "assistant");
+      await streamReply(msgEl);
     } else {
-      for await (const _ of stream) {
+      // No tool calls — display response directly
+      appendChatMsg("Agent", content, "assistant");
+      chatHistory.push({ role: "assistant", content });
+      if (onReplyStream) {
+        onReplyStream(
+          (async function* () {
+            yield content;
+          })(),
+        );
       }
     }
-
-    chatHistory.push({ role: "assistant", content: reply });
   } catch (err) {
     appendChatMsg("System", "Error: " + (err as Error).message, "user");
   }
@@ -122,24 +180,15 @@ async function sendChat() {
   chatInput.focus();
 }
 
-function appendChatMsg(sender: string, text: string, role: string) {
-  const div = document.createElement("div");
-  div.className = `chat-msg chat-msg-${role}`;
-  div.innerHTML = `<b>${sender}:</b> <span class="chat-msg-text"></span>`;
-  div.querySelector(".chat-msg-text").textContent = text;
-  chatMessages.appendChild(div);
-  chatMessages.scrollTop = chatMessages.scrollHeight;
-  return div;
-}
-
 chatSend.addEventListener("click", sendChat);
 chatInput.addEventListener("keydown", (e) => {
   if (e.key === "Enter") sendChat();
 });
 
 // Speech-to-text (optional, Chrome/Edge)
 const SpeechRecognition =
-  (globalThis as any).SpeechRecognition || (globalThis as any).webkitSpeechRecognition;
+  (globalThis as any).SpeechRecognition ||
+  (globalThis as any).webkitSpeechRecognition;
 
 if (SpeechRecognition) {
   chatMic.style.display = "";
@@ -207,3 +256,78 @@ if (SpeechRecognition) {
     setInterval(pollTTS, 200);
   }
 }
+
+// --- Internal helpers ---
+
+function appendChatMsg(sender: string, text: string, role: string) {
+  const div = document.createElement("div");
+  div.className = `chat-msg chat-msg-${role}`;
+  div.innerHTML = `<b>${sender}:</b> <span class="chat-msg-text"></span>`;
+  div.querySelector(".chat-msg-text").textContent = text;
+  chatMessages.appendChild(div);
+  chatMessages.scrollTop = chatMessages.scrollHeight;
+  return div;
+}
+
+async function streamReply(msgEl: HTMLElement) {
+  const chunks = await engine!.chat.completions.create({
+    messages: chatHistory as any,
+    stream: true,
+  });
+
+  let reply = "";
+
+  async function* deltaStream() {
+    for await (const chunk of chunks) {
+      const delta = chunk.choices[0]?.delta?.content || "";
+      if (!delta) continue;
+      reply += delta;
+      msgEl.querySelector(".chat-msg-text").textContent = reply;
+      chatMessages.scrollTop = chatMessages.scrollHeight;
+      yield delta;
+    }
+  }
+
+  const stream = deltaStream();
+  if (onReplyStream) {
+    onReplyStream(stream);
+  } else {
+    for await (const _ of stream) {
+      /* drain */
+    }
+  }
+
+  chatHistory.push({ role: "assistant", content: reply });
+}
+
+function parseToolCalls(content: string): ToolInvocation[] {
+  const regex = /<tool_call>\s*(\{[\s\S]*?\})\s*<\/tool_call>/g;
+  const calls: ToolInvocation[] = [];
+  let match: RegExpExecArray | null;
+  while ((match = regex.exec(content)) !== null) {
+    try {
+      const payload = JSON.parse(match[1]);
+      if (typeof payload.name === "string" && payload.arguments !== undefined) {
+        calls.push({ name: payload.name, arguments: payload.arguments });
+      }
+    } catch {
+      // skip malformed tool calls
+    }
+  }
+  return calls;
+}
+
+function executeTool(call: ToolInvocation): Record<string, unknown> {
+  if (call.name === "get_current_time") {
+    const timezone = String(call.arguments.timezone || "UTC");
+    try {
+      return {
+        timezone,
+        time: new Date().toLocaleString("en-US", { timeZone: timezone }),
+      };
+    } catch {
+      return { timezone: "UTC", time: new Date().toISOString() };
+    }
+  }
+  return { error: `Unknown tool: ${call.name}` };
+}