feat: pass context to runTools callbacks

src/lib/AbstractChatCompletionRunner.ts

@@ -285,16 +285,17 @@ export class AbstractChatCompletionRunner<
     return await this._createChatCompletion(client, params, options);
   }
 
-  protected async _runTools<FunctionsArgs extends BaseFunctionsArgs>(
+  protected async _runTools<FunctionsArgs extends BaseFunctionsArgs, ToolContext>(
     client: OpenAI,
     params:
-      | ChatCompletionToolRunnerParams<FunctionsArgs>
-      | ChatCompletionStreamingToolRunnerParams<FunctionsArgs>,
+      | ChatCompletionToolRunnerParams<FunctionsArgs, ToolContext>
+      | ChatCompletionStreamingToolRunnerParams<FunctionsArgs, ToolContext>,
     runner: ChatCompletionRunner<any> | ChatCompletionStreamingRunner<any>,
     options?: RunnerOptions,
   ) {
     const role = 'tool' as const;
-    const { tool_choice = 'auto', stream, ...restParams } = params;
+    const { tool_choice = 'auto', stream, toolContext: inputToolContext, ...restParams } = params;
+    const toolContext = inputToolContext as ToolContext;
     const singleFunctionToCall =
       typeof tool_choice !== 'string' && tool_choice.type === 'function' && tool_choice?.function?.name;
     const { maxChatCompletions = DEFAULT_MAX_CHAT_COMPLETIONS, afterCompletion } = options || {};
@@ -322,7 +323,7 @@ export class AbstractChatCompletionRunner<
       return tool as any as RunnableToolFunction<any>;
     });
 
-    const functionsByName: Record<string, RunnableFunction<any>> = {};
+    const functionsByName: Record<string, RunnableFunction<any, ToolContext>> = {};
     for (const f of inputTools) {
       if (f.type === 'function') {
         functionsByName[f.function.name || f.function.function.name] = f.function;
@@ -389,9 +390,9 @@ export class AbstractChatCompletionRunner<
           const content = error instanceof Error ? error.message : String(error);
           return { message: { role, tool_call_id, content }, functionCalled: false };
         }
-        rawContent = await fn.function(parsed, runner);
+        rawContent = await fn.function(parsed, runner, toolContext);
       } else {
-        rawContent = await fn.function(args, runner);
+        rawContent = await fn.function(args, runner, toolContext);
       }
 
       const content = this.#stringifyFunctionCallResult(rawContent);

src/lib/ChatCompletionRunner.ts

@@ -16,20 +16,24 @@ export interface ChatCompletionRunnerEvents extends AbstractChatCompletionRunner
   content: (content: string) => void;
 }
 
-export type ChatCompletionToolRunnerParams<FunctionsArgs extends BaseFunctionsArgs> = Omit<
-  ChatCompletionCreateParamsNonStreaming,
-  'tools'
-> & {
-  tools: RunnableTools<FunctionsArgs> | AutoParseableTool<any, true>[];
+export type ChatCompletionToolRunnerParams<
+  FunctionsArgs extends BaseFunctionsArgs,
+  ToolContext = unknown,
+> = Omit<ChatCompletionCreateParamsNonStreaming, 'tools'> & {
+  tools: RunnableTools<FunctionsArgs, ToolContext> | AutoParseableTool<any, true>[];
+  /**
+   * Context to pass to each tool callback during this run.
+   */
+  toolContext?: ToolContext;
 };
 
 export class ChatCompletionRunner<ParsedT = null> extends AbstractChatCompletionRunner<
   ChatCompletionRunnerEvents,
   ParsedT
 > {
-  static runTools<ParsedT>(
+  static runTools<ParsedT, ToolContext = unknown>(
     client: OpenAI,
-    params: ChatCompletionToolRunnerParams<any[]>,
+    params: ChatCompletionToolRunnerParams<any[], ToolContext>,
     options?: RunnerOptions,
   ): ChatCompletionRunner<ParsedT> {
     const runner = new ChatCompletionRunner<ParsedT>();

src/lib/ChatCompletionStreamingRunner.ts

@@ -14,11 +14,15 @@ export interface ChatCompletionStreamEvents extends AbstractChatCompletionRunner
   chunk: (chunk: ChatCompletionChunk, snapshot: ChatCompletionSnapshot) => void;
 }
 
-export type ChatCompletionStreamingToolRunnerParams<FunctionsArgs extends BaseFunctionsArgs> = Omit<
-  ChatCompletionCreateParamsStreaming,
-  'tools'
-> & {
-  tools: RunnableTools<FunctionsArgs> | AutoParseableTool<any, true>[];
+export type ChatCompletionStreamingToolRunnerParams<
+  FunctionsArgs extends BaseFunctionsArgs,
+  ToolContext = unknown,
+> = Omit<ChatCompletionCreateParamsStreaming, 'tools'> & {
+  tools: RunnableTools<FunctionsArgs, ToolContext> | AutoParseableTool<any, true>[];
+  /**
+   * Context to pass to each tool callback during this run.
+   */
+  toolContext?: ToolContext;
 };
 
 export class ChatCompletionStreamingRunner<ParsedT = null>
@@ -31,9 +35,9 @@ export class ChatCompletionStreamingRunner<ParsedT = null>
     return runner;
   }
 
-  static runTools<T extends (string | object)[], ParsedT = null>(
+  static runTools<T extends (string | object)[], ParsedT = null, ToolContext = unknown>(
     client: OpenAI,
-    params: ChatCompletionStreamingToolRunnerParams<T>,
+    params: ChatCompletionStreamingToolRunnerParams<T, ToolContext>,
     options?: RunnerOptions,
   ): ChatCompletionStreamingRunner<ParsedT> {
     const runner = new ChatCompletionStreamingRunner<ParsedT>(

src/lib/RunnableFunction.ts

@@ -4,15 +4,17 @@ import { JSONSchema } from './jsonschema';
 
 type PromiseOrValue<T> = T | Promise<T>;
 
-export type RunnableFunctionWithParse<Args extends object> = {
+export type RunnableFunctionWithParse<Args extends object, ToolContext = unknown> = {
   /**
    * @param args the return value from `parse`.
    * @param runner the runner evaluating this callback.
+   * @param toolContext the context passed to `.runTools()`.
    * @returns a string to send back to OpenAI.
    */
   function: (
     args: Args,
     runner: ChatCompletionRunner<unknown> | ChatCompletionStreamingRunner<unknown>,
+    toolContext: ToolContext,
   ) => PromiseOrValue<unknown>;
   /**
    * @param input the raw args from the OpenAI function call.
@@ -34,14 +36,17 @@ export type RunnableFunctionWithParse<Args extends object> = {
   strict?: boolean | undefined;
 };
 
-export type RunnableFunctionWithoutParse = {
+export type RunnableFunctionWithoutParse<ToolContext = unknown> = {
   /**
    * @param args the raw args from the OpenAI function call.
+   * @param runner the runner evaluating this callback.
+   * @param toolContext the context passed to `.runTools()`.
    * @returns a string to send back to OpenAI
    */
   function: (
     args: string,
     runner: ChatCompletionRunner<unknown> | ChatCompletionStreamingRunner<unknown>,
+    toolContext: ToolContext,
   ) => PromiseOrValue<unknown>;
   /**
    * The parameters the function accepts, describes as a JSON Schema object.
@@ -58,56 +63,80 @@ export type RunnableFunctionWithoutParse = {
   strict?: boolean | undefined;
 };
 
-export type RunnableFunction<Args extends object | string> =
-  Args extends string ? RunnableFunctionWithoutParse
-  : Args extends object ? RunnableFunctionWithParse<Args>
-  : never;
+export type RunnableFunction<Args extends object | string, ToolContext = unknown> = Args extends string ?
+  RunnableFunctionWithoutParse<ToolContext>
+: Args extends object ? RunnableFunctionWithParse<Args, ToolContext>
+: never;
 
-export type RunnableToolFunction<Args extends object | string> =
-  Args extends string ? RunnableToolFunctionWithoutParse
-  : Args extends object ? RunnableToolFunctionWithParse<Args>
-  : never;
+export type RunnableToolFunction<Args extends object | string, ToolContext = unknown> = Args extends string ?
+  RunnableToolFunctionWithoutParse<ToolContext>
+: Args extends object ? RunnableToolFunctionWithParse<Args, ToolContext>
+: never;
 
-export type RunnableToolFunctionWithoutParse = {
+export type RunnableToolFunctionWithoutParse<ToolContext = unknown> = {
+  type: 'function';
+  function: RunnableFunctionWithoutParse<ToolContext>;
+};
+export type RunnableToolFunctionWithParse<Args extends object, ToolContext = unknown> = {
   type: 'function';
-  function: RunnableFunctionWithoutParse;
+  function: RunnableFunctionWithParse<Args, ToolContext>;
 };
-export type RunnableToolFunctionWithParse<Args extends object> = {
+
+/**
+ * A broad tool shape for contextually typing callbacks when the argument type is inferred.
+ */
+export type RunnableToolFunctionWithContext<ToolContext> = {
   type: 'function';
-  function: RunnableFunctionWithParse<Args>;
+  function: {
+    function: (
+      args: any,
+      runner: ChatCompletionRunner<unknown> | ChatCompletionStreamingRunner<unknown>,
+      toolContext: ToolContext,
+    ) => PromiseOrValue<unknown>;
+    parse?: (input: string) => PromiseOrValue<any>;
+    parameters: JSONSchema;
+    description: string;
+    name?: string | undefined;
+    strict?: boolean | undefined;
+  };
 };
 
-export function isRunnableFunctionWithParse<Args extends object>(
+export function isRunnableFunctionWithParse<Args extends object, ToolContext = unknown>(
   fn: any,
-): fn is RunnableFunctionWithParse<Args> {
+): fn is RunnableFunctionWithParse<Args, ToolContext> {
   return typeof (fn as any).parse === 'function';
 }
 
 export type BaseFunctionsArgs = readonly (object | string)[];
 
-export type RunnableFunctions<FunctionsArgs extends BaseFunctionsArgs> =
-  [any[]] extends [FunctionsArgs] ? readonly RunnableFunction<any>[]
-  : {
-      [Index in keyof FunctionsArgs]: Index extends number ? RunnableFunction<FunctionsArgs[Index]>
-      : FunctionsArgs[Index];
-    };
+export type RunnableFunctions<FunctionsArgs extends BaseFunctionsArgs, ToolContext = unknown> = [
+  any[],
+] extends [FunctionsArgs] ?
+  readonly RunnableFunction<any, ToolContext>[]
+: {
+    [Index in keyof FunctionsArgs]: Index extends number ? RunnableFunction<FunctionsArgs[Index], ToolContext>
+    : FunctionsArgs[Index];
+  };
 
-export type RunnableTools<FunctionsArgs extends BaseFunctionsArgs> =
-  [any[]] extends [FunctionsArgs] ? readonly RunnableToolFunction<any>[]
-  : {
-      [Index in keyof FunctionsArgs]: Index extends number ? RunnableToolFunction<FunctionsArgs[Index]>
-      : FunctionsArgs[Index];
-    };
+export type RunnableTools<FunctionsArgs extends BaseFunctionsArgs, ToolContext = unknown> = [any[]] extends (
+  [FunctionsArgs]
+) ?
+  readonly RunnableToolFunction<any, ToolContext>[]
+: {
+    [Index in keyof FunctionsArgs]: Index extends number ?
+      RunnableToolFunction<FunctionsArgs[Index], ToolContext>
+    : FunctionsArgs[Index];
+  };
 
 /**
  * This is helper class for passing a `function` and `parse` where the `function`
  * argument type matches the `parse` return type.
  */
-export class ParsingToolFunction<Args extends object> {
+export class ParsingToolFunction<Args extends object, ToolContext = unknown> {
   type: 'function';
-  function: RunnableFunctionWithParse<Args>;
+  function: RunnableFunctionWithParse<Args, ToolContext>;
 
-  constructor(input: RunnableFunctionWithParse<Args>) {
+  constructor(input: RunnableFunctionWithParse<Args, ToolContext>) {
     this.type = 'function';
     this.function = input;
   }

src/lib/parser.ts

@@ -18,8 +18,8 @@ import { ResponseFormatJSONSchema } from '../resources/shared';
 
 type AnyChatCompletionCreateParams =
   | ChatCompletionCreateParams
-  | ChatCompletionToolRunnerParams<any>
-  | ChatCompletionStreamingToolRunnerParams<any>
+  | ChatCompletionToolRunnerParams<any, any>
+  | ChatCompletionStreamingToolRunnerParams<any, any>
   | ChatCompletionStreamParams;
 
 type Unpacked<T> = T extends (infer U)[] ? U : T;

src/resources/chat/completions/completions.ts

@@ -17,6 +17,7 @@ import { ChatCompletionStreamingRunner } from '../../../lib/ChatCompletionStream
 import { RunnerOptions } from '../../../lib/AbstractChatCompletionRunner';
 import { ChatCompletionToolRunnerParams } from '../../../lib/ChatCompletionRunner';
 import { ChatCompletionStreamingToolRunnerParams } from '../../../lib/ChatCompletionStreamingRunner';
+import { type RunnableToolFunctionWithContext } from '../../../lib/RunnableFunction';
 import { ChatCompletionStream, type ChatCompletionStreamParams } from '../../../lib/ChatCompletionStream';
 import { ExtractParsedContentFromParams, parseChatCompletion, validateInputTools } from '../../../lib/parser';
 
@@ -187,6 +188,38 @@ export class Completions extends APIResource {
    * For more details and examples, see
    * [the docs](https://github.com/openai/openai-node#automated-function-calls)
    */
+  runTools<
+    ToolContext,
+    Params extends Omit<ChatCompletionToolRunnerParams<any, ToolContext>, 'toolContext' | 'tools'> = Omit<
+      ChatCompletionToolRunnerParams<any, ToolContext>,
+      'toolContext' | 'tools'
+    >,
+    ParsedT = ExtractParsedContentFromParams<Params & ChatCompletionToolRunnerParams<any, ToolContext>>,
+  >(
+    body: Params & {
+      toolContext: ToolContext;
+      tools: readonly RunnableToolFunctionWithContext<ToolContext>[];
+    },
+    options?: RunnerOptions,
+  ): ChatCompletionRunner<ParsedT>;
+
+  runTools<
+    ToolContext,
+    Params extends Omit<
+      ChatCompletionStreamingToolRunnerParams<any, ToolContext>,
+      'toolContext' | 'tools'
+    > = Omit<ChatCompletionStreamingToolRunnerParams<any, ToolContext>, 'toolContext' | 'tools'>,
+    ParsedT = ExtractParsedContentFromParams<
+      Params & ChatCompletionStreamingToolRunnerParams<any, ToolContext>
+    >,
+  >(
+    body: Params & {
+      toolContext: ToolContext;
+      tools: readonly RunnableToolFunctionWithContext<ToolContext>[];
+    },
+    options?: RunnerOptions,
+  ): ChatCompletionStreamingRunner<ParsedT>;
+
   runTools<
     Params extends ChatCompletionToolRunnerParams<any>,
     ParsedT = ExtractParsedContentFromParams<Params>,