Skip to content

Commit f2834e9

Browse files
404Wolf404Wolf
committed
fix more tests
1 parent c4f95d4 commit f2834e9

File tree

1 file changed

+73
-44
lines changed

1 file changed

+73
-44
lines changed

‎tests/lib/tools/ToolRunner.test.ts‎

Lines changed: 73 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,15 @@ import OpenAI from 'openai';
55
// import { Fetch } from 'openai/sdk/internal/builtin-types';
66
import { mockFetch } from '../../utils/mock-fetch';
77
import { BetaRunnableTool } from 'openai/lib/beta/BetaRunnableTool';
8-
import { ChatCompletion, ChatCompletionChunk, ChatCompletionMessage, ChatCompletionTool, ChatCompletionToolMessageParam } from 'openai/resources';
8+
import {
9+
ChatCompletion,
10+
ChatCompletionChunk,
11+
ChatCompletionFunctionTool,
12+
ChatCompletionMessage,
13+
ChatCompletionMessageToolCall,
14+
ChatCompletionTool,
15+
ChatCompletionToolMessageParam,
16+
} from 'openai/resources';
917
import { Fetch } from 'openai/internal/builtin-types';
1018
import { ChatCompletionStream } from 'openai/lib/ChatCompletionStream';
1119

@@ -48,25 +56,34 @@ const calculatorTool: BetaRunnableTool<{ a: number; b: number; operation: string
4856
};
4957

5058
// Helper functions to create content blocks
51-
function getWeatherToolUse(location: string, id: string = 'tool_1'): ChatCompletionFunctionTool {
52-
return { type: 'tool_use', id, name: 'getWeather', input: { location } };
59+
function getWeatherToolUse(location: string, id: string = 'tool_1'): ChatCompletionMessageToolCall {
60+
return {
61+
id: id,
62+
type: 'function',
63+
function: {
64+
name: 'getWeather',
65+
arguments: JSON.stringify({ location }),
66+
},
67+
};
5368
}
5469

5570
function getWeatherToolResult(location: string, id: string = 'tool_1'): ChatCompletionToolMessageParam {
56-
return { role: 'tool', tool_use_id: id, content: `Sunny in ${location}` };
71+
return { role: 'tool', tool_call_id: id, content: `Sunny in ${location}` };
5772
}
5873

5974
function getCalculatorToolUse(
6075
a: number,
6176
b: number,
6277
operation: string,
6378
id: string = 'tool_2',
64-
): BetaContentBlock {
79+
): ChatCompletionMessageToolCall {
6580
return {
66-
type: 'tool_use',
67-
id,
68-
name: 'calculate',
69-
input: { a, b, operation },
81+
id: id,
82+
type: 'function',
83+
function: {
84+
name: 'calculate',
85+
arguments: JSON.stringify({ a, b, operation }),
86+
},
7087
};
7188
}
7289

@@ -75,7 +92,7 @@ function getCalculatorToolResult(
7592
b: number,
7693
operation: string,
7794
id: string = 'tool_2',
78-
): BetaToolResultBlockParam {
95+
): ChatCompletionToolMessageParam {
7996
let result: string;
8097
if (operation === 'add') {
8198
result = String(a + b);
@@ -85,13 +102,13 @@ function getCalculatorToolResult(
85102
result = `Error: Unknown operation: ${operation}`;
86103
}
87104
return {
88-
type: 'tool_result',
89-
tool_use_id: id,
105+
role: 'tool',
106+
tool_call_id: id,
90107
content: result,
91108
};
92109
}
93110

94-
function getTextContent(text?: string): BetaContentBlock {
111+
function getTextContent(text?: string): ChatCompletionMessage {
95112
return {
96113
type: 'text',
97114
text: text || 'Some text content',
@@ -213,21 +230,34 @@ interface SetupTestResult<Stream extends boolean> {
213230
runner: OpenAI.Beta.Chat.Completions.BetaToolRunner<Stream>;
214231
fetch: ReturnType<typeof mockFetch>['fetch'];
215232
handleRequest: (fetch: Fetch) => void;
216-
handleAssistantMessage: (...content: ChatCompletionMessage[]) => ChatCompletion;
233+
handleAssistantMessage: (firstChoice: ToolCallsOrMessage) => ChatCompletion;
217234
handleAssistantMessageStream: (toolCalls?: ChatCompletionChunk[]) => ChatCompletionStream;
218235
}
219236

220237
type ToolRunnerParams = Parameters<typeof OpenAI.Beta.Chat.Completions.prototype.toolRunner>[0];
221238

239+
type ToolCallsOrMessage = ChatCompletionMessageToolCall[] | ChatCompletionMessage;
240+
222241
function setupTest(params?: Partial<ToolRunnerParams> & { stream?: false }): SetupTestResult<false>;
223242
function setupTest(params: Partial<ToolRunnerParams> & { stream: true }): SetupTestResult<true>;
224243
function setupTest(params: Partial<ToolRunnerParams> = {}): SetupTestResult<boolean> {
225244
const { handleRequest, handleStreamEvents, fetch } = mockFetch();
226245
let messageIdCounter = 0;
227246

228-
const handleAssistantMessage: SetupTestResult<false>['handleAssistantMessage'] = (...content) => {
229-
const hasToolUse = content.some((block) => (block.tool_calls?.length ?? 0) > 0);
230-
const stop_reason = hasToolUse ? 'tool_use' : 'end_turn';
247+
const handleAssistantMessage: SetupTestResult<false>['handleAssistantMessage'] = (
248+
messageContentOrToolCalls,
249+
) => {
250+
const isToolCalls = Array.isArray(messageContentOrToolCalls);
251+
252+
const messageContent =
253+
isToolCalls ?
254+
{
255+
role: 'assistant' as const,
256+
tool_calls: messageContentOrToolCalls,
257+
refusal: null,
258+
content: null,
259+
}
260+
: (messageContentOrToolCalls as ChatCompletionMessage); // TODO: check that this is right
231261

232262
const message: ChatCompletion = {
233263
id: `msg_${messageIdCounter++}`,
@@ -237,12 +267,8 @@ function setupTest(params: Partial<ToolRunnerParams> = {}): SetupTestResult<bool
237267
choices: [
238268
{
239269
index: 0,
240-
message: {
241-
role: 'assistant',
242-
content: JSON.stringify(content),
243-
refusal: null,
244-
},
245-
finish_reason: stop_reason === 'tool_use' ? 'tool_calls' : 'stop',
270+
message: messageContent,
271+
finish_reason: 'stop',
246272
logprobs: null,
247273
},
248274
],
@@ -252,6 +278,7 @@ function setupTest(params: Partial<ToolRunnerParams> = {}): SetupTestResult<bool
252278
total_tokens: 30,
253279
},
254280
};
281+
255282
handleRequest(async () => {
256283
return new Response(JSON.stringify(message), {
257284
status: 200,
@@ -261,11 +288,9 @@ function setupTest(params: Partial<ToolRunnerParams> = {}): SetupTestResult<bool
261288
return message;
262289
};
263290

264-
const handleAssistantMessageStream: SetupTestResult<true>['handleAssistantMessageStream'] = (
265-
...content
266-
) => {
267-
const hasToolUse = content.some(
268-
(block) => (block?.at(0)?.choices?.at(0)?.delta?.tool_calls?.length ?? 0) > 0,
291+
const handleAssistantMessageStream: SetupTestResult<true>['handleAssistantMessageStream'] = (...chunks) => {
292+
const hasToolUse = chunks.some(
293+
(chunk) => (chunk?.at(0)?.choices?.at(0)?.delta?.tool_calls?.length ?? 0) > 0,
269294
);
270295
const stop_reason = hasToolUse ? 'tool_use' : 'end_turn';
271296

@@ -279,7 +304,7 @@ function setupTest(params: Partial<ToolRunnerParams> = {}): SetupTestResult<bool
279304
index: 0,
280305
message: {
281306
role: 'assistant',
282-
content: JSON.stringify(content),
307+
content: JSON.stringify(chunks),
283308
refusal: null,
284309
},
285310
finish_reason: stop_reason === 'tool_use' ? 'tool_calls' : 'stop',
@@ -348,19 +373,21 @@ describe('ToolRunner', () => {
348373
});
349374

350375
describe('iterator.next()', () => {
351-
it('yields BetaMessage', async () => {
376+
it('yields CompletionMessage', async () => {
352377
const { runner, handleAssistantMessage } = setupTest();
353378

354379
const iterator = runner[Symbol.asyncIterator]();
355380

356-
handleAssistantMessage(getWeatherToolUse('SF'));
381+
handleAssistantMessage([getWeatherToolUse('SF')]);
382+
357383
await expectEvent(iterator, (message) => {
358-
expect(message.content).toMatchObject([getWeatherToolUse('SF')]);
384+
expect(message?.choices[0]?.message.tool_calls).toMatchObject([getWeatherToolUse('SF')]);
359385
});
360386

361387
handleAssistantMessage(getTextContent());
388+
362389
await expectEvent(iterator, (message) => {
363-
expect(message.content).toMatchObject([getTextContent()]);
390+
expect(message?.choices[0]?.message).toMatchObject(getTextContent());
364391
});
365392

366393
await expectDone(iterator);
@@ -411,32 +438,34 @@ describe('ToolRunner', () => {
411438

412439
const iterator = runner[Symbol.asyncIterator]();
413440

414-
handleAssistantMessage(getWeatherToolUse('NYC'), getCalculatorToolUse(2, 3, 'add'));
441+
handleAssistantMessage([getWeatherToolUse('NYC'), getCalculatorToolUse(2, 3, 'add')]);
415442
await expectEvent(iterator, (message) => {
416-
expect(message.content).toHaveLength(2);
417-
expect(message.content).toMatchObject([getWeatherToolUse('NYC'), getCalculatorToolUse(2, 3, 'add')]);
443+
expect(message?.choices).toHaveLength(1);
444+
expect(message?.choices[0]?.message.tool_calls).toHaveLength(2);
445+
expect(message?.choices[0]?.message.tool_calls).toMatchObject([
446+
getWeatherToolUse('NYC'),
447+
getCalculatorToolUse(2, 3, 'add'),
448+
]);
418449
});
419450

420451
// Assistant provides final response
421452
handleAssistantMessage(getTextContent());
422453
await expectEvent(iterator, (message) => {
423-
expect(message.content).toMatchObject([getTextContent()]);
454+
expect(message?.choices).toHaveLength(1);
455+
expect(message?.choices[0]?.message).toMatchObject(getTextContent());
424456
});
425457

426458
// Check that we have both tool results in the messages
427459
// Second message should be assistant with tool uses
428460
// Third message should be user with both tool results
429461
const messages = runner.params.messages;
430-
expect(messages).toHaveLength(3); // user message, assistant with tools, user with results
462+
expect(messages).toHaveLength(4); // user message, assistant with tools, tool result 1, tool result 2, assistant final
431463
expect(messages[1]).toMatchObject({
432464
role: 'assistant',
433-
content: [getWeatherToolUse('NYC'), getCalculatorToolUse(2, 3, 'add')],
465+
tool_calls: [getWeatherToolUse('NYC'), getCalculatorToolUse(2, 3, 'add')],
434466
});
435-
expect(messages[2]).toMatchObject({
436-
role: 'user',
437-
content: [getWeatherToolResult('NYC'), getCalculatorToolResult(2, 3, 'add', 'tool_2')],
438-
});
439-
467+
expect(messages[2]).toMatchObject(getWeatherToolResult('NYC'));
468+
expect(messages[3]).toMatchObject(getCalculatorToolResult(2, 3, 'add', 'tool_2'));
440469
await expectDone(iterator);
441470
});
442471

0 commit comments

Comments
 (0)