Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,9 @@ class MockOpenAIToolCalls {

// If stream is requested, return an async generator
if (params.stream) {
if (params.model === 'gemini-compat') {
return this._createGeminiStyleToolCallsStream(params);
}
return this._createChatCompletionToolCallsStream(params);
}

Expand Down Expand Up @@ -139,6 +142,38 @@ class MockOpenAIToolCalls {
};
}

// Gemini's OpenAI-compatible endpoint streams whole, parallel tool calls without an `index`
async *_createGeminiStyleToolCallsStream(params) {
yield {
id: 'chatcmpl-stream-gemini-tools-123',
object: 'chat.completion.chunk',
created: 1677652310,
model: params.model,
choices: [
{
index: 0,
delta: {
role: 'assistant',
tool_calls: [
{
id: 'call_paris',
type: 'function',
function: { name: 'get_weather', arguments: '{"latitude":48.8566,"longitude":2.3522}' },
},
{
id: 'call_london',
type: 'function',
function: { name: 'get_weather', arguments: '{"latitude":51.5072,"longitude":-0.1276}' },
},
],
},
finish_reason: 'tool_calls',
},
],
usage: { prompt_tokens: 15, completion_tokens: 25, total_tokens: 40 },
};
}

// Create a mock streaming response for responses API with tool calls
async *_createResponsesApiToolCallsStream(params) {
const responseId = 'resp_stream_tools_789';
Expand Down Expand Up @@ -275,6 +310,15 @@ async function run() {
});
for await (const chunk of stream1) void chunk;

// Test 2b: Chat completion streaming whole parallel tool calls without an index (Gemini style)
const geminiStream = await client.chat.completions.create({
model: 'gemini-compat',
messages: [message],
tools: [weatherTool],
stream: true,
});
for await (const chunk of geminiStream) void chunk;

// Test 3: Responses API with tools (non-streaming)
await client.responses.create({
model: 'gpt-4',
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,25 @@ describe('OpenAI Tool Calls integration', () => {
},
]);

const CHAT_GEMINI_STREAM_TOOL_CALLS = JSON.stringify([
{
id: 'call_paris',
type: 'function',
function: {
name: 'get_weather',
arguments: '{"latitude":48.8566,"longitude":2.3522}',
},
},
{
id: 'call_london',
type: 'function',
function: {
name: 'get_weather',
arguments: '{"latitude":51.5072,"longitude":-0.1276}',
},
},
]);

const RESPONSES_TOOL_CALLS = JSON.stringify([
{
type: 'function_call',
Expand All @@ -84,7 +103,7 @@ describe('OpenAI Tool Calls integration', () => {
const segment = container.items.find(span => span.is_segment && span.name === 'main');
expect(segment).toBeDefined();
const spans = container.items.filter(span => span.attributes[SENTRY_ORIGIN]?.value === 'auto.ai.openai');
expect(spans).toHaveLength(4);
expect(spans).toHaveLength(5);
const chatToolsSpan = spans.find(
span => span.attributes[GEN_AI_RESPONSE_ID]?.value === 'chatcmpl-tools-123',
);
Expand Down Expand Up @@ -193,6 +212,12 @@ describe('OpenAI Tool Calls integration', () => {
value: 40,
});

const geminiStreamSpan = spans.find(
span => span.attributes[GEN_AI_RESPONSE_ID]?.value === 'chatcmpl-stream-gemini-tools-123',
);
expect(geminiStreamSpan).toBeDefined();
expect(geminiStreamSpan!.attributes[GEN_AI_RESPONSE_TOOL_CALLS]).toBeUndefined();

const responsesToolsSpan = spans.find(
span => span.attributes[GEN_AI_RESPONSE_ID]?.value === 'resp_tools_789',
);
Expand Down Expand Up @@ -318,7 +343,7 @@ describe('OpenAI Tool Calls integration', () => {
const segment = container.items.find(span => span.is_segment && span.name === 'main');
expect(segment).toBeDefined();
const spans = container.items.filter(span => span.attributes[SENTRY_ORIGIN]?.value === 'auto.ai.openai');
expect(spans).toHaveLength(4);
expect(spans).toHaveLength(5);
const chatToolsSpan = spans.find(
span => span.attributes[GEN_AI_RESPONSE_ID]?.value === 'chatcmpl-tools-123',
);
Expand Down Expand Up @@ -452,6 +477,15 @@ describe('OpenAI Tool Calls integration', () => {
value: 40,
});

const geminiStreamSpan = spans.find(
span => span.attributes[GEN_AI_RESPONSE_ID]?.value === 'chatcmpl-stream-gemini-tools-123',
);
expect(geminiStreamSpan).toBeDefined();
expect(geminiStreamSpan!.attributes[GEN_AI_RESPONSE_TOOL_CALLS]).toEqual({
type: 'string',
value: CHAT_GEMINI_STREAM_TOOL_CALLS,
});

const responsesToolsSpan = spans.find(
span => span.attributes[GEN_AI_RESPONSE_ID]?.value === 'resp_tools_789',
);
Expand Down
45 changes: 42 additions & 3 deletions packages/server-utils/src/ai/openai/streaming.ts
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,11 @@ 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<number, ChatCompletionToolCall>;
chatCompletionToolCalls: Record<number | string, ChatCompletionToolCall>;
/** Key of the latest tool call, which unindexed argument fragments continue. */
lastChatCompletionToolCallKey?: number | string;
/** Key of the call currently open for a reported `index`, when a new `id` reused that index. */
chatCompletionToolCallKeyByIndex?: Record<number, number | string>;
/**
* Accumulated function calls from Responses API streaming.
* @see https://platform.openai.com/docs/guides/function-calling?api-mode=responses#streaming
Expand All @@ -56,8 +60,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)) {
Expand All @@ -75,9 +79,44 @@ 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`
* (or with `null`), or reuse the same `index` for every call. A call without a usable `index` is
* matched by a known `id`; a new `id` or a function name opens a new call, and anything else
* continues the latest call. A different non-empty `id` on an `index` that is already taken also
* opens a new call.
*/
function getToolCallKey(toolCall: ChatCompletionToolCall, state: StreamingState): number | string {
const calls = state.chatCompletionToolCalls;
const findById = (id: string): string | undefined => Object.keys(calls).find(key => calls[key]?.id === id);

if (typeof toolCall.index === 'number') {
const keyByIndex = (state.chatCompletionToolCallKeyByIndex ??= {});
const key = keyByIndex[toolCall.index] ?? toolCall.index;
const existingId = calls[key]?.id;
if (existingId && toolCall.id && existingId !== toolCall.id) {
const next = findById(toolCall.id) ?? `index-${toolCall.index}-${Object.keys(calls).length}`;
keyByIndex[toolCall.index] = next;
return next;
}
return key;
}
if (toolCall.id) {
const known = findById(toolCall.id);
if (known !== undefined) {
return known;
}
}
if (toolCall.id || toolCall.function?.name || state.lastChatCompletionToolCallKey === undefined) {
return `no-index-${Object.keys(calls).length}`;
}
return state.lastChatCompletionToolCallKey;
}

/**
* Processes a single OpenAI ChatCompletionChunk event, updating the streaming state.
*
Expand Down
2 changes: 1 addition & 1 deletion packages/server-utils/src/ai/openai/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -176,7 +176,7 @@ interface ResponseOutputItemDoneEvent {
* Tool call object for Chat Completion streaming
*/
export interface ChatCompletionToolCall {
index?: number; // Present for streaming responses
index?: number | null; // Present for streaming responses; some compatible providers send null
id: string;
type?: string; // Could be missing for streaming responses
function?: {
Expand Down

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

could we also get a node-integration test for this? a Gemini style stream in dev-packages/node-integration-tests/suites/tracing/openai/openai-tool-calls/scenario.mjs (whole parallel calls, no index) so it goes through the real openai SDK, the unit tests call instrumentStream directly

Original file line number Diff line number Diff line change
@@ -0,0 +1,150 @@
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<string, unknown> } {
const attributes: Record<string, unknown> = {};
let isEnded = false;
const span = {
isRecording: () => !isEnded,
setAttribute: (key: string, value: unknown) => {
attributes[key] = value;
},
setAttributes: (attrs: Record<string, unknown>) => {
Object.assign(attributes, attrs);
},
setStatus: () => {},
end: () => {
isEnded = true;
},
} as unknown as Span;
return { span, attributes };
}

function chunkWith(toolCalls: Array<Partial<ChatCompletionToolCall>>): 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<Array<Record<string, unknown>>> {
const { span, attributes } = createMockSpan();
async function* stream(): AsyncGenerator<ChatCompletionChunk> {
yield* chunks;
}
for await (const _ of instrumentStream(stream(), span, true)) {
void _;
}
return JSON.parse(attributes[GEN_AI_RESPONSE_TOOL_CALLS] as string) as Array<Record<string, unknown>>;
}

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: '{}' } }),
]);
});

it('treats a null index like a missing one', async () => {
const toolCalls = await recordedToolCalls([
chunkWith([
{ index: null, id: 'call_a', type: 'function', function: { name: 'lookup', arguments: '{"q":"a"}' } },
{ index: null, id: 'call_b', type: 'function', function: { name: 'search', arguments: '{"q":"b"}' } },
]),
chunkWith([{ index: null, function: { name: '', arguments: ' ' } }]),
]);

expect(toolCalls.map(call => call.id)).toEqual(['call_a', 'call_b']);
expect((toolCalls[1]?.function as { arguments: string }).arguments).toBe('{"q":"b"} ');
});

it('opens a new call when a different id lands on an index that is already taken', async () => {
const toolCalls = await recordedToolCalls([
chunkWith([
{ index: 0, id: 'call_a', type: 'function', function: { name: 'lookup', arguments: '{"q":"a"}' } },
{ index: 0, id: 'call_b', type: 'function', function: { name: 'search', arguments: '{"q":"b"}' } },
]),
]);

expect(toolCalls).toEqual([
{ index: 0, id: 'call_a', type: 'function', function: { name: 'lookup', arguments: '{"q":"a"}' } },
{ index: 0, id: 'call_b', type: 'function', function: { name: 'search', arguments: '{"q":"b"}' } },
]);
});

it('continues the latest call opened on a reused index when the next delta has no id', async () => {
const toolCalls = await recordedToolCalls([
chunkWith([{ index: 0, id: 'call_a', function: { name: 'lookup', arguments: '{}' } }]),
chunkWith([{ index: 0, id: 'call_b', function: { name: 'search', arguments: '{"q":' } }]),
chunkWith([{ index: 0, function: { name: '', arguments: '"b"}' } }]),
]);

expect(toolCalls.map(call => (call.function as { arguments: string }).arguments)).toEqual(['{}', '{"q":"b"}']);
});

it('does not split a call whose id is repeated on every delta of its index', async () => {
const toolCalls = await recordedToolCalls([
chunkWith([{ index: 0, id: 'call_a', function: { name: 'lookup', arguments: '{"q":' } }]),
chunkWith([{ index: 0, id: 'call_a', function: { name: '', arguments: '"a"}' } }]),
]);

expect(toolCalls).toHaveLength(1);
expect((toolCalls[0]?.function as { arguments: string }).arguments).toBe('{"q":"a"}');
});
});
Loading