@assistant-ui/react
Version:
TypeScript/React library for AI Chat
201 lines • 9.34 kB
JavaScript
import { useEffect, useRef, useState } from "react";
import { createAssistantStreamController, ToolResponse, unstable_toolResultStream, } from "assistant-stream";
import { AssistantMetaTransformStream, } from "assistant-stream/utils";
const isArgsTextComplete = (argsText) => {
try {
JSON.parse(argsText);
return true;
}
catch {
return false;
}
};
export function useToolInvocations({ state, getTools, onResult, setToolStatuses, }) {
const lastToolStates = useRef({});
const humanInputRef = useRef(new Map());
const acRef = useRef(new AbortController());
const executingCountRef = useRef(0);
const settledResolversRef = useRef([]);
const [controller] = useState(() => {
const [stream, controller] = createAssistantStreamController();
const transform = unstable_toolResultStream(getTools, () => acRef.current?.signal ?? new AbortController().signal, (toolCallId, payload) => {
return new Promise((resolve, reject) => {
// Reject previous human input request if it exists
const previous = humanInputRef.current.get(toolCallId);
if (previous) {
previous.reject(new Error("Human input request was superseded by a new request"));
}
humanInputRef.current.set(toolCallId, { resolve, reject });
setToolStatuses((prev) => ({
...prev,
[toolCallId]: {
type: "interrupt",
payload: { type: "human", payload },
},
}));
});
}, {
onExecutionStart: (toolCallId) => {
executingCountRef.current++;
setToolStatuses((prev) => ({
...prev,
[toolCallId]: { type: "executing" },
}));
},
onExecutionEnd: (toolCallId) => {
executingCountRef.current--;
setToolStatuses((prev) => {
const next = { ...prev };
delete next[toolCallId];
return next;
});
// Resolve any waiting abort promises when all tools have settled
if (executingCountRef.current === 0) {
settledResolversRef.current.forEach((resolve) => resolve());
settledResolversRef.current = [];
}
},
});
stream
.pipeThrough(transform)
.pipeThrough(new AssistantMetaTransformStream())
.pipeTo(new WritableStream({
write(chunk) {
if (chunk.type === "result") {
// the tool call result was already set by the backend
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 controller;
});
const ignoredToolIds = useRef(new Set());
const isInitialState = useRef(true);
useEffect(() => {
const processMessages = (messages) => {
messages.forEach((message) => {
message.content.forEach((content) => {
if (content.type === "tool-call") {
if (isInitialState.current) {
ignoredToolIds.current.add(content.toolCallId);
}
else {
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,
argsComplete: false,
controller: toolCallController,
};
lastToolStates.current[content.toolCallId] = lastState;
}
if (content.argsText !== lastState.argsText) {
if (lastState.argsComplete) {
if (process.env["NODE_ENV"] !== "production") {
console.warn("argsText updated after controller was closed:", {
previous: lastState.argsText,
next: content.argsText,
});
}
}
else {
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);
const shouldClose = isArgsTextComplete(content.argsText);
if (shouldClose) {
lastState.controller.argsText.close();
}
lastToolStates.current[content.toolCallId] = {
argsText: content.argsText,
hasResult: lastState.hasResult,
argsComplete: shouldClose,
controller: lastState.controller,
};
}
}
if (content.result !== undefined && !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,
argsComplete: true,
argsText: lastState.argsText,
controller: lastState.controller,
};
}
}
// Recursively process nested messages
if (content.messages) {
processMessages(content.messages);
}
}
});
});
};
processMessages(state.messages);
if (isInitialState.current) {
isInitialState.current = false;
}
}, [state, controller, onResult]);
const abort = () => {
humanInputRef.current.forEach(({ reject }) => {
reject(new Error("Tool execution aborted"));
});
humanInputRef.current.clear();
acRef.current.abort();
acRef.current = new AbortController();
// Return a promise that resolves when all executing tools have settled
if (executingCountRef.current === 0) {
return Promise.resolve();
}
return new Promise((resolve) => {
settledResolversRef.current.push(resolve);
});
};
return {
reset: () => {
void abort();
isInitialState.current = true;
},
abort,
resume: (toolCallId, payload) => {
const handlers = humanInputRef.current.get(toolCallId);
if (handlers) {
humanInputRef.current.delete(toolCallId);
setToolStatuses((prev) => ({
...prev,
[toolCallId]: { type: "executing" },
}));
handlers.resolve(payload);
}
else {
throw new Error(`Tool call ${toolCallId} is not waiting for human input`);
}
},
};
}
//# sourceMappingURL=useToolInvocations.js.map