@assistant-ui/react/src/runtimes/edge/streams/toolResultStream.ts

Typescript/React library for AI Chat

import { Tool } from "../../../model-context/ModelContextTypes";
import { z } from "zod";
import { ToolExecutionStream } from "assistant-stream";

export function toolResultStream(
  tools: Record<string, Tool<any, any>> | undefined,
  abortSignal: AbortSignal,
) {
  return new ToolExecutionStream(({ toolCallId, toolName, args }) => {
    const tool = tools?.[toolName];
    if (!tool || !tool.execute) return undefined;

    let executeFn = tool.execute;

    if (tool.parameters instanceof z.ZodType) {
      const result = tool.parameters.safeParse(args);
      if (!result.success) {
        executeFn =
          tool.experimental_onSchemaValidationError ??
          (() => {
            throw (
              "Function parameter validation failed. " +
              JSON.stringify(result.error.issues)
            );
          });
      }
    }

    const getResult = async () => {
      const result = await executeFn(args, {
        toolCallId,
        abortSignal,
      });
      return result === undefined ? "<no result>" : result;
    };

    return getResult();
  });
}