@assistant-ui/react 0.4.4 → 0.4.5

package/dist/index.js

@@ -1972,7 +1972,7 @@ var ThreadPrimitiveSuggestion = createActionButton(
 );
 
 // src/runtimes/local/useLocalRuntime.tsx
-var import_react53 = require("react");
+var import_react54 = require("react");
 
 // src/utils/idUtils.tsx
 var import_non_secure = require("nanoid/non-secure");
@@ -2273,163 +2273,8 @@ var TooltipIconButton = (0, import_react52.forwardRef)(({ children, tooltip, sid
 });
 TooltipIconButton.displayName = "TooltipIconButton";
 
-// src/runtimes/local/LocalRuntime.tsx
-var LocalRuntime = class extends BaseAssistantRuntime {
-  _proxyConfigProvider;
-  constructor(adapter) {
-    const proxyConfigProvider = new ProxyConfigProvider();
-    super(new LocalThreadRuntime(proxyConfigProvider, adapter));
-    this._proxyConfigProvider = proxyConfigProvider;
-  }
-  set adapter(adapter) {
-    this.thread.adapter = adapter;
-  }
-  registerModelConfigProvider(provider) {
-    return this._proxyConfigProvider.registerModelConfigProvider(provider);
-  }
-  switchToThread(threadId) {
-    if (threadId) {
-      throw new Error("LocalRuntime does not yet support switching threads");
-    }
-    return this.thread = new LocalThreadRuntime(
-      this._proxyConfigProvider,
-      this.thread.adapter
-    );
-  }
-};
-var CAPABILITIES = Object.freeze({
-  edit: true,
-  reload: true,
-  cancel: true,
-  copy: true
-});
-var LocalThreadRuntime = class {
-  constructor(configProvider, adapter) {
-    this.configProvider = configProvider;
-    this.adapter = adapter;
-  }
-  _subscriptions = /* @__PURE__ */ new Set();
-  abortController = null;
-  repository = new MessageRepository();
-  capabilities = CAPABILITIES;
-  get messages() {
-    return this.repository.getMessages();
-  }
-  get isRunning() {
-    return this.abortController != null;
-  }
-  getBranches(messageId) {
-    return this.repository.getBranches(messageId);
-  }
-  switchToBranch(branchId) {
-    this.repository.switchToBranch(branchId);
-    this.notifySubscribers();
-  }
-  async append(message) {
-    if (message.role !== "user")
-      throw new Error(
-        "Only appending user messages are supported in LocalRuntime. This is likely an internal bug in assistant-ui."
-      );
-    const userMessageId = generateId();
-    const userMessage = {
-      id: userMessageId,
-      role: "user",
-      content: message.content,
-      createdAt: /* @__PURE__ */ new Date()
-    };
-    this.repository.addOrUpdateMessage(message.parentId, userMessage);
-    await this.startRun(userMessageId);
-  }
-  async startRun(parentId) {
-    this.repository.resetHead(parentId);
-    const messages = this.repository.getMessages();
-    const message = {
-      id: generateId(),
-      role: "assistant",
-      status: { type: "in_progress" },
-      content: [{ type: "text", text: "" }],
-      createdAt: /* @__PURE__ */ new Date()
-    };
-    this.repository.addOrUpdateMessage(parentId, { ...message });
-    this.abortController?.abort();
-    this.abortController = new AbortController();
-    this.notifySubscribers();
-    try {
-      const updateHandler = ({ content }) => {
-        message.content = content;
-        this.repository.addOrUpdateMessage(parentId, { ...message });
-        this.notifySubscribers();
-      };
-      const result = await this.adapter.run({
-        messages,
-        abortSignal: this.abortController.signal,
-        config: this.configProvider.getModelConfig(),
-        onUpdate: updateHandler
-      });
-      if (result !== void 0) {
-        updateHandler(result);
-      }
-      if (result.status?.type === "in_progress")
-        throw new Error(
-          "Unexpected in_progress status returned from ChatModelAdapter"
-        );
-      message.status = result.status ?? { type: "done" };
-      this.repository.addOrUpdateMessage(parentId, { ...message });
-    } catch (e) {
-      message.status = { type: "error", error: e };
-      this.repository.addOrUpdateMessage(parentId, { ...message });
-      throw e;
-    } finally {
-      this.abortController = null;
-      this.notifySubscribers();
-    }
-  }
-  cancelRun() {
-    if (!this.abortController) return;
-    this.abortController.abort();
-    this.abortController = null;
-  }
-  notifySubscribers() {
-    for (const callback of this._subscriptions) callback();
-  }
-  subscribe(callback) {
-    this._subscriptions.add(callback);
-    return () => this._subscriptions.delete(callback);
-  }
-  addToolResult({ messageId, toolCallId, result }) {
-    const { parentId, message } = this.repository.getMessage(messageId);
-    if (message.role !== "assistant")
-      throw new Error("Tried to add tool result to non-assistant message");
-    let found = false;
-    const newContent = message.content.map((c) => {
-      if (c.type !== "tool-call") return c;
-      if (c.toolCallId !== toolCallId) return c;
-      found = true;
-      return {
-        ...c,
-        result
-      };
-    });
-    if (!found)
-      throw new Error("Tried to add tool result to non-existing tool call");
-    this.repository.addOrUpdateMessage(parentId, {
-      ...message,
-      content: newContent
-    });
-  }
-};
-
-// src/runtimes/local/useLocalRuntime.tsx
-var useLocalRuntime = (adapter) => {
-  const [runtime] = (0, import_react53.useState)(() => new LocalRuntime(adapter));
-  (0, import_react53.useInsertionEffect)(() => {
-    runtime.adapter = adapter;
-  });
-  return runtime;
-};
-
 // src/runtimes/edge/useEdgeRuntime.ts
-var import_react54 = require("react");
+var import_react53 = require("react");
 
 // src/runtimes/edge/converters/toCoreMessages.ts
 var toCoreMessages = (message) => {
@@ -2459,6 +2304,7 @@ var toLanguageModelTools = (tools) => {
 
 // src/runtimes/edge/streams/assistantDecoderStream.ts
 function assistantDecoderStream() {
+  const toolCallNames = /* @__PURE__ */ new Map();
   let currentToolCall;
   return new TransformStream({
     transform(chunk, controller) {
@@ -2483,6 +2329,7 @@ function assistantDecoderStream() {
         }
         case "1" /* ToolCallBegin */: {
           const { id, name } = value;
+          toolCallNames.set(id, name);
           currentToolCall = { id, name, argsText: "" };
           break;
         }
@@ -2498,6 +2345,16 @@ function assistantDecoderStream() {
           });
           break;
         }
+        case "3" /* ToolCallResult */: {
+          controller.enqueue({
+            type: "tool-result",
+            toolCallType: "function",
+            toolCallId: value.id,
+            toolName: toolCallNames.get(value.id),
+            result: value.result
+          });
+          break;
+        }
         case "F" /* Finish */: {
           controller.enqueue({
             type: "finish",
@@ -2884,9 +2741,9 @@ var parsePartialJson = (json) => {
 };
 
 // src/runtimes/edge/streams/runResultStream.ts
-function runResultStream() {
+function runResultStream(initialContent) {
   let message = {
-    content: []
+    content: initialContent
   };
   const currentToolCall = { toolCallId: "", argsText: "" };
   return new TransformStream({
@@ -3069,6 +2926,7 @@ function toolResultStream(tools) {
         }
         case "text-delta":
         case "tool-call-delta":
+        case "tool-result":
         case "finish":
         case "error":
           break;
@@ -3102,7 +2960,7 @@ var EdgeChatAdapter = class {
   constructor(options) {
     this.options = options;
   }
-  async run({ messages, abortSignal, config, onUpdate }) {
+  async roundtrip(initialContent, { messages, abortSignal, config, onUpdate }) {
     const result = await fetch(this.options.api, {
       method: "POST",
       headers: {
@@ -3117,21 +2975,56 @@ var EdgeChatAdapter = class {
       }),
       signal: abortSignal
     });
-    const stream = result.body.pipeThrough(new TextDecoderStream()).pipeThrough(chunkByLineStream()).pipeThrough(assistantDecoderStream()).pipeThrough(toolResultStream(config.tools)).pipeThrough(runResultStream());
+    const stream = result.body.pipeThrough(new TextDecoderStream()).pipeThrough(chunkByLineStream()).pipeThrough(assistantDecoderStream()).pipeThrough(toolResultStream(config.tools)).pipeThrough(runResultStream(initialContent));
+    let message;
     let update;
     for await (update of asAsyncIterable(stream)) {
-      onUpdate(update);
+      message = onUpdate(update);
     }
     if (update === void 0)
       throw new Error("No data received from Edge Runtime");
-    return update;
+    return [message, update];
+  }
+  async run({ messages, abortSignal, config, onUpdate }) {
+    let roundtripAllowance = this.options.maxToolRoundtrips ?? 1;
+    let usage = {
+      promptTokens: 0,
+      completionTokens: 0
+    };
+    let result;
+    let assistantMessage;
+    do {
+      [assistantMessage, result] = await this.roundtrip(result?.content ?? [], {
+        messages: assistantMessage ? [...messages, assistantMessage] : messages,
+        abortSignal,
+        config,
+        onUpdate
+      });
+      if (result.status?.type === "done") {
+        usage.promptTokens += result.status.usage?.promptTokens ?? 0;
+        usage.completionTokens += result.status.usage?.completionTokens ?? 0;
+      }
+    } while (result.status?.type === "done" && result.status.finishReason === "tool-calls" && result.content.every((c) => c.type !== "tool-call" || !!c.result) && roundtripAllowance-- > 0);
+    if (result.status?.type === "done" && usage.promptTokens > 0) {
+      result = {
+        ...result,
+        status: {
+          ...result.status,
+          usage
+        }
+      };
+    }
+    return result;
   }
 };
 
 // src/runtimes/edge/useEdgeRuntime.ts
-var useEdgeRuntime = (options) => {
-  const [adapter] = (0, import_react54.useState)(() => new EdgeChatAdapter(options));
-  return useLocalRuntime(adapter);
+var useEdgeRuntime = ({
+  initialMessages,
+  ...options
+}) => {
+  const [adapter] = (0, import_react53.useState)(() => new EdgeChatAdapter(options));
+  return useLocalRuntime(adapter, { initialMessages });
 };
 
 // src/runtimes/edge/converters/toLanguageModelMessages.ts
@@ -3358,14 +3251,179 @@ var fromLanguageModelMessages = (lm, mergeRoundtrips) => {
 var fromCoreMessages = (message) => {
   return message.map((message2) => {
     return {
-      ...message2,
       id: generateId(),
       createdAt: /* @__PURE__ */ new Date(),
       ...message2.role === "assistant" ? {
         status: { type: "done" }
-      } : void 0
+      } : void 0,
+      ...message2
+    };
+  });
+};
+
+// src/runtimes/local/LocalRuntime.tsx
+var LocalRuntime = class extends BaseAssistantRuntime {
+  _proxyConfigProvider;
+  constructor(adapter, options) {
+    const proxyConfigProvider = new ProxyConfigProvider();
+    super(new LocalThreadRuntime(proxyConfigProvider, adapter, options));
+    this._proxyConfigProvider = proxyConfigProvider;
+  }
+  set adapter(adapter) {
+    this.thread.adapter = adapter;
+  }
+  registerModelConfigProvider(provider) {
+    return this._proxyConfigProvider.registerModelConfigProvider(provider);
+  }
+  switchToThread(threadId) {
+    if (threadId) {
+      throw new Error("LocalRuntime does not yet support switching threads");
+    }
+    return this.thread = new LocalThreadRuntime(
+      this._proxyConfigProvider,
+      this.thread.adapter
+    );
+  }
+};
+var CAPABILITIES = Object.freeze({
+  edit: true,
+  reload: true,
+  cancel: true,
+  copy: true
+});
+var LocalThreadRuntime = class {
+  constructor(configProvider, adapter, options) {
+    this.configProvider = configProvider;
+    this.adapter = adapter;
+    if (options?.initialMessages) {
+      let parentId = null;
+      const messages = fromCoreMessages(options.initialMessages);
+      for (const message of messages) {
+        this.repository.addOrUpdateMessage(parentId, message);
+        parentId = message.id;
+      }
+    }
+  }
+  _subscriptions = /* @__PURE__ */ new Set();
+  abortController = null;
+  repository = new MessageRepository();
+  capabilities = CAPABILITIES;
+  get messages() {
+    return this.repository.getMessages();
+  }
+  get isRunning() {
+    return this.abortController != null;
+  }
+  getBranches(messageId) {
+    return this.repository.getBranches(messageId);
+  }
+  switchToBranch(branchId) {
+    this.repository.switchToBranch(branchId);
+    this.notifySubscribers();
+  }
+  async append(message) {
+    if (message.role !== "user")
+      throw new Error(
+        "Only appending user messages are supported in LocalRuntime. This is likely an internal bug in assistant-ui."
+      );
+    const userMessageId = generateId();
+    const userMessage = {
+      id: userMessageId,
+      role: "user",
+      content: message.content,
+      createdAt: /* @__PURE__ */ new Date()
+    };
+    this.repository.addOrUpdateMessage(message.parentId, userMessage);
+    await this.startRun(userMessageId);
+  }
+  async startRun(parentId) {
+    this.repository.resetHead(parentId);
+    const messages = this.repository.getMessages();
+    const message = {
+      id: generateId(),
+      role: "assistant",
+      status: { type: "in_progress" },
+      content: [{ type: "text", text: "" }],
+      createdAt: /* @__PURE__ */ new Date()
     };
+    this.repository.addOrUpdateMessage(parentId, { ...message });
+    this.abortController?.abort();
+    this.abortController = new AbortController();
+    this.notifySubscribers();
+    try {
+      const updateHandler = ({ content }) => {
+        message.content = content;
+        const newMessage = { ...message };
+        this.repository.addOrUpdateMessage(parentId, newMessage);
+        this.notifySubscribers();
+        return newMessage;
+      };
+      const result = await this.adapter.run({
+        messages,
+        abortSignal: this.abortController.signal,
+        config: this.configProvider.getModelConfig(),
+        onUpdate: updateHandler
+      });
+      if (result !== void 0) {
+        updateHandler(result);
+      }
+      if (result.status?.type === "in_progress")
+        throw new Error(
+          "Unexpected in_progress status returned from ChatModelAdapter"
+        );
+      message.status = result.status ?? { type: "done" };
+      this.repository.addOrUpdateMessage(parentId, { ...message });
+    } catch (e) {
+      message.status = { type: "error", error: e };
+      this.repository.addOrUpdateMessage(parentId, { ...message });
+      throw e;
+    } finally {
+      this.abortController = null;
+      this.notifySubscribers();
+    }
+  }
+  cancelRun() {
+    if (!this.abortController) return;
+    this.abortController.abort();
+    this.abortController = null;
+  }
+  notifySubscribers() {
+    for (const callback of this._subscriptions) callback();
+  }
+  subscribe(callback) {
+    this._subscriptions.add(callback);
+    return () => this._subscriptions.delete(callback);
+  }
+  addToolResult({ messageId, toolCallId, result }) {
+    const { parentId, message } = this.repository.getMessage(messageId);
+    if (message.role !== "assistant")
+      throw new Error("Tried to add tool result to non-assistant message");
+    let found = false;
+    const newContent = message.content.map((c) => {
+      if (c.type !== "tool-call") return c;
+      if (c.toolCallId !== toolCallId) return c;
+      found = true;
+      return {
+        ...c,
+        result
+      };
+    });
+    if (!found)
+      throw new Error("Tried to add tool result to non-existing tool call");
+    this.repository.addOrUpdateMessage(parentId, {
+      ...message,
+      content: newContent
+    });
+  }
+};
+
+// src/runtimes/local/useLocalRuntime.tsx
+var useLocalRuntime = (adapter, options) => {
+  const [runtime] = (0, import_react54.useState)(() => new LocalRuntime(adapter, options));
+  (0, import_react54.useInsertionEffect)(() => {
+    runtime.adapter = adapter;
   });
+  return runtime;
 };
 
 // src/ui/thread-config.tsx
@@ -3963,10 +4021,7 @@ var ThreadScrollToBottom = (0, import_react64.forwardRef)((props, ref) => {
       thread: { scrollToBottom: { tooltip = "Scroll to bottom" } = {} } = {}
     } = {}
   } = useThreadConfig();
-  return /* @__PURE__ */ (0, import_jsx_runtime41.jsx)(thread_exports.ScrollToBottom, { asChild: true, children: /* @__PURE__ */ (0, import_jsx_runtime41.jsxs)(ThreadScrollToBottomIconButton, { tooltip, ...props, ref, children: [
-    "|",
-    props.children ?? /* @__PURE__ */ (0, import_jsx_runtime41.jsx)(import_lucide_react5.ArrowDownIcon, {})
-  ] }) });
+  return /* @__PURE__ */ (0, import_jsx_runtime41.jsx)(thread_exports.ScrollToBottom, { asChild: true, children: /* @__PURE__ */ (0, import_jsx_runtime41.jsx)(ThreadScrollToBottomIconButton, { tooltip, ...props, ref, children: props.children ?? /* @__PURE__ */ (0, import_jsx_runtime41.jsx)(import_lucide_react5.ArrowDownIcon, {}) }) });
 });
 ThreadScrollToBottom.displayName = "ThreadScrollToBottom";
 var exports11 = {