@assistant-ui/react 0.4.3 → 0.4.5

package/dist/edge.mjs

@@ -34,22 +34,25 @@ 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(
-            formatStreamPart(
-              "F" /* Finish */,
-              JSON.stringify(rest)
-            )
+            formatStreamPart("F" /* Finish */, rest)
           );
           break;
         }
         case "error": {
           controller.enqueue(
-            formatStreamPart(
-              "E" /* Error */,
-              JSON.stringify(chunk.error)
-            )
+            formatStreamPart("E" /* Error */, chunk.error)
           );
           break;
         }
@@ -61,64 +64,12 @@ function assistantEncoderStream() {
     }
   });
 }
-function formatStreamPart(code, value) {
+function formatStreamPart(...[code, value]) {
   return `${code}:${JSON.stringify(value)}
 `;
 }
 
-// src/runtimes/edge/createEdgeRuntimeAPI.ts
-var createEdgeRuntimeAPI = ({ model }) => {
-  const POST = async (request) => {
-    const { system, messages, tools } = await request.json();
-    const { stream } = await streamMessage({
-      model,
-      abortSignal: request.signal,
-      ...system ? { system } : void 0,
-      messages,
-      tools
-    });
-    return new Response(stream, {
-      headers: {
-        contentType: "text/plain; charset=utf-8"
-      }
-    });
-  };
-  return { POST };
-};
-async function streamMessage({
-  model,
-  system,
-  messages,
-  tools,
-  toolChoice,
-  ...options
-}) {
-  const { stream, warnings, rawResponse } = await model.doStream({
-    inputFormat: "messages",
-    mode: {
-      type: "regular",
-      ...tools ? { tools } : void 0,
-      ...toolChoice ? { toolChoice } : void 0
-    },
-    prompt: convertToLanguageModelPrompt(system, messages),
-    ...options
-  });
-  return {
-    stream: stream.pipeThrough(assistantEncoderStream()).pipeThrough(new TextEncoderStream()),
-    warnings,
-    rawResponse
-  };
-}
-function convertToLanguageModelPrompt(system, messages) {
-  const languageModelMessages = [];
-  if (system != null) {
-    languageModelMessages.push({ role: "system", content: system });
-  }
-  languageModelMessages.push(
-    ...messages.flatMap(convertToLanguageModelMessage)
-  );
-  return languageModelMessages;
-}
+// src/runtimes/edge/converters/toLanguageModelMessages.ts
 var assistantMessageSplitter = () => {
   const stash = [];
   let assistantMessage = {
@@ -170,66 +121,752 @@ var assistantMessageSplitter = () => {
     }
   };
 };
-function convertToLanguageModelMessage(message) {
-  const role = message.role;
-  switch (role) {
-    case "system": {
-      return [{ role: "system", content: message.content[0].text }];
-    }
-    case "user": {
-      const msg = {
-        role: "user",
-        content: message.content.map(
-          (part) => {
-            const type = part.type;
-            switch (type) {
-              case "text": {
-                return part;
-              }
-              case "image": {
-                return {
-                  type: "image",
-                  image: new URL(part.image)
-                };
-              }
-              default: {
-                const unhandledType = type;
-                throw new Error(
-                  `Unspported content part type: ${unhandledType}`
-                );
+function toLanguageModelMessages(message) {
+  return message.flatMap((message2) => {
+    const role = message2.role;
+    switch (role) {
+      case "system": {
+        return [{ role: "system", content: message2.content[0].text }];
+      }
+      case "user": {
+        const msg = {
+          role: "user",
+          content: message2.content.map(
+            (part) => {
+              const type = part.type;
+              switch (type) {
+                case "text": {
+                  return part;
+                }
+                case "image": {
+                  return {
+                    type: "image",
+                    image: new URL(part.image)
+                  };
+                }
+                default: {
+                  const unhandledType = type;
+                  throw new Error(
+                    `Unspported content part type: ${unhandledType}`
+                  );
+                }
               }
             }
+          )
+        };
+        return [msg];
+      }
+      case "assistant": {
+        const splitter = assistantMessageSplitter();
+        for (const part of message2.content) {
+          const type = part.type;
+          switch (type) {
+            case "text": {
+              splitter.addTextContentPart(part);
+              break;
+            }
+            case "tool-call": {
+              splitter.addToolCallPart(part);
+              break;
+            }
+            default: {
+              const unhandledType = type;
+              throw new Error(`Unhandled content part type: ${unhandledType}`);
+            }
+          }
+        }
+        return splitter.getMessages();
+      }
+      default: {
+        const unhandledRole = role;
+        throw new Error(`Unknown message role: ${unhandledRole}`);
+      }
+    }
+  });
+}
+
+// src/runtimes/edge/createEdgeRuntimeAPI.ts
+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);
+                  }
+                })()
+              );
+            }
           }
-        )
-      };
-      return [msg];
+          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());
     }
-    case "assistant": {
-      const splitter = assistantMessageSplitter();
-      for (const part of message.content) {
-        const type = part.type;
-        switch (type) {
-          case "text": {
-            splitter.addTextContentPart(part);
+  });
+}
+
+// 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;
           }
-          case "tool-call": {
-            splitter.addToolCallPart(part);
+        }
+        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: {
-            const _exhaustiveCheck = type;
-            throw new Error(`Unhandled content part type: ${_exhaustiveCheck}`);
+            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;
       }
-      return splitter.getMessages();
     }
-    default: {
-      const unhandledRole = role;
-      throw new Error(`Unknown message role: ${unhandledRole}`);
+  }
+  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: 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,
+      messages,
+      tools: lmServerTools.concat(clientTools),
+      ...toolChoice ? { toolChoice } : void 0
+    });
+    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 };
+};
+async function streamMessage({
+  model,
+  system,
+  messages,
+  tools,
+  toolChoice,
+  ...options
+}) {
+  return model.doStream({
+    inputFormat: "messages",
+    mode: {
+      type: "regular",
+      ...tools ? { tools } : void 0,
+      ...toolChoice ? { toolChoice } : void 0
+    },
+    prompt: convertToLanguageModelPrompt(system, messages),
+    ...options
+  });
+}
+function convertToLanguageModelPrompt(system, messages) {
+  const languageModelMessages = [];
+  if (system != null) {
+    languageModelMessages.push({ role: "system", content: system });
+  }
+  languageModelMessages.push(...toLanguageModelMessages(messages));
+  return languageModelMessages;
 }
 export {
   createEdgeRuntimeAPI