diff --git a/py/src/braintrust/integrations/openai/test_openai.py b/py/src/braintrust/integrations/openai/test_openai.py index 4f5d592d..43b4646e 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,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() diff --git a/py/src/braintrust/integrations/openai/tracing.py b/py/src/braintrust/integrations/openai/tracing.py index 3d4538ee..73e46252 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,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