From 5b9e5da08b27644f999027ef12d0066b2069a275 Mon Sep 17 00:00:00 2001 From: Alex Castillo Date: Fri, 16 Oct 2026 16:47:36 -0400 Subject: [PATCH 1/2] fix(openai): keep every streamed tool call when a chunk carries several The chat completion stream postprocessor read only delta.tool_calls[0] and appended argument fragments to the last call opened. Providers that stream parallel tool calls in one delta (OpenAI-compatible endpoints such as Gemini's, Together, vLLM) lost every call after the first in the logged span, and fragments interleaved across indexes went to the wrong call. Walk every entry of delta.tool_calls and continue the call named by its index, falling back to the last call when a provider omits index. A new id still opens a new call, so providers that reuse index 0 keep working. Co-Authored-By: Claude Opus 5.5 --- .../integrations/openai/test_openai.py | 99 +++++++++++++++++++ .../braintrust/integrations/openai/tracing.py | 51 +++++----- 2 files changed, 126 insertions(+), 24 deletions(-) diff --git a/py/src/braintrust/integrations/openai/test_openai.py b/py/src/braintrust/integrations/openai/test_openai.py index 4f5d592d..1adc711e 100644 --- a/py/src/braintrust/integrations/openai/test_openai.py +++ b/py/src/braintrust/integrations/openai/test_openai.py @@ -2268,6 +2268,29 @@ def test_wrap_openai_and_setup_use_same_wrappers(): ) +def _tool_call_delta(index, *, call_id=None, name=None, arguments=None): + function = {key: value for key, value in (("name", name), ("arguments", arguments)) if value is not None} + delta = {"index": index, "function": function} + if call_id is not None: + delta.update(id=call_id, type="function") + return delta + + +def _tool_call_chunk(tool_calls, *, role=None, finish_reason=None): + delta = {"tool_calls": tool_calls} if tool_calls else {} + if role is not None: + delta["role"] = role + return openai.types.chat.ChatCompletionChunk.model_validate( + { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 0, + "model": "gemini-3.5-flash", + "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], + } + ) + + class TestZAICompatibleOpenAI: """Tests for validating some ZAI compatibility with OpenAI wrapper.""" @@ -2367,6 +2390,82 @@ def test_chat_completion_streaming_none_arguments(self, memory_logger): # No spans should be generated from this unit test assert not memory_logger.pop() + def test_chat_completion_streaming_parallel_tool_calls_in_one_chunk(self, memory_logger): + """Several tool calls in one delta (Gemini's OpenAI-compatible API, Together, vLLM) are all kept.""" + assert not memory_logger.pop() + + all_results = [ + _tool_call_chunk( + [ + _tool_call_delta(0, call_id="call_1", name="get_weather", arguments='{"city": "Paris"}'), + _tool_call_delta(1, call_id="call_2", name="get_weather", arguments='{"city": "Rome"}'), + ], + role="assistant", + ), + _tool_call_chunk([], finish_reason="tool_calls"), + ] + + wrapper = ChatCompletionWrapper(None, None) + message = wrapper._postprocess_streaming_results(all_results)["output"][0]["message"] + + assert message["tool_calls"] == [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + }, + {"id": "call_2", "type": "function", "function": {"name": "get_weather", "arguments": '{"city": "Rome"}'}}, + ] + assert not memory_logger.pop() + + def test_chat_completion_streaming_interleaved_tool_call_arguments(self, memory_logger): + """Argument fragments go to the call named by their index, not to the last call opened.""" + assert not memory_logger.pop() + + all_results = [ + _tool_call_chunk( + [_tool_call_delta(0, call_id="call_1", name="get_weather", arguments="")], role="assistant" + ), + _tool_call_chunk([_tool_call_delta(1, call_id="call_2", name="get_time", arguments="")]), + _tool_call_chunk([_tool_call_delta(0, arguments='{"city": ')]), + _tool_call_chunk([_tool_call_delta(1, arguments='{"tz": ')]), + _tool_call_chunk([_tool_call_delta(0, arguments='"Paris"}'), _tool_call_delta(1, arguments='"CET"}')]), + _tool_call_chunk([], finish_reason="tool_calls"), + ] + + wrapper = ChatCompletionWrapper(None, None) + message = wrapper._postprocess_streaming_results(all_results)["output"][0]["message"] + + assert 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"}'}}, + ] + assert not memory_logger.pop() + + def test_chat_completion_streaming_new_tool_call_id_under_reused_index(self, memory_logger): + """Some providers send every call at index 0; a new id still starts a new call.""" + assert not memory_logger.pop() + + all_results = [ + _tool_call_chunk( + [_tool_call_delta(0, call_id="call_1", name="get_weather", arguments='{"city": "Paris"}')], + role="assistant", + ), + _tool_call_chunk([_tool_call_delta(0, call_id="call_2", name="get_time", arguments='{"tz": "CET"}')]), + _tool_call_chunk([], finish_reason="tool_calls"), + ] + + wrapper = ChatCompletionWrapper(None, None) + message = wrapper._postprocess_streaming_results(all_results)["output"][0]["message"] + + assert [tool_call["id"] for tool_call in message["tool_calls"]] == ["call_1", "call_2"] + assert message["tool_calls"][1]["function"] == {"name": "get_time", "arguments": '{"tz": "CET"}'} + assert not memory_logger.pop() + def test_chat_completion_streaming_audio_is_materialized_as_attachment(self, memory_logger): assert not memory_logger.pop() diff --git a/py/src/braintrust/integrations/openai/tracing.py b/py/src/braintrust/integrations/openai/tracing.py index 3d4538ee..12febe50 100644 --- a/py/src/braintrust/integrations/openai/tracing.py +++ b/py/src/braintrust/integrations/openai/tracing.py @@ -755,6 +755,7 @@ def _postprocess_streaming_results( refusal = None audio: dict[str, Any] | None = None tool_calls: list[Any] | None = None + tool_calls_by_index: dict[int, dict[str, Any]] = {} finish_reason = None logprobs_content: list[Any] | None = None logprobs_refusal: list[Any] | None = None @@ -829,34 +830,36 @@ def _postprocess_streaming_results( if delta_audio_expires_at is not None: audio["expires_at"] = delta_audio_expires_at - if _get_attr_or_item(delta, "tool_calls") is not None: - delta_tool_calls = _get_attr_or_item(delta, "tool_calls") - if not delta_tool_calls: - continue - tool_delta = delta_tool_calls[0] - - # pylint: disable=unsubscriptable-object + # One delta can carry several calls, and a fragment continues the call + # with its `index`; providers that omit `index` continue the last call. + for tool_delta in _get_attr_or_item(delta, "tool_calls") or []: tool_id = _get_attr_or_item(tool_delta, "id") - if not tool_calls or (tool_id and tool_calls[-1]["id"] != tool_id): - function_arg = _get_attr_or_item(tool_delta, "function") - tool_calls = (tool_calls or []) + [ - { - "id": tool_id, - "type": _get_attr_or_item(tool_delta, "type"), - "function": { - "name": _get_attr_or_item(function_arg, "name"), - "arguments": _get_attr_or_item(function_arg, "arguments") or "", - }, - } - ] + tool_index = _get_attr_or_item(tool_delta, "index") + function_arg = _get_attr_or_item(tool_delta, "function") + if tool_index is not None: + tool_call = tool_calls_by_index.get(tool_index) + else: + tool_call = tool_calls[-1] if tool_calls else None + # A new id always opens a new call, even under an index already in use. + if tool_call is not None and tool_id and tool_call["id"] and tool_call["id"] != tool_id: + tool_call = None + + if tool_call is None: + tool_call = { + "id": tool_id, + "type": _get_attr_or_item(tool_delta, "type"), + "function": { + "name": _get_attr_or_item(function_arg, "name"), + "arguments": _get_attr_or_item(function_arg, "arguments") or "", + }, + } + tool_calls = (tool_calls or []) + [tool_call] + if tool_index is not None: + tool_calls_by_index[tool_index] = tool_call else: - # pylint: disable=unsubscriptable-object - # append to existing tool call - function_arg = _get_attr_or_item(tool_delta, "function") args = _get_attr_or_item(function_arg, "arguments") or "" if isinstance(args, str): - # pylint: disable=unsubscriptable-object - tool_calls[-1]["function"]["arguments"] += args + tool_call["function"]["arguments"] += args processed_audio = ( _process_chat_completion_audio(audio, audio_format=audio_format) if audio is not None else None From f8972eef52700e93f789539fe644d5128b087d43 Mon Sep 17 00:00:00 2001 From: Alex Castillo Date: Fri, 16 Oct 2026 17:27:58 -0400 Subject: [PATCH 2/2] fix(openai): open a new streamed tool call when an id follows an id-less call A call that arrived without an id stored id=None, so a later call with an id under the same (reused) index was merged into it. A non-empty id that differs from the stored one now always opens a new call, as before this branch. Fragments without an id still continue the current call (GLM-4.6 behavior). Co-Authored-By: Claude Opus 5.5 --- .../integrations/openai/test_openai.py | 21 +++++++++++++++++++ .../braintrust/integrations/openai/tracing.py | 5 +++-- 2 files changed, 24 insertions(+), 2 deletions(-) diff --git a/py/src/braintrust/integrations/openai/test_openai.py b/py/src/braintrust/integrations/openai/test_openai.py index 1adc711e..43b4646e 100644 --- a/py/src/braintrust/integrations/openai/test_openai.py +++ b/py/src/braintrust/integrations/openai/test_openai.py @@ -2466,6 +2466,27 @@ def test_chat_completion_streaming_new_tool_call_id_under_reused_index(self, mem assert message["tool_calls"][1]["function"] == {"name": "get_time", "arguments": '{"tz": "CET"}'} assert not memory_logger.pop() + def test_chat_completion_streaming_call_with_id_after_call_without_id_under_reused_index(self, memory_logger): + """A call that arrives without an id is not merged into the next call that has one.""" + assert not memory_logger.pop() + + all_results = [ + _tool_call_chunk( + [_tool_call_delta(0, name="get_weather", arguments='{"city": "Paris"}')], + role="assistant", + ), + _tool_call_chunk([_tool_call_delta(0, call_id="call_2", name="get_time", arguments='{"tz": "CET"}')]), + _tool_call_chunk([], finish_reason="tool_calls"), + ] + + wrapper = ChatCompletionWrapper(None, None) + message = wrapper._postprocess_streaming_results(all_results)["output"][0]["message"] + + assert [tool_call["id"] for tool_call in message["tool_calls"]] == [None, "call_2"] + assert message["tool_calls"][0]["function"] == {"name": "get_weather", "arguments": '{"city": "Paris"}'} + assert message["tool_calls"][1]["function"] == {"name": "get_time", "arguments": '{"tz": "CET"}'} + assert not memory_logger.pop() + def test_chat_completion_streaming_audio_is_materialized_as_attachment(self, memory_logger): assert not memory_logger.pop() diff --git a/py/src/braintrust/integrations/openai/tracing.py b/py/src/braintrust/integrations/openai/tracing.py index 12febe50..73e46252 100644 --- a/py/src/braintrust/integrations/openai/tracing.py +++ b/py/src/braintrust/integrations/openai/tracing.py @@ -840,8 +840,9 @@ def _postprocess_streaming_results( tool_call = tool_calls_by_index.get(tool_index) else: tool_call = tool_calls[-1] if tool_calls else None - # A new id always opens a new call, even under an index already in use. - if tool_call is not None and tool_id and tool_call["id"] and tool_call["id"] != tool_id: + # A new id always opens a new call, even under an index already in use + # and even when the call there arrived without an id. + if tool_call is not None and tool_id and tool_call["id"] != tool_id: tool_call = None if tool_call is None: