diff --git a/packages/server-utils/src/ai/openai/streaming.ts b/packages/server-utils/src/ai/openai/streaming.ts index 0828e755e4cb..c706111fce97 100644 --- a/packages/server-utils/src/ai/openai/streaming.ts +++ b/packages/server-utils/src/ai/openai/streaming.ts @@ -37,7 +37,9 @@ interface StreamingState { * Accumulated tool calls from Chat Completion streaming, indexed by tool call index. * @see https://platform.openai.com/docs/guides/function-calling?api-mode=chat#streaming */ - chatCompletionToolCalls: Record; + chatCompletionToolCalls: Record; + /** Key of the latest tool call, which unindexed argument fragments continue. */ + lastChatCompletionToolCallKey?: number | string; /** * Accumulated function calls from Responses API streaming. * @see https://platform.openai.com/docs/guides/function-calling?api-mode=responses#streaming @@ -56,8 +58,8 @@ interface StreamingState { */ function processChatCompletionToolCalls(toolCalls: ChatCompletionToolCall[], state: StreamingState): void { for (const toolCall of toolCalls) { - const index = toolCall.index; - if (index === undefined || !toolCall.function) continue; + if (!toolCall.function) continue; + const index = getToolCallKey(toolCall, state); // Initialize tool call if this is the first chunk for this index if (!(index in state.chatCompletionToolCalls)) { @@ -75,9 +77,33 @@ function processChatCompletionToolCalls(toolCalls: ChatCompletionToolCall[], sta existingToolCall.function.arguments += toolCall.function.arguments; } } + state.lastChatCompletionToolCallKey = index; } } +/** + * OpenAI-compatible providers such as Gemini's OpenAI endpoint send tool calls without an `index`. + * Those are matched by a known `id`; a new `id` or a function name opens a new call, and anything + * else continues the latest call. + */ +function getToolCallKey(toolCall: ChatCompletionToolCall, state: StreamingState): number | string { + if (toolCall.index !== undefined) { + return toolCall.index; + } + if (toolCall.id) { + const known = Object.keys(state.chatCompletionToolCalls).find( + key => state.chatCompletionToolCalls[key]?.id === toolCall.id, + ); + if (known !== undefined) { + return known; + } + } + if (toolCall.id || toolCall.function?.name || state.lastChatCompletionToolCallKey === undefined) { + return `no-index-${Object.keys(state.chatCompletionToolCalls).length}`; + } + return state.lastChatCompletionToolCallKey; +} + /** * Processes a single OpenAI ChatCompletionChunk event, updating the streaming state. * diff --git a/packages/server-utils/test/ai/lib/tracing/openai-streaming-tool-calls.test.ts b/packages/server-utils/test/ai/lib/tracing/openai-streaming-tool-calls.test.ts new file mode 100644 index 000000000000..f68a09e41546 --- /dev/null +++ b/packages/server-utils/test/ai/lib/tracing/openai-streaming-tool-calls.test.ts @@ -0,0 +1,103 @@ +import { describe, expect, it } from 'vitest'; +import type { Span } from '@sentry/core'; +import { GEN_AI_RESPONSE_TOOL_CALLS } from '@sentry/conventions/attributes'; +import { instrumentStream } from '../../../../src/ai/openai/streaming'; +import type { ChatCompletionChunk, ChatCompletionToolCall } from '../../../../src/ai/openai/types'; + +function createMockSpan(): { span: Span; attributes: Record } { + const attributes: Record = {}; + let isEnded = false; + const span = { + isRecording: () => !isEnded, + setAttribute: (key: string, value: unknown) => { + attributes[key] = value; + }, + setAttributes: (attrs: Record) => { + Object.assign(attributes, attrs); + }, + setStatus: () => {}, + end: () => { + isEnded = true; + }, + } as unknown as Span; + return { span, attributes }; +} + +function chunkWith(toolCalls: Array>): ChatCompletionChunk { + return { + id: 'chatcmpl-1', + object: 'chat.completion.chunk', + model: 'gemini-test', + choices: [{ index: 0, delta: { tool_calls: toolCalls } }], + } as unknown as ChatCompletionChunk; +} + +async function recordedToolCalls(chunks: ChatCompletionChunk[]): Promise>> { + const { span, attributes } = createMockSpan(); + async function* stream(): AsyncGenerator { + yield* chunks; + } + for await (const _ of instrumentStream(stream(), span, true)) { + void _; + } + return JSON.parse(attributes[GEN_AI_RESPONSE_TOOL_CALLS] as string) as Array>; +} + +describe('instrumentStream (openai chat completion tool calls)', () => { + it('keeps whole tool calls that carry an id and no index', async () => { + const toolCalls = await recordedToolCalls([ + chunkWith([ + { id: 'call_a', type: 'function', function: { name: 'lookup', arguments: '{"q":"a"}' } }, + { id: 'call_b', type: 'function', function: { name: 'search', arguments: '{"q":"b"}' } }, + ]), + ]); + + expect(toolCalls).toEqual([ + { id: 'call_a', type: 'function', function: { name: 'lookup', arguments: '{"q":"a"}' } }, + { id: 'call_b', type: 'function', function: { name: 'search', arguments: '{"q":"b"}' } }, + ]); + }); + + it('joins argument fragments by id, then onto the latest call when there is no id', async () => { + const toolCalls = await recordedToolCalls([ + chunkWith([{ id: 'call_a', type: 'function', function: { name: 'lookup', arguments: '{"q":' } }]), + chunkWith([{ id: 'call_a', function: { name: '', arguments: '"a"' } }]), + chunkWith([{ function: { name: '', arguments: '}' } }]), + ]); + + expect(toolCalls).toEqual([ + { id: 'call_a', type: 'function', function: { name: 'lookup', arguments: '{"q":"a"}' } }, + ]); + }); + + it('keeps named calls that have neither id nor index', async () => { + const toolCalls = await recordedToolCalls([ + chunkWith([{ function: { name: 'lookup', arguments: '{}' } }]), + chunkWith([{ function: { name: 'search', arguments: '{}' } }]), + ]); + + expect(toolCalls.map(call => (call.function as { name: string }).name)).toEqual(['lookup', 'search']); + }); + + it('does not let a later real index overwrite an earlier call without one', async () => { + const toolCalls = await recordedToolCalls([ + chunkWith([{ id: 'call_b', type: 'function', function: { name: 'search', arguments: '{}' } }]), + chunkWith([{ index: 0, id: 'call_a', function: { name: 'lookup', arguments: '{}' } }]), + ]); + + expect(toolCalls.map(call => call.id).sort()).toEqual(['call_a', 'call_b']); + }); + + it('still groups indexed deltas by index', async () => { + const toolCalls = await recordedToolCalls([ + chunkWith([{ index: 1, id: 'call_b', function: { name: 'search', arguments: '{' } }]), + chunkWith([{ index: 0, id: 'call_a', function: { name: 'lookup', arguments: '{}' } }]), + chunkWith([{ index: 1, function: { name: '', arguments: '}' } }]), + ]); + + expect(toolCalls).toEqual([ + expect.objectContaining({ id: 'call_a', function: { name: 'lookup', arguments: '{}' } }), + expect.objectContaining({ id: 'call_b', function: { name: 'search', arguments: '{}' } }), + ]); + }); +});