diff --git a/src/lib/AbstractChatCompletionRunner.ts b/src/lib/AbstractChatCompletionRunner.ts index c96e19d91..34b73e090 100644 --- a/src/lib/AbstractChatCompletionRunner.ts +++ b/src/lib/AbstractChatCompletionRunner.ts @@ -285,16 +285,17 @@ export class AbstractChatCompletionRunner< return await this._createChatCompletion(client, params, options); } - protected async _runTools( + protected async _runTools( client: OpenAI, params: - | ChatCompletionToolRunnerParams - | ChatCompletionStreamingToolRunnerParams, + | ChatCompletionToolRunnerParams + | ChatCompletionStreamingToolRunnerParams, runner: ChatCompletionRunner | ChatCompletionStreamingRunner, 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; }); - const functionsByName: Record> = {}; + const functionsByName: Record> = {}; 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); diff --git a/src/lib/ChatCompletionRunner.ts b/src/lib/ChatCompletionRunner.ts index 669c6b0e7..f16897555 100644 --- a/src/lib/ChatCompletionRunner.ts +++ b/src/lib/ChatCompletionRunner.ts @@ -16,20 +16,24 @@ export interface ChatCompletionRunnerEvents extends AbstractChatCompletionRunner content: (content: string) => void; } -export type ChatCompletionToolRunnerParams = Omit< - ChatCompletionCreateParamsNonStreaming, - 'tools' -> & { - tools: RunnableTools | AutoParseableTool[]; +export type ChatCompletionToolRunnerParams< + FunctionsArgs extends BaseFunctionsArgs, + ToolContext = unknown, +> = Omit & { + tools: RunnableTools | AutoParseableTool[]; + /** + * Context to pass to each tool callback during this run. + */ + toolContext?: ToolContext; }; export class ChatCompletionRunner extends AbstractChatCompletionRunner< ChatCompletionRunnerEvents, ParsedT > { - static runTools( + static runTools( client: OpenAI, - params: ChatCompletionToolRunnerParams, + params: ChatCompletionToolRunnerParams, options?: RunnerOptions, ): ChatCompletionRunner { const runner = new ChatCompletionRunner(); diff --git a/src/lib/ChatCompletionStreamingRunner.ts b/src/lib/ChatCompletionStreamingRunner.ts index 4809a3774..493d7a31f 100644 --- a/src/lib/ChatCompletionStreamingRunner.ts +++ b/src/lib/ChatCompletionStreamingRunner.ts @@ -14,11 +14,15 @@ export interface ChatCompletionStreamEvents extends AbstractChatCompletionRunner chunk: (chunk: ChatCompletionChunk, snapshot: ChatCompletionSnapshot) => void; } -export type ChatCompletionStreamingToolRunnerParams = Omit< - ChatCompletionCreateParamsStreaming, - 'tools' -> & { - tools: RunnableTools | AutoParseableTool[]; +export type ChatCompletionStreamingToolRunnerParams< + FunctionsArgs extends BaseFunctionsArgs, + ToolContext = unknown, +> = Omit & { + tools: RunnableTools | AutoParseableTool[]; + /** + * Context to pass to each tool callback during this run. + */ + toolContext?: ToolContext; }; export class ChatCompletionStreamingRunner @@ -31,9 +35,9 @@ export class ChatCompletionStreamingRunner return runner; } - static runTools( + static runTools( client: OpenAI, - params: ChatCompletionStreamingToolRunnerParams, + params: ChatCompletionStreamingToolRunnerParams, options?: RunnerOptions, ): ChatCompletionStreamingRunner { const runner = new ChatCompletionStreamingRunner( diff --git a/src/lib/RunnableFunction.ts b/src/lib/RunnableFunction.ts index d387245cd..beaed93d8 100644 --- a/src/lib/RunnableFunction.ts +++ b/src/lib/RunnableFunction.ts @@ -4,15 +4,17 @@ import { JSONSchema } from './jsonschema'; type PromiseOrValue = T | Promise; -export type RunnableFunctionWithParse = { +export type RunnableFunctionWithParse = { /** * @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 | ChatCompletionStreamingRunner, + toolContext: ToolContext, ) => PromiseOrValue; /** * @param input the raw args from the OpenAI function call. @@ -34,14 +36,17 @@ export type RunnableFunctionWithParse = { strict?: boolean | undefined; }; -export type RunnableFunctionWithoutParse = { +export type RunnableFunctionWithoutParse = { /** * @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 | ChatCompletionStreamingRunner, + toolContext: ToolContext, ) => PromiseOrValue; /** * 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 string ? RunnableFunctionWithoutParse - : Args extends object ? RunnableFunctionWithParse - : never; +export type RunnableFunction = Args extends string ? + RunnableFunctionWithoutParse +: Args extends object ? RunnableFunctionWithParse +: never; -export type RunnableToolFunction = - Args extends string ? RunnableToolFunctionWithoutParse - : Args extends object ? RunnableToolFunctionWithParse - : never; +export type RunnableToolFunction = Args extends string ? + RunnableToolFunctionWithoutParse +: Args extends object ? RunnableToolFunctionWithParse +: never; -export type RunnableToolFunctionWithoutParse = { +export type RunnableToolFunctionWithoutParse = { + type: 'function'; + function: RunnableFunctionWithoutParse; +}; +export type RunnableToolFunctionWithParse = { type: 'function'; - function: RunnableFunctionWithoutParse; + function: RunnableFunctionWithParse; }; -export type RunnableToolFunctionWithParse = { + +/** + * A broad tool shape for contextually typing callbacks when the argument type is inferred. + */ +export type RunnableToolFunctionWithContext = { type: 'function'; - function: RunnableFunctionWithParse; + function: { + function: ( + args: any, + runner: ChatCompletionRunner | ChatCompletionStreamingRunner, + toolContext: ToolContext, + ) => PromiseOrValue; + parse?: (input: string) => PromiseOrValue; + parameters: JSONSchema; + description: string; + name?: string | undefined; + strict?: boolean | undefined; + }; }; -export function isRunnableFunctionWithParse( +export function isRunnableFunctionWithParse( fn: any, -): fn is RunnableFunctionWithParse { +): fn is RunnableFunctionWithParse { return typeof (fn as any).parse === 'function'; } export type BaseFunctionsArgs = readonly (object | string)[]; -export type RunnableFunctions = - [any[]] extends [FunctionsArgs] ? readonly RunnableFunction[] - : { - [Index in keyof FunctionsArgs]: Index extends number ? RunnableFunction - : FunctionsArgs[Index]; - }; +export type RunnableFunctions = [ + any[], +] extends [FunctionsArgs] ? + readonly RunnableFunction[] +: { + [Index in keyof FunctionsArgs]: Index extends number ? RunnableFunction + : FunctionsArgs[Index]; + }; -export type RunnableTools = - [any[]] extends [FunctionsArgs] ? readonly RunnableToolFunction[] - : { - [Index in keyof FunctionsArgs]: Index extends number ? RunnableToolFunction - : FunctionsArgs[Index]; - }; +export type RunnableTools = [any[]] extends ( + [FunctionsArgs] +) ? + readonly RunnableToolFunction[] +: { + [Index in keyof FunctionsArgs]: Index extends number ? + RunnableToolFunction + : 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 { +export class ParsingToolFunction { type: 'function'; - function: RunnableFunctionWithParse; + function: RunnableFunctionWithParse; - constructor(input: RunnableFunctionWithParse) { + constructor(input: RunnableFunctionWithParse) { this.type = 'function'; this.function = input; } diff --git a/src/lib/parser.ts b/src/lib/parser.ts index 78e116073..db23cb239 100644 --- a/src/lib/parser.ts +++ b/src/lib/parser.ts @@ -18,8 +18,8 @@ import { ResponseFormatJSONSchema } from '../resources/shared'; type AnyChatCompletionCreateParams = | ChatCompletionCreateParams - | ChatCompletionToolRunnerParams - | ChatCompletionStreamingToolRunnerParams + | ChatCompletionToolRunnerParams + | ChatCompletionStreamingToolRunnerParams | ChatCompletionStreamParams; type Unpacked = T extends (infer U)[] ? U : T; diff --git a/src/resources/chat/completions/completions.ts b/src/resources/chat/completions/completions.ts index 05cae402e..375e7791f 100644 --- a/src/resources/chat/completions/completions.ts +++ b/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, 'toolContext' | 'tools'> = Omit< + ChatCompletionToolRunnerParams, + 'toolContext' | 'tools' + >, + ParsedT = ExtractParsedContentFromParams>, + >( + body: Params & { + toolContext: ToolContext; + tools: readonly RunnableToolFunctionWithContext[]; + }, + options?: RunnerOptions, + ): ChatCompletionRunner; + + runTools< + ToolContext, + Params extends Omit< + ChatCompletionStreamingToolRunnerParams, + 'toolContext' | 'tools' + > = Omit, 'toolContext' | 'tools'>, + ParsedT = ExtractParsedContentFromParams< + Params & ChatCompletionStreamingToolRunnerParams + >, + >( + body: Params & { + toolContext: ToolContext; + tools: readonly RunnableToolFunctionWithContext[]; + }, + options?: RunnerOptions, + ): ChatCompletionStreamingRunner; + runTools< Params extends ChatCompletionToolRunnerParams, ParsedT = ExtractParsedContentFromParams, diff --git a/tests/lib/ChatCompletionRunFunctions.test.ts b/tests/lib/ChatCompletionRunFunctions.test.ts index 888ffdd88..023bbf2dc 100644 --- a/tests/lib/ChatCompletionRunFunctions.test.ts +++ b/tests/lib/ChatCompletionRunFunctions.test.ts @@ -394,6 +394,26 @@ function _typeTests() { }, ], }); + openai.chat.completions.runTools({ + messages: [{ role: 'user', content: 'update my event' }], + model: 'gpt-3.5-turbo', + toolContext: { eventId: 'event_123' }, + tools: [ + { + type: 'function', + function: { + function: (_args, _runner, context) => { + const eventId: string = context.eventId; + // @ts-expect-error tool context only includes eventId + context.missing; + return eventId; + }, + parameters: {}, + description: 'updates an event', + }, + }, + ], + }); openai.chat.completions.runTools({ messages: [ { role: 'user', content: 'can you tell me how many properties are in {"a": 1, "b": 2, "c": 3}' }, @@ -662,6 +682,96 @@ describe('resource completions', () => { expect(listener.functionCallResults).toEqual([`it's raining`]); await listener.sanityCheck({ ignoredMessages: new Set([injectedMessage]) }); }); + + test('passes tool context to callbacks without sending it to the API', async () => { + const { fetch, handleRequest } = mockChatCompletionFetch(); + + const openai = new OpenAI({ apiKey: 'something1234', baseURL: 'http://127.0.0.1:4010', fetch }); + const toolContext = { eventId: 'event_123' }; + const receivedContexts: (typeof toolContext)[] = []; + const runner = openai.chat.completions.runTools({ + messages: [{ role: 'user', content: 'update my event' }], + model: 'gpt-3.5-turbo', + toolContext, + tools: [ + { + type: 'function', + function: { + function: function updateEvent(_args, _runner, context) { + receivedContexts.push(context); + return context.eventId; + }, + parameters: {}, + description: 'updates an event', + }, + }, + ], + }); + + await handleRequest(async (request) => { + expect(request).not.toHaveProperty('toolContext'); + return { + id: '1', + choices: [ + { + index: 0, + finish_reason: 'tool_calls', + logprobs: null, + message: { + role: 'assistant', + content: null, + refusal: null, + tool_calls: [ + { + type: 'function', + id: '123', + function: { + arguments: '', + name: 'updateEvent', + }, + }, + ], + }, + }, + ], + created: Math.floor(Date.now() / 1000), + model: 'gpt-3.5-turbo', + object: 'chat.completion', + }; + }); + + await handleRequest(async (request) => { + expect(request.messages[2]).toEqual({ + role: 'tool', + content: 'event_123', + tool_call_id: '123', + }); + return { + id: '2', + choices: [ + { + index: 0, + finish_reason: 'stop', + logprobs: null, + message: { + role: 'assistant', + content: 'updated', + refusal: null, + }, + }, + ], + created: Math.floor(Date.now() / 1000), + model: 'gpt-3.5-turbo', + object: 'chat.completion', + }; + }); + + await runner.done(); + + expect(receivedContexts).toEqual([toolContext]); + expect(receivedContexts[0]).toBe(toolContext); + }); + test('generates unique IDs for parallel tool calls with empty IDs', async () => { const { fetch, handleRequest } = mockChatCompletionFetch();