@assistant-ui/react/dist/legacy-runtime/runtime-cores/assistant-transport/useToolInvocations.js

TypeScript/React library for AI Chat

// src/legacy-runtime/runtime-cores/assistant-transport/useToolInvocations.ts
import { useEffect, useRef, useState } from "react";
import {
  createAssistantStreamController,
  ToolResponse,
  unstable_toolResultStream
} from "assistant-stream";
import {
  AssistantMetaTransformStream
} from "assistant-stream/utils";
var isArgsTextComplete = (argsText) => {
  try {
    JSON.parse(argsText);
    return true;
  } catch {
    return false;
  }
};
function useToolInvocations({
  state,
  getTools,
  onResult
}) {
  const lastToolStates = useRef({});
  const acRef = useRef(new AbortController());
  const [controller] = useState(() => {
    const [stream, controller2] = createAssistantStreamController();
    const transform = unstable_toolResultStream(
      getTools,
      () => acRef.current?.signal ?? new AbortController().signal
    );
    stream.pipeThrough(transform).pipeThrough(new AssistantMetaTransformStream()).pipeTo(
      new WritableStream({
        write(chunk) {
          if (chunk.type === "result") {
            if (lastToolStates.current[chunk.meta.toolCallId]?.hasResult)
              return;
            onResult({
              type: "add-tool-result",
              toolCallId: chunk.meta.toolCallId,
              toolName: chunk.meta.toolName,
              result: chunk.result,
              isError: chunk.isError,
              ...chunk.artifact && { artifact: chunk.artifact }
            });
          }
        }
      })
    );
    return controller2;
  });
  const ignoredToolIds = useRef(/* @__PURE__ */ new Set());
  const isInititialState = useRef(true);
  useEffect(() => {
    if (isInititialState.current) {
      state.messages.forEach((message) => {
        message.content.forEach((content) => {
          if (content.type === "tool-call") {
            ignoredToolIds.current.add(content.toolCallId);
          }
        });
      });
      isInititialState.current = false;
    } else {
      state.messages.forEach((message) => {
        message.content.forEach((content) => {
          if (content.type === "tool-call") {
            if (ignoredToolIds.current.has(content.toolCallId)) {
              return;
            }
            let lastState = lastToolStates.current[content.toolCallId];
            if (!lastState) {
              const toolCallController = controller.addToolCallPart({
                toolName: content.toolName,
                toolCallId: content.toolCallId
              });
              lastState = {
                argsText: "",
                hasResult: false,
                controller: toolCallController
              };
              lastToolStates.current[content.toolCallId] = lastState;
            }
            if (content.argsText !== lastState.argsText) {
              if (!content.argsText.startsWith(lastState.argsText)) {
                throw new Error(
                  `Tool call argsText can only be appended, not updated: ${content.argsText} does not start with ${lastState.argsText}`
                );
              }
              const argsTextDelta = content.argsText.slice(
                lastState.argsText.length
              );
              lastState.controller.argsText.append(argsTextDelta);
              if (isArgsTextComplete(content.argsText)) {
                lastState.controller.argsText.close();
              }
              lastToolStates.current[content.toolCallId] = {
                argsText: content.argsText,
                hasResult: lastState.hasResult,
                controller: lastState.controller
              };
            }
            if (content.result !== void 0 && !lastState.hasResult) {
              lastState.controller.setResponse(
                new ToolResponse({
                  result: content.result,
                  artifact: content.artifact,
                  isError: content.isError
                })
              );
              lastState.controller.close();
              lastToolStates.current[content.toolCallId] = {
                hasResult: true,
                argsText: lastState.argsText,
                controller: lastState.controller
              };
            }
          }
        });
      });
    }
  }, [state]);
  return {
    reset: () => {
      acRef.current.abort();
      acRef.current = new AbortController();
      isInititialState.current = true;
    },
    abort: () => {
      acRef.current.abort();
      acRef.current = new AbortController();
    }
  };
}
export {
  useToolInvocations
};
//# sourceMappingURL=useToolInvocations.js.map