fix(openai): keep every streamed tool call when a chunk carries several

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,103 @@ 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_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()
 

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,37 @@ 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
+                # 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:
+                    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