@assistant-ui/react 0.4.4 → 0.4.5

package/dist/edge.mjs

@@ -34,6 +34,15 @@ function assistantEncoderStream() {
         }
         case "tool-call":
           break;
+        case "tool-result": {
+          controller.enqueue(
+            formatStreamPart("3" /* ToolCallResult */, {
+              id: chunk.toolCallId,
+              result: chunk.result
+            })
+          );
+          break;
+        }
         case "finish": {
           const { type, ...rest } = chunk;
           controller.enqueue(
@@ -177,21 +186,658 @@ function toLanguageModelMessages(message) {
 }
 
 // src/runtimes/edge/createEdgeRuntimeAPI.ts
-var createEdgeRuntimeAPI = ({ model }) => {
+import { z as z3 } from "zod";
+
+// src/runtimes/edge/converters/toLanguageModelTools.ts
+import { z } from "zod";
+import zodToJsonSchema from "zod-to-json-schema";
+var toLanguageModelTools = (tools) => {
+  if (!tools) return [];
+  return Object.entries(tools).map(([name, tool]) => ({
+    type: "function",
+    name,
+    ...tool.description ? { description: tool.description } : void 0,
+    parameters: tool.parameters instanceof z.ZodType ? zodToJsonSchema(tool.parameters) : tool.parameters
+  }));
+};
+
+// src/runtimes/edge/streams/toolResultStream.ts
+import { z as z2 } from "zod";
+import sjson from "secure-json-parse";
+function toolResultStream(tools) {
+  const toolCallExecutions = /* @__PURE__ */ new Map();
+  return new TransformStream({
+    transform(chunk, controller) {
+      controller.enqueue(chunk);
+      const chunkType = chunk.type;
+      switch (chunkType) {
+        case "tool-call": {
+          const { toolCallId, toolCallType, toolName, args: argsText } = chunk;
+          const tool = tools?.[toolName];
+          if (!tool || !tool.execute) return;
+          const args = sjson.parse(argsText);
+          if (tool.parameters instanceof z2.ZodType) {
+            const result = tool.parameters.safeParse(args);
+            if (!result.success) {
+              controller.enqueue({
+                type: "error",
+                error: new Error("Invalid tool call arguments")
+              });
+              return;
+            } else {
+              toolCallExecutions.set(
+                toolCallId,
+                (async () => {
+                  try {
+                    const result2 = await tool.execute(args);
+                    controller.enqueue({
+                      type: "tool-result",
+                      toolCallType,
+                      toolCallId,
+                      toolName,
+                      result: result2
+                    });
+                  } catch (error) {
+                    controller.enqueue({
+                      type: "error",
+                      error
+                    });
+                  } finally {
+                    toolCallExecutions.delete(toolCallId);
+                  }
+                })()
+              );
+            }
+          }
+          break;
+        }
+        case "text-delta":
+        case "tool-call-delta":
+        case "tool-result":
+        case "finish":
+        case "error":
+          break;
+        default: {
+          const unhandledType = chunkType;
+          throw new Error(`Unhandled chunk type: ${unhandledType}`);
+        }
+      }
+    },
+    async flush() {
+      await Promise.all(toolCallExecutions.values());
+    }
+  });
+}
+
+// src/runtimes/edge/partial-json/parse-partial-json.ts
+import sjson2 from "secure-json-parse";
+
+// src/runtimes/edge/partial-json/fix-json.ts
+function fixJson(input) {
+  const stack = ["ROOT"];
+  let lastValidIndex = -1;
+  let literalStart = null;
+  function processValueStart(char, i, swapState) {
+    {
+      switch (char) {
+        case '"': {
+          lastValidIndex = i;
+          stack.pop();
+          stack.push(swapState);
+          stack.push("INSIDE_STRING");
+          break;
+        }
+        case "f":
+        case "t":
+        case "n": {
+          lastValidIndex = i;
+          literalStart = i;
+          stack.pop();
+          stack.push(swapState);
+          stack.push("INSIDE_LITERAL");
+          break;
+        }
+        case "-": {
+          stack.pop();
+          stack.push(swapState);
+          stack.push("INSIDE_NUMBER");
+          break;
+        }
+        case "0":
+        case "1":
+        case "2":
+        case "3":
+        case "4":
+        case "5":
+        case "6":
+        case "7":
+        case "8":
+        case "9": {
+          lastValidIndex = i;
+          stack.pop();
+          stack.push(swapState);
+          stack.push("INSIDE_NUMBER");
+          break;
+        }
+        case "{": {
+          lastValidIndex = i;
+          stack.pop();
+          stack.push(swapState);
+          stack.push("INSIDE_OBJECT_START");
+          break;
+        }
+        case "[": {
+          lastValidIndex = i;
+          stack.pop();
+          stack.push(swapState);
+          stack.push("INSIDE_ARRAY_START");
+          break;
+        }
+      }
+    }
+  }
+  function processAfterObjectValue(char, i) {
+    switch (char) {
+      case ",": {
+        stack.pop();
+        stack.push("INSIDE_OBJECT_AFTER_COMMA");
+        break;
+      }
+      case "}": {
+        lastValidIndex = i;
+        stack.pop();
+        break;
+      }
+    }
+  }
+  function processAfterArrayValue(char, i) {
+    switch (char) {
+      case ",": {
+        stack.pop();
+        stack.push("INSIDE_ARRAY_AFTER_COMMA");
+        break;
+      }
+      case "]": {
+        lastValidIndex = i;
+        stack.pop();
+        break;
+      }
+    }
+  }
+  for (let i = 0; i < input.length; i++) {
+    const char = input[i];
+    const currentState = stack[stack.length - 1];
+    switch (currentState) {
+      case "ROOT":
+        processValueStart(char, i, "FINISH");
+        break;
+      case "INSIDE_OBJECT_START": {
+        switch (char) {
+          case '"': {
+            stack.pop();
+            stack.push("INSIDE_OBJECT_KEY");
+            break;
+          }
+          case "}": {
+            lastValidIndex = i;
+            stack.pop();
+            break;
+          }
+        }
+        break;
+      }
+      case "INSIDE_OBJECT_AFTER_COMMA": {
+        switch (char) {
+          case '"': {
+            stack.pop();
+            stack.push("INSIDE_OBJECT_KEY");
+            break;
+          }
+        }
+        break;
+      }
+      case "INSIDE_OBJECT_KEY": {
+        switch (char) {
+          case '"': {
+            stack.pop();
+            stack.push("INSIDE_OBJECT_AFTER_KEY");
+            break;
+          }
+        }
+        break;
+      }
+      case "INSIDE_OBJECT_AFTER_KEY": {
+        switch (char) {
+          case ":": {
+            stack.pop();
+            stack.push("INSIDE_OBJECT_BEFORE_VALUE");
+            break;
+          }
+        }
+        break;
+      }
+      case "INSIDE_OBJECT_BEFORE_VALUE": {
+        processValueStart(char, i, "INSIDE_OBJECT_AFTER_VALUE");
+        break;
+      }
+      case "INSIDE_OBJECT_AFTER_VALUE": {
+        processAfterObjectValue(char, i);
+        break;
+      }
+      case "INSIDE_STRING": {
+        switch (char) {
+          case '"': {
+            stack.pop();
+            lastValidIndex = i;
+            break;
+          }
+          case "\\": {
+            stack.push("INSIDE_STRING_ESCAPE");
+            break;
+          }
+          default: {
+            lastValidIndex = i;
+          }
+        }
+        break;
+      }
+      case "INSIDE_ARRAY_START": {
+        switch (char) {
+          case "]": {
+            lastValidIndex = i;
+            stack.pop();
+            break;
+          }
+          default: {
+            lastValidIndex = i;
+            processValueStart(char, i, "INSIDE_ARRAY_AFTER_VALUE");
+            break;
+          }
+        }
+        break;
+      }
+      case "INSIDE_ARRAY_AFTER_VALUE": {
+        switch (char) {
+          case ",": {
+            stack.pop();
+            stack.push("INSIDE_ARRAY_AFTER_COMMA");
+            break;
+          }
+          case "]": {
+            lastValidIndex = i;
+            stack.pop();
+            break;
+          }
+          default: {
+            lastValidIndex = i;
+            break;
+          }
+        }
+        break;
+      }
+      case "INSIDE_ARRAY_AFTER_COMMA": {
+        processValueStart(char, i, "INSIDE_ARRAY_AFTER_VALUE");
+        break;
+      }
+      case "INSIDE_STRING_ESCAPE": {
+        stack.pop();
+        lastValidIndex = i;
+        break;
+      }
+      case "INSIDE_NUMBER": {
+        switch (char) {
+          case "0":
+          case "1":
+          case "2":
+          case "3":
+          case "4":
+          case "5":
+          case "6":
+          case "7":
+          case "8":
+          case "9": {
+            lastValidIndex = i;
+            break;
+          }
+          case "e":
+          case "E":
+          case "-":
+          case ".": {
+            break;
+          }
+          case ",": {
+            stack.pop();
+            if (stack[stack.length - 1] === "INSIDE_ARRAY_AFTER_VALUE") {
+              processAfterArrayValue(char, i);
+            }
+            if (stack[stack.length - 1] === "INSIDE_OBJECT_AFTER_VALUE") {
+              processAfterObjectValue(char, i);
+            }
+            break;
+          }
+          case "}": {
+            stack.pop();
+            if (stack[stack.length - 1] === "INSIDE_OBJECT_AFTER_VALUE") {
+              processAfterObjectValue(char, i);
+            }
+            break;
+          }
+          case "]": {
+            stack.pop();
+            if (stack[stack.length - 1] === "INSIDE_ARRAY_AFTER_VALUE") {
+              processAfterArrayValue(char, i);
+            }
+            break;
+          }
+          default: {
+            stack.pop();
+            break;
+          }
+        }
+        break;
+      }
+      case "INSIDE_LITERAL": {
+        const partialLiteral = input.substring(literalStart, i + 1);
+        if (!"false".startsWith(partialLiteral) && !"true".startsWith(partialLiteral) && !"null".startsWith(partialLiteral)) {
+          stack.pop();
+          if (stack[stack.length - 1] === "INSIDE_OBJECT_AFTER_VALUE") {
+            processAfterObjectValue(char, i);
+          } else if (stack[stack.length - 1] === "INSIDE_ARRAY_AFTER_VALUE") {
+            processAfterArrayValue(char, i);
+          }
+        } else {
+          lastValidIndex = i;
+        }
+        break;
+      }
+    }
+  }
+  let result = input.slice(0, lastValidIndex + 1);
+  for (let i = stack.length - 1; i >= 0; i--) {
+    const state = stack[i];
+    switch (state) {
+      case "INSIDE_STRING": {
+        result += '"';
+        break;
+      }
+      case "INSIDE_OBJECT_KEY":
+      case "INSIDE_OBJECT_AFTER_KEY":
+      case "INSIDE_OBJECT_AFTER_COMMA":
+      case "INSIDE_OBJECT_START":
+      case "INSIDE_OBJECT_BEFORE_VALUE":
+      case "INSIDE_OBJECT_AFTER_VALUE": {
+        result += "}";
+        break;
+      }
+      case "INSIDE_ARRAY_START":
+      case "INSIDE_ARRAY_AFTER_COMMA":
+      case "INSIDE_ARRAY_AFTER_VALUE": {
+        result += "]";
+        break;
+      }
+      case "INSIDE_LITERAL": {
+        const partialLiteral = input.substring(literalStart, input.length);
+        if ("true".startsWith(partialLiteral)) {
+          result += "true".slice(partialLiteral.length);
+        } else if ("false".startsWith(partialLiteral)) {
+          result += "false".slice(partialLiteral.length);
+        } else if ("null".startsWith(partialLiteral)) {
+          result += "null".slice(partialLiteral.length);
+        }
+      }
+    }
+  }
+  return result;
+}
+
+// src/runtimes/edge/partial-json/parse-partial-json.ts
+var parsePartialJson = (json) => {
+  try {
+    return sjson2.parse(json);
+  } catch {
+    try {
+      return sjson2.parse(fixJson(json));
+    } catch {
+      return void 0;
+    }
+  }
+};
+
+// src/runtimes/edge/streams/runResultStream.ts
+function runResultStream(initialContent) {
+  let message = {
+    content: initialContent
+  };
+  const currentToolCall = { toolCallId: "", argsText: "" };
+  return new TransformStream({
+    transform(chunk, controller) {
+      const chunkType = chunk.type;
+      switch (chunkType) {
+        case "text-delta": {
+          message = appendOrUpdateText(message, chunk.textDelta);
+          controller.enqueue(message);
+          break;
+        }
+        case "tool-call-delta": {
+          const { toolCallId, toolName, argsTextDelta } = chunk;
+          if (currentToolCall.toolCallId !== toolCallId) {
+            currentToolCall.toolCallId = toolCallId;
+            currentToolCall.argsText = argsTextDelta;
+          } else {
+            currentToolCall.argsText += argsTextDelta;
+          }
+          message = appendOrUpdateToolCall(
+            message,
+            toolCallId,
+            toolName,
+            currentToolCall.argsText
+          );
+          controller.enqueue(message);
+          break;
+        }
+        case "tool-call": {
+          break;
+        }
+        case "tool-result": {
+          message = appendOrUpdateToolResult(
+            message,
+            chunk.toolCallId,
+            chunk.toolName,
+            chunk.result
+          );
+          controller.enqueue(message);
+          break;
+        }
+        case "finish": {
+          message = appendOrUpdateFinish(message, chunk);
+          controller.enqueue(message);
+          break;
+        }
+        case "error": {
+          throw chunk.error;
+        }
+        default: {
+          const unhandledType = chunkType;
+          throw new Error(`Unhandled chunk type: ${unhandledType}`);
+        }
+      }
+    }
+  });
+}
+var appendOrUpdateText = (message, textDelta) => {
+  let contentParts = message.content;
+  let contentPart = message.content.at(-1);
+  if (contentPart?.type !== "text") {
+    contentPart = { type: "text", text: textDelta };
+  } else {
+    contentParts = contentParts.slice(0, -1);
+    contentPart = { type: "text", text: contentPart.text + textDelta };
+  }
+  return {
+    ...message,
+    content: contentParts.concat([contentPart])
+  };
+};
+var appendOrUpdateToolCall = (message, toolCallId, toolName, argsText) => {
+  let contentParts = message.content;
+  let contentPart = message.content.at(-1);
+  if (contentPart?.type !== "tool-call" || contentPart.toolCallId !== toolCallId) {
+    contentPart = {
+      type: "tool-call",
+      toolCallId,
+      toolName,
+      argsText,
+      args: parsePartialJson(argsText)
+    };
+  } else {
+    contentParts = contentParts.slice(0, -1);
+    contentPart = {
+      ...contentPart,
+      argsText,
+      args: parsePartialJson(argsText)
+    };
+  }
+  return {
+    ...message,
+    content: contentParts.concat([contentPart])
+  };
+};
+var appendOrUpdateToolResult = (message, toolCallId, toolName, result) => {
+  let found = false;
+  const newContentParts = message.content.map((part) => {
+    if (part.type !== "tool-call" || part.toolCallId !== toolCallId)
+      return part;
+    found = true;
+    if (part.toolName !== toolName)
+      throw new Error(
+        `Tool call ${toolCallId} found with tool name ${part.toolName}, but expected ${toolName}`
+      );
+    return {
+      ...part,
+      result
+    };
+  });
+  if (!found)
+    throw new Error(
+      `Received tool result for unknown tool call "${toolName}" / "${toolCallId}". This is likely an internal bug in assistant-ui.`
+    );
+  return {
+    ...message,
+    content: newContentParts
+  };
+};
+var appendOrUpdateFinish = (message, chunk) => {
+  const { type, ...rest } = chunk;
+  return {
+    ...message,
+    status: {
+      type: "done",
+      ...rest
+    }
+  };
+};
+
+// src/runtimes/edge/createEdgeRuntimeAPI.ts
+var LanguageModelSettingsSchema = z3.object({
+  maxTokens: z3.number().int().positive().optional(),
+  temperature: z3.number().optional(),
+  topP: z3.number().optional(),
+  presencePenalty: z3.number().optional(),
+  frequencyPenalty: z3.number().optional(),
+  seed: z3.number().int().optional(),
+  headers: z3.record(z3.string().optional()).optional()
+});
+var voidStream = () => {
+  return new WritableStream({
+    abort(reason) {
+      console.error("Server stream processing aborted:", reason);
+    }
+  });
+};
+var createEdgeRuntimeAPI = ({
+  model,
+  system: serverSystem,
+  tools: serverTools = {},
+  toolChoice,
+  onFinish,
+  ...unsafeSettings
+}) => {
+  const settings = LanguageModelSettingsSchema.parse(unsafeSettings);
+  const lmServerTools = toLanguageModelTools(serverTools);
+  const hasServerTools = Object.values(serverTools).some((v) => !!v.execute);
   const POST = async (request) => {
-    const { system, messages, tools } = await request.json();
-    const { stream } = await streamMessage({
+    const {
+      system: clientSystem,
+      tools: clientTools,
+      messages
+    } = await request.json();
+    const systemMessages = [];
+    if (serverSystem) systemMessages.push(serverSystem);
+    if (clientSystem) systemMessages.push(clientSystem);
+    const system = systemMessages.join("\n\n");
+    for (const clientTool of clientTools) {
+      if (serverTools?.[clientTool.name]) {
+        throw new Error(
+          `Tool ${clientTool.name} was defined in both the client and server tools. This is not allowed.`
+        );
+      }
+    }
+    let stream;
+    const streamResult = await streamMessage({
+      ...settings,
       model,
       abortSignal: request.signal,
-      ...system ? { system } : void 0,
+      ...!!system ? { system } : void 0,
       messages,
-      tools
+      tools: lmServerTools.concat(clientTools),
+      ...toolChoice ? { toolChoice } : void 0
     });
-    return new Response(stream, {
-      headers: {
-        contentType: "text/plain; charset=utf-8"
+    stream = streamResult.stream;
+    const canExecuteTools = hasServerTools && toolChoice?.type !== "none";
+    if (canExecuteTools) {
+      stream = stream.pipeThrough(toolResultStream(serverTools));
+    }
+    if (canExecuteTools || onFinish) {
+      const tees = stream.tee();
+      stream = tees[0];
+      let serverStream = tees[1];
+      if (onFinish) {
+        serverStream = serverStream.pipeThrough(runResultStream([])).pipeThrough(
+          new TransformStream({
+            transform(chunk) {
+              if (chunk.status?.type !== "done") return;
+              const resultingMessages = [
+                ...messages,
+                {
+                  role: "assistant",
+                  content: chunk.content
+                }
+              ];
+              onFinish({
+                finishReason: chunk.status.finishReason,
+                usage: chunk.status.usage,
+                messages: resultingMessages,
+                logProbs: chunk.status.logprops,
+                warnings: streamResult.warnings,
+                rawCall: streamResult.rawCall,
+                rawResponse: streamResult.rawResponse
+              });
+            }
+          })
+        );
       }
-    });
+      serverStream.pipeTo(voidStream()).catch((e) => {
+        console.error("Server stream processing error:", e);
+      });
+    }
+    return new Response(
+      stream.pipeThrough(assistantEncoderStream()).pipeThrough(new TextEncoderStream()),
+      {
+        headers: {
+          contentType: "text/plain; charset=utf-8"
+        }
+      }
+    );
   };
   return { POST };
 };
@@ -203,7 +849,7 @@ async function streamMessage({
   toolChoice,
   ...options
 }) {
-  const { stream, warnings, rawResponse } = await model.doStream({
+  return model.doStream({
     inputFormat: "messages",
     mode: {
       type: "regular",
@@ -213,11 +859,6 @@ async function streamMessage({
     prompt: convertToLanguageModelPrompt(system, messages),
     ...options
   });
-  return {
-    stream: stream.pipeThrough(assistantEncoderStream()).pipeThrough(new TextEncoderStream()),
-    warnings,
-    rawResponse
-  };
 }
 function convertToLanguageModelPrompt(system, messages) {
   const languageModelMessages = [];