[WIP] Add new tool call streaming system

examples/tool-calls-beta-zod.ts

@@ -0,0 +1,106 @@
+#!/usr/bin/env -S npm run tsn -T
+
+import OpenAI from 'openai';
+import { betaZodTool } from 'openai/helpers/beta/zod';
+// import { BetaToolUseBlock } from 'openai/helpers/beta';
+import { z } from 'zod';
+
+const client = new OpenAI();
+
+async function main() {
+  const runner = client.beta.chat.completions.toolRunner({
+    messages: [
+      {
+        role: 'user',
+        content: `I'm planning a trip to San Francisco and I need some information. Can you help me with the weather, current time, and currency exchange rates (from EUR)? Please use parallel tool use`,
+      },
+    ],
+    tools: [
+      betaZodTool({
+        name: 'getWeather',
+        description: 'Get the weather at a specific location',
+        inputSchema: z.object({
+          location: z.string().describe('The city and state, e.g. San Francisco, CA'),
+        }),
+        run: ({ location }) => {
+          return `The weather is sunny with a temperature of 20°C in ${location}.`;
+        },
+      }),
+      betaZodTool({
+        name: 'getTime',
+        description: 'Get the current time in a specific timezone',
+        inputSchema: z.object({
+          timezone: z.string().describe('The timezone, e.g. America/Los_Angeles'),
+        }),
+        run: ({ timezone }) => {
+          return `The current time in ${timezone} is 3:00 PM.`;
+        },
+      }),
+      betaZodTool({
+        name: 'getCurrencyExchangeRate',
+        description: 'Get the exchange rate between two currencies',
+        inputSchema: z.object({
+          from_currency: z.string().describe('The currency to convert from, e.g. USD'),
+          to_currency: z.string().describe('The currency to convert to, e.g. EUR'),
+        }),
+        run: ({ from_currency, to_currency }) => {
+          return `The exchange rate from ${from_currency} to ${to_currency} is 0.85.`;
+        },
+      }),
+    ],
+    model: 'gpt-4o',
+    max_tokens: 1024,
+    // This limits the conversation to at most 10 back and forth between the API.
+    max_iterations: 10,
+  });
+
+  console.log(`\n🚀 Running tools...\n`);
+
+  for await (const message of runner) {
+    console.log(`┌─ Message ${message.id} `.padEnd(process.stdout.columns, '─'));
+    console.log();
+
+    const { choices } = message;
+    const firstChoice = choices.at(0)!;
+
+    // When we get a tool call request it's null
+    if (firstChoice.message.content !== null) {
+      console.log(`${firstChoice.message.content}\n`);
+    } else {
+      // each tool call (could be many)
+      for (const toolCall of firstChoice.message.tool_calls ?? []) {
+        if (toolCall.type === 'function') {
+          console.log(`${toolCall.function.name}(${JSON.stringify(toolCall.function.arguments, null, 2)})\n`);
+        }
+      }
+    }
+
+    console.log(`└─`.padEnd(process.stdout.columns, '─'));
+    console.log();
+    console.log();
+
+    const defaultResponse = await runner.generateToolResponse();
+    if (defaultResponse && Array.isArray(defaultResponse)) {
+      console.log(`┌─ Response `.padEnd(process.stdout.columns, '─'));
+      console.log();
+
+      for (const toolResponse of defaultResponse) {
+        if (toolResponse.role === 'tool') {
+          const toolCall = firstChoice.message.tool_calls?.find((tc) => tc.id === toolResponse.tool_call_id);
+          if (toolCall && toolCall.type === 'function') {
+            console.log(`${toolCall.function.name}(): ${toolResponse.content}`);
+          }
+        }
+      }
+
+      console.log();
+      console.log(`└─`.padEnd(process.stdout.columns, '─'));
+      console.log();
+      console.log();
+    }
+  }
+
+  console.log(JSON.stringify(runner.params, null, 2));
+}
+
+main();

src/helpers/beta/zod.ts

@@ -0,0 +1,111 @@
+import { transformJSONSchema } from '../..//lib/transform-json-schema';
+import type { infer as zodInfer, ZodType } from 'zod';
+import * as z from 'zod';
+import { OpenAIError } from '../../core/error';
+import type { BetaRunnableTool, Promisable } from '../../lib/beta/BetaRunnableTool';
+import type { AutoParseableBetaOutputFormat } from '../../lib/beta-parser';
+import { FunctionTool } from '../../resources/beta';
+// import { AutoParseableBetaOutputFormat } from '../../lib/beta-parser';
+// import { BetaRunnableTool, Promisable } from '../../lib/tools/BetaRunnableTool';
+// import { BetaToolResultContentBlockParam } from '../../resources/beta';
+/**
+ * Creates a JSON schema output format object from the given Zod schema.
+ *
+ * If this is passed to the `.parse()` method then the response message will contain a
+ * `.parsed` property that is the result of parsing the content with the given Zod object.
+ *
+ * This can be passed directly to the `.create()` method but will not
+ * result in any automatic parsing, you'll have to parse the response yourself.
+ */
+export function betaZodOutputFormat<ZodInput extends ZodType>(
+  zodObject: ZodInput,
+): AutoParseableBetaOutputFormat<zodInfer<ZodInput>> {
+  let jsonSchema = z.toJSONSchema(zodObject, { reused: 'ref' });
+
+  jsonSchema = transformJSONSchema(jsonSchema);
+
+  return {
+    type: 'json_schema',
+    schema: {
+      ...jsonSchema,
+    },
+    parse: (content) => {
+      const output = zodObject.safeParse(JSON.parse(content));
+
+      if (!output.success) {
+        throw new OpenAIError(
+          `Failed to parse structured output: ${output.error.message} cause: ${output.error.issues}`,
+        );
+      }
+
+      return output.data;
+    },
+  };
+}
+
+/**
+ * Creates a tool using the provided Zod schema that can be passed
+ * into the `.toolRunner()` method. The Zod schema will automatically be
+ * converted into JSON Schema when passed to the API. The provided function's
+ * input arguments will also be validated against the provided schema.
+ */
+export function betaZodTool<InputSchema extends ZodType>(options: {
+  name: string;
+  inputSchema: InputSchema;
+  description: string;
+  run: (args: zodInfer<InputSchema>) => Promisable<string | Array<FunctionTool>>; // TODO: I changed this but double check
+}): BetaRunnableTool<zodInfer<InputSchema>> {
+  const jsonSchema = z.toJSONSchema(options.inputSchema, { reused: 'ref' });
+
+  if (jsonSchema.type !== 'object') {
+    throw new Error(`Zod schema for tool "${options.name}" must be an object, but got ${jsonSchema.type}`);
+  }
+
+  // TypeScript doesn't narrow the type after the runtime check, so we need to assert it
+  const objectSchema = jsonSchema as typeof jsonSchema & { type: 'object' };
+
+  return {
+    type: 'function',
+    function: {
+      name: options.name,
+      description: options.description,
+      parameters: {
+        type: 'object',
+        properties: objectSchema.properties,
+      },
+    },
+    run: options.run,
+    parse: (args: unknown) => options.inputSchema.parse(args) as zodInfer<InputSchema>,
+  };
+}
+
+// /**
+//  * Creates a tool using the provided Zod schema that can be passed
+//  * into the `.toolRunner()` method. The Zod schema will automatically be
+//  * converted into JSON Schema when passed to the API. The provided function's
+//  * input arguments will also be validated against the provided schema.
+//  */
+// export function betaZodTool<InputSchema extends ZodType>(options: {
+//   name: string;
+//   inputSchema: InputSchema;
+//   description: string;
+//   run: (args: zodInfer<InputSchema>) => Promisable<string | Array<BetaToolResultContentBlockParam>>;
+// }): BetaRunnableTool<zodInfer<InputSchema>> {
+//   const jsonSchema = z.toJSONSchema(options.inputSchema, { reused: 'ref' });
+
+//   if (jsonSchema.type !== 'object') {
+//     throw new Error(`Zod schema for tool "${options.name}" must be an object, but got ${jsonSchema.type}`);
+//   }
+
+//   // TypeScript doesn't narrow the type after the runtime check, so we need to assert it
+//   const objectSchema = jsonSchema as typeof jsonSchema & { type: 'object' };
+
+//   return {
+//     type: 'custom',
+//     name: options.name,
+//     input_schema: objectSchema,
+//     description: options.description,
+//     run: options.run,
+//     parse: (args: unknown) => options.inputSchema.parse(args) as zodInfer<InputSchema>,
+//   };
+// }

src/internal/utils/values.ts

@@ -103,3 +103,10 @@ export const safeJSON = (text: string) => {
     return undefined;
   }
 };
+
+// Gets a value from an object, deletes the key, and returns the value (or undefined if not found)
+export const pop = <T extends Record<string, any>, K extends string>(obj: T, key: K): T[K] => {
+  const value = obj[key];
+  delete obj[key];
+  return value;
+};

src/lib/beta-parser.ts

@@ -0,0 +1,105 @@
+import { OpenAIError } from '../core/error';
+import {
+  BetaContentBlock,
+  BetaJSONOutputFormat,
+  BetaMessage,
+  BetaTextBlock,
+  MessageCreateParams,
+} from '../resources/beta/messages/messages';
+
+// vendored from typefest just to make things look a bit nicer on hover
+type Simplify<T> = { [KeyType in keyof T]: T[KeyType] } & {};
+
+export type BetaParseableMessageCreateParams = Simplify<
+  Omit<MessageCreateParams, 'output_format'> & {
+    output_format?: BetaJSONOutputFormat | AutoParseableBetaOutputFormat<any> | null;
+  }
+>;
+
+export type ExtractParsedContentFromBetaParams<Params extends BetaParseableMessageCreateParams> =
+  Params['output_format'] extends AutoParseableBetaOutputFormat<infer P> ? P : null;
+
+export type AutoParseableBetaOutputFormat<ParsedT> = BetaJSONOutputFormat & {
+  parse(content: string): ParsedT;
+};
+
+export type ParsedBetaMessage<ParsedT> = BetaMessage & {
+  content: Array<ParsedBetaContentBlock<ParsedT>>;
+  parsed_output: ParsedT | null;
+};
+
+export type ParsedBetaContentBlock<ParsedT> =
+  | (BetaTextBlock & { parsed: ParsedT | null })
+  | Exclude<BetaContentBlock, BetaTextBlock>;
+
+export function maybeParseBetaMessage<Params extends BetaParseableMessageCreateParams | null>(
+  message: BetaMessage,
+  params: Params,
+): ParsedBetaMessage<ExtractParsedContentFromBetaParams<NonNullable<Params>>> {
+  if (!params || !('parse' in (params.output_format ?? {}))) {
+    return {
+      ...message,
+      content: message.content.map((block) => {
+        if (block.type === 'text') {
+          return {
+            ...block,
+            parsed: null,
+          };
+        }
+        return block;
+      }),
+      parsed_output: null,
+    } as ParsedBetaMessage<ExtractParsedContentFromBetaParams<NonNullable<Params>>>;
+  }
+
+  return parseBetaMessage(message, params);
+}
+
+export function parseBetaMessage<Params extends BetaParseableMessageCreateParams>(
+  message: BetaMessage,
+  params: Params,
+): ParsedBetaMessage<ExtractParsedContentFromBetaParams<Params>> {
+  let firstParsed: ReturnType<typeof parseBetaOutputFormat<Params>> | null = null;
+
+  const content: Array<ParsedBetaContentBlock<ExtractParsedContentFromBetaParams<Params>>> =
+    message.content.map((block) => {
+      if (block.type === 'text') {
+        const parsed = parseBetaOutputFormat(params, block.text);
+
+        if (firstParsed === null) {
+          firstParsed = parsed;
+        }
+
+        return {
+          ...block,
+          parsed,
+        };
+      }
+      return block;
+    });
+
+  return {
+    ...message,
+    content,
+    parsed_output: firstParsed,
+  } as ParsedBetaMessage<ExtractParsedContentFromBetaParams<Params>>;
+}
+
+function parseBetaOutputFormat<Params extends BetaParseableMessageCreateParams>(
+  params: Params,
+  content: string,
+): ExtractParsedContentFromBetaParams<Params> | null {
+  if (params.output_format?.type !== 'json_schema') {
+    return null;
+  }
+
+  try {
+    if ('parse' in params.output_format) {
+      return params.output_format.parse(content);
+    }
+
+    return JSON.parse(content);
+  } catch (error) {
+    throw new OpenAIError(`Failed to parse structured output: ${error}`);
+  }
+}

src/lib/beta/BetaRunnableTool.ts

@@ -0,0 +1,11 @@
+import { FunctionTool } from '../../resources/beta';
+import type { ChatCompletionTool } from '../../resources';
+
+export type Promisable<T> = T | Promise<T>;
+
+// this type is just an extension of BetaTool with a run and parse method
+// that will be called by `toolRunner()` helpers
+export type BetaRunnableTool<Input = any> = ChatCompletionTool & {
+  run: (args: Input) => Promisable<string | Array<FunctionTool>>;
+  parse: (content: unknown) => Input;
+};