diff --git a/.changeset/quiet-owls-count.md b/.changeset/quiet-owls-count.md new file mode 100644 index 000000000..04ec1f96b --- /dev/null +++ b/.changeset/quiet-owls-count.md @@ -0,0 +1,5 @@ +--- +"braintrust": patch +--- + +fix(openai,openrouter,huggingface): keep streamed tool calls apart when the provider omits `index` diff --git a/js/src/instrumentation/plugins/huggingface-plugin.test.ts b/js/src/instrumentation/plugins/huggingface-plugin.test.ts index 8b2b80286..2612e47e2 100644 --- a/js/src/instrumentation/plugins/huggingface-plugin.test.ts +++ b/js/src/instrumentation/plugins/huggingface-plugin.test.ts @@ -89,6 +89,98 @@ describe("aggregateChatCompletionChunks", () => { }); }); + it("keeps tool calls apart when the provider omits index", () => { + const chunk = (toolCall: Record) => ({ + choices: [{ delta: { tool_calls: [toolCall] } }], + }); + const result = aggregateChatCompletionChunks([ + chunk({ + id: "call_1", + type: "function", + function: { name: "get_weather", arguments: '{"city":' }, + }), + chunk({ function: { arguments: '"Paris"}' } }), + chunk({ + id: "call_2", + type: "function", + function: { name: "get_time", arguments: '{"tz":"CET"}' }, + }), + ]); + + expect(result).toMatchObject({ + choices: [ + { + message: { + tool_calls: [ + { + id: "call_1", + type: "function", + function: { + name: "get_weather", + arguments: '{"city":"Paris"}', + }, + }, + { + id: "call_2", + type: "function", + function: { name: "get_time", arguments: '{"tz":"CET"}' }, + }, + ], + }, + }, + ], + }); + }); + + it.each([undefined, null])( + "merges interleaved tool calls by ID when index is %s", + (index) => { + const chunks = [ + { + id: "call_1", + type: "function", + function: { name: "get_weather", arguments: '{"city":' }, + }, + { + id: "call_2", + type: "function", + function: { name: "get_time", arguments: '{"tz":' }, + }, + { id: "call_1", function: { arguments: '"Paris"}' } }, + { id: "call_2", function: { arguments: '"CET"}' } }, + ].map((toolCall) => ({ + choices: [ + { index: 0, delta: { tool_calls: [{ ...toolCall, index }] } }, + ], + })); + const result = aggregateChatCompletionChunks(chunks); + + expect(result).toMatchObject({ + choices: [ + { + message: { + tool_calls: [ + { + id: "call_1", + type: "function", + function: { + name: "get_weather", + arguments: '{"city":"Paris"}', + }, + }, + { + id: "call_2", + type: "function", + function: { name: "get_time", arguments: '{"tz":"CET"}' }, + }, + ], + }, + }, + ], + }); + }, + ); + it("preserves full tool calls emitted in a single streamed chunk", () => { expect( aggregateChatCompletionChunks([ diff --git a/js/src/instrumentation/plugins/huggingface-plugin.ts b/js/src/instrumentation/plugins/huggingface-plugin.ts index cbbcec366..830c12525 100644 --- a/js/src/instrumentation/plugins/huggingface-plugin.ts +++ b/js/src/instrumentation/plugins/huggingface-plugin.ts @@ -560,7 +560,7 @@ function mergeChatToolCallDeltas( const toolIndex = typeof toolDelta.index === "number" && toolDelta.index >= 0 ? toolDelta.index - : 0; + : resolveMissingToolCallIndex(toolCallsByIndex, toolDelta); const existing = toolCallsByIndex.get(toolIndex); if (!existing) { @@ -572,6 +572,31 @@ function mergeChatToolCallDeltas( } } +// Some providers omit `index`: match existing ids before starting a new call. +// A delta without an id continues the latest one. +function resolveMissingToolCallIndex( + toolCallsByIndex: Map>, + toolDelta: Record, +): number { + if (typeof toolDelta.id === "string") { + for (const [index, toolCall] of toolCallsByIndex) { + if (toolCall.id === toolDelta.id) { + return index; + } + } + } + if (toolCallsByIndex.size === 0) { + return 0; + } + const lastIndex = Math.max(...toolCallsByIndex.keys()); + const lastId = toolCallsByIndex.get(lastIndex)?.id; + return typeof toolDelta.id === "string" && + typeof lastId === "string" && + lastId !== toolDelta.id + ? lastIndex + 1 + : lastIndex; +} + function createChatToolCall( toolDelta: Record, ): Record { diff --git a/js/src/instrumentation/plugins/openai-plugin.test.ts b/js/src/instrumentation/plugins/openai-plugin.test.ts index 8d0764cf4..86d31d15e 100644 --- a/js/src/instrumentation/plugins/openai-plugin.test.ts +++ b/js/src/instrumentation/plugins/openai-plugin.test.ts @@ -1028,6 +1028,114 @@ describe("aggregateChatCompletionChunks", () => { '{"a":1}', ); }); + + it("should keep parallel tool calls apart when the provider omits index", () => { + const toolCallChunk = (toolCall: Record) => ({ + choices: [{ index: 0, delta: { tool_calls: [toolCall] } }], + }); + const chunks = [ + toolCallChunk({ + id: "call_1", + type: "function", + function: { name: "get_weather", arguments: '{"city":' }, + }), + toolCallChunk({ function: { arguments: '"Paris"}' } }), + toolCallChunk({ + id: "call_2", + type: "function", + function: { name: "get_time", arguments: '{"tz":"CET"}' }, + }), + ]; + + const result = aggregateChatCompletionChunks(chunks as any); + const toolCalls = (result.output as any[])[0].message.tool_calls; + + expect(toolCalls).toEqual([ + { + id: "call_1", + type: "function", + function: { name: "get_weather", arguments: '{"city":"Paris"}' }, + }, + { + id: "call_2", + type: "function", + function: { name: "get_time", arguments: '{"tz":"CET"}' }, + }, + ]); + }); + + it.each([undefined, null])( + "merges interleaved tool calls by ID when index is %s", + (index) => { + const chunks = [ + { + id: "call_1", + type: "function", + function: { name: "get_weather", arguments: '{"city":' }, + }, + { + id: "call_2", + type: "function", + function: { name: "get_time", arguments: '{"tz":' }, + }, + { id: "call_1", function: { arguments: '"Paris"}' } }, + { id: "call_2", function: { arguments: '"CET"}' } }, + ].map((toolCall) => ({ + choices: [ + { index: 0, delta: { tool_calls: [{ ...toolCall, index }] } }, + ], + })); + // @ts-expect-error Compatible providers may omit or null the required index. + const result = aggregateChatCompletionChunks(chunks); + + expect(result.output[0].message.tool_calls).toMatchObject([ + { + id: "call_1", + type: "function", + function: { name: "get_weather", arguments: '{"city":"Paris"}' }, + }, + { + id: "call_2", + type: "function", + function: { name: "get_time", arguments: '{"tz":"CET"}' }, + }, + ]); + }, + ); + + it("should keep every tool call when one chunk carries several without index", () => { + const chunks = [ + { + choices: [ + { + index: 0, + delta: { + tool_calls: [ + { + id: "call_1", + type: "function", + function: { name: "a", arguments: "{}" }, + }, + { + id: "call_2", + type: "function", + function: { name: "b", arguments: "{}" }, + }, + ], + }, + }, + ], + }, + ]; + + const result = aggregateChatCompletionChunks(chunks as any); + const toolCalls = (result.output as any[])[0].message.tool_calls; + + expect(toolCalls.map((toolCall: any) => toolCall.id)).toEqual([ + "call_1", + "call_2", + ]); + }); }); describe("finish reason", () => { diff --git a/js/src/instrumentation/plugins/openai-plugin.ts b/js/src/instrumentation/plugins/openai-plugin.ts index 7399184d8..7db400678 100644 --- a/js/src/instrumentation/plugins/openai-plugin.ts +++ b/js/src/instrumentation/plugins/openai-plugin.ts @@ -640,6 +640,33 @@ function toChatChoice(choice: AggregatedChatChoice): OpenAIChatChoice { }; } +// Some OpenAI-compatible providers omit `index` on streamed tool call deltas. +// Match existing ids before starting a new call; a delta without an id +// continues the latest one. +function resolveToolCallIndex( + toolCallsByIndex: AggregatedChatChoice["toolCallsByIndex"], + toolDelta: { index?: number | null; id?: string }, +): number { + if (typeof toolDelta.index === "number") { + return toolDelta.index; + } + if (toolDelta.id) { + for (const [index, toolCall] of toolCallsByIndex) { + if (toolCall.id === toolDelta.id) { + return index; + } + } + } + if (toolCallsByIndex.size === 0) { + return 0; + } + const lastIndex = Math.max(...toolCallsByIndex.keys()); + const lastId = toolCallsByIndex.get(lastIndex)?.id; + return toolDelta.id && lastId && lastId !== toolDelta.id + ? lastIndex + 1 + : lastIndex; +} + /** * Aggregate chat completion chunks into a single response. * Combines role (first), content, reasoning_content and audio transcripts (concatenated), @@ -740,15 +767,18 @@ export function aggregateChatCompletionChunks( if (delta.tool_calls) { for (const toolDelta of delta.tool_calls) { - let aggregatedToolCall = aggregatedChoice.toolCallsByIndex.get( - toolDelta.index, + const toolIndex = resolveToolCallIndex( + aggregatedChoice.toolCallsByIndex, + toolDelta, ); + let aggregatedToolCall = + aggregatedChoice.toolCallsByIndex.get(toolIndex); if (!aggregatedToolCall) { aggregatedToolCall = { function: { arguments: "" }, }; aggregatedChoice.toolCallsByIndex.set( - toolDelta.index, + toolIndex, aggregatedToolCall, ); } diff --git a/js/src/instrumentation/plugins/openrouter-plugin.test.ts b/js/src/instrumentation/plugins/openrouter-plugin.test.ts index 4acfc7edc..6cd52838b 100644 --- a/js/src/instrumentation/plugins/openrouter-plugin.test.ts +++ b/js/src/instrumentation/plugins/openrouter-plugin.test.ts @@ -196,6 +196,75 @@ describe("OpenRouter Plugin", () => { }); }); + it("should keep tool calls apart when the provider omits index", () => { + const chunk = (toolCall: Record) => ({ + choices: [{ delta: { tool_calls: [toolCall] } }], + }); + const result = aggregateOpenRouterChatChunks([ + chunk({ + id: "call_1", + type: "function", + function: { name: "lookup", arguments: '{"city":' }, + }), + chunk({ function: { arguments: '"Vienna"}' } }), + chunk({ + id: "call_2", + type: "function", + function: { name: "weather", arguments: '{"tz":"CET"}' }, + }), + ]); + + expect( + (result.output[0].message?.tool_calls as any[]).map((toolCall) => [ + toolCall.id, + toolCall.function.name, + toolCall.function.arguments, + ]), + ).toEqual([ + ["call_1", "lookup", '{"city":"Vienna"}'], + ["call_2", "weather", '{"tz":"CET"}'], + ]); + }); + + it.each([undefined, null])( + "merges interleaved tool calls by ID when index is %s", + (index) => { + const chunks = [ + { + id: "call_1", + type: "function", + function: { name: "get_weather", arguments: '{"city":' }, + }, + { + id: "call_2", + type: "function", + function: { name: "get_time", arguments: '{"tz":' }, + }, + { id: "call_1", function: { arguments: '"Paris"}' } }, + { id: "call_2", function: { arguments: '"CET"}' } }, + ].map((toolCall) => ({ + choices: [ + { index: 0, delta: { tool_calls: [{ ...toolCall, index }] } }, + ], + })); + // @ts-expect-error Providers may emit null indexes despite the SDK type. + const result = aggregateOpenRouterChatChunks(chunks); + + expect(result.output[0].message?.tool_calls).toMatchObject([ + { + id: "call_1", + type: "function", + function: { name: "get_weather", arguments: '{"city":"Paris"}' }, + }, + { + id: "call_2", + type: "function", + function: { name: "get_time", arguments: '{"tz":"CET"}' }, + }, + ]); + }, + ); + it("should return an empty assistant message for empty chunks", () => { expect(aggregateOpenRouterChatChunks([])).toEqual({ output: [ diff --git a/js/src/instrumentation/plugins/openrouter-plugin.ts b/js/src/instrumentation/plugins/openrouter-plugin.ts index 98db68801..bf85143a6 100644 --- a/js/src/instrumentation/plugins/openrouter-plugin.ts +++ b/js/src/instrumentation/plugins/openrouter-plugin.ts @@ -944,7 +944,27 @@ export function aggregateOpenRouterChatChunks( continue; } - const toolIndex = toolDelta.index ?? 0; + // Some providers omit `index`: match existing ids before starting a new + // call. A delta without an id continues the latest one. + let toolIndex = toolDelta.index; + if (toolIndex == null && toolDelta.id) { + const existingIndex = toolCalls?.findIndex( + (toolCall) => toolCall?.id === toolDelta.id, + ); + if (existingIndex !== undefined && existingIndex >= 0) { + toolIndex = existingIndex; + } + } + const lastToolIndex = (toolCalls?.length ?? 0) - 1; + toolIndex ??= + lastToolIndex >= 0 && + !( + toolDelta.id && + toolCalls?.[lastToolIndex]?.id && + toolCalls[lastToolIndex].id !== toolDelta.id + ) + ? lastToolIndex + : lastToolIndex + 1; const existingToolCall = toolCalls?.[toolIndex]; if (