diff --git a/src/agents/models/chatcmpl_stream_handler.py b/src/agents/models/chatcmpl_stream_handler.py index 289e116b06..07c39cdf92 100644 --- a/src/agents/models/chatcmpl_stream_handler.py +++ b/src/agents/models/chatcmpl_stream_handler.py @@ -149,7 +149,7 @@ def _open_thinking_block(self) -> dict[str, Any]: class _BufferedToolCall: """Accumulates a streamed Chat Completions function tool call.""" - index: int + index: int | None call_id: str | None = None name: str | None = None arguments: str = "" @@ -307,14 +307,67 @@ def _delta_has_passthrough_output(delta: ChoiceDelta | None) -> bool: @staticmethod def _accumulate_tool_call_delta( - buffered_calls: dict[int, _BufferedToolCall], + buffered_calls: dict[int | None, _BufferedToolCall], tool_call_delta: ChoiceDeltaToolCall, ) -> None: + tool_call_index = tool_call_delta.index + function_name = tool_call_delta.function.name if tool_call_delta.function else None + matching_indexes = [ + index + for index, buffered_call in buffered_calls.items() + if tool_call_delta.id and buffered_call.call_id == tool_call_delta.id + ] + + if isinstance(tool_call_index, int): + if matching_indexes and matching_indexes[0] != tool_call_index: + raise ModelBehaviorError( + "Chat Completions tool call index and ID identified different buffered calls." + ) + elif matching_indexes: + tool_call_index = matching_indexes[0] + elif tool_call_delta.id: + if None in buffered_calls and buffered_calls[None].call_id: + raise ModelBehaviorError( + "Chat Completions tool call delta omitted an index for a new call while " + "another index-less call was being buffered." + ) + tool_call_index = None + elif function_name: + if None in buffered_calls: + if buffered_calls[None].name: + raise ModelBehaviorError( + "Chat Completions tool call delta omitted both index and ID, so its " + "function-call owner could not be attributed safely." + ) + elif any(call.name == function_name for call in buffered_calls.values()): + raise ModelBehaviorError( + "Chat Completions tool call delta omitted both index and ID, so its " + "function-call owner could not be attributed safely." + ) + else: + tool_call_index = None + elif len(buffered_calls) == 1: + tool_call_index = next(iter(buffered_calls)) + elif len(buffered_calls) > 1: + raise ModelBehaviorError( + "Chat Completions tool call delta omitted both index and ID while multiple " + "function calls were being buffered." + ) + buffered_call = buffered_calls.setdefault( - tool_call_delta.index, - _BufferedToolCall(index=tool_call_delta.index), + tool_call_index, + _BufferedToolCall(index=tool_call_index), ) + if tool_call_delta.id and buffered_call.call_id not in (None, tool_call_delta.id): + raise ModelBehaviorError( + "Chat Completions tool call delta changed the ID of a buffered call." + ) + if function_name and buffered_call.name not in (None, function_name): + raise ModelBehaviorError( + "Chat Completions tool call delta changed the name of a buffered function call." + ) + if tool_call_delta.id: buffered_call.call_id = tool_call_delta.id @@ -341,6 +394,8 @@ def _accumulate_tool_call_delta( @staticmethod def _buffered_tool_call_delta( buffered_call: _BufferedToolCall, + *, + fallback_index: int = 0, ) -> ChoiceDeltaToolCall: if not buffered_call.call_id: raise ModelBehaviorError( @@ -353,7 +408,7 @@ def _buffered_tool_call_delta( ) tool_call_delta = ChoiceDeltaToolCall( - index=buffered_call.index, + index=buffered_call.index if isinstance(buffered_call.index, int) else fallback_index, id=buffered_call.call_id, function=ChoiceDeltaToolCallFunction( name=buffered_call.name, @@ -374,11 +429,23 @@ def _buffered_tool_call_delta( def _buffered_tool_calls_chunk( cls, template_chunk: ChatCompletionChunk, - buffered_calls: dict[int, _BufferedToolCall], + buffered_calls: dict[int | None, _BufferedToolCall], + passthrough_tool_call_indexes: set[int | None], ) -> ChatCompletionChunk: + ordered_calls = sorted( + buffered_calls.values(), + key=lambda call: ( + not isinstance(call.index, int), + call.index if isinstance(call.index, int) else 0, + ), + ) + used_indexes = { + index for index in passthrough_tool_call_indexes if isinstance(index, int) + } | {call.index for call in ordered_calls if isinstance(call.index, int)} + fallback_index = max(used_indexes, default=-1) + 1 tool_call_deltas = [ - cls._buffered_tool_call_delta(buffered_call) - for _, buffered_call in sorted(buffered_calls.items()) + cls._buffered_tool_call_delta(buffered_call, fallback_index=fallback_index) + for buffered_call in ordered_calls ] choice = Choice( index=0, @@ -393,8 +460,8 @@ async def buffer_tool_call_stream( stream: AsyncIterator[ChatCompletionChunk], ) -> AsyncIterator[ChatCompletionChunk]: """Buffer streamed function tool-call deltas until they are complete.""" - buffered_calls: dict[int, _BufferedToolCall] = {} - passthrough_tool_call_indexes: set[int] = set() + buffered_calls: dict[int | None, _BufferedToolCall] = {} + passthrough_tool_call_indexes: set[int | None] = set() saw_passthrough_tool_call = False last_chunk: ChatCompletionChunk | None = None @@ -418,12 +485,50 @@ async def buffer_tool_call_stream( if tool_call_deltas := (delta.tool_calls if delta and delta.tool_calls else None): remaining_tool_calls: list[ChoiceDeltaToolCall] = [] for tool_call_delta in tool_call_deltas: - if tool_call_delta.index in passthrough_tool_call_indexes: + matches_buffered_id = bool( + tool_call_delta.id + and any( + buffered_call.call_id == tool_call_delta.id + for buffered_call in buffered_calls.values() + ) + ) + is_unindexed_function_delta = tool_call_delta.index is None and ( + tool_call_delta.function is not None + or tool_call_delta.type == "function" + or matches_buffered_id + ) + if ( + tool_call_delta.index in passthrough_tool_call_indexes + and not is_unindexed_function_delta + ): + if ( + tool_call_delta.index is None + and tool_call_delta.id is None + and tool_call_delta.function is None + and None in buffered_calls + ): + raise ModelBehaviorError( + "Chat Completions tool call delta could not be attributed " + "safely between buffered and passthrough calls." + ) + if matches_buffered_id: + raise ModelBehaviorError( + "Chat Completions tool call index and ID identified different " + "tool call owners." + ) saw_passthrough_tool_call = True remaining_tool_calls.append(tool_call_delta) elif cls._should_buffer_tool_call_delta(tool_call_delta): cls._accumulate_tool_call_delta(buffered_calls, tool_call_delta) else: + if ( + isinstance(tool_call_delta.index, int) + and tool_call_delta.index in buffered_calls + ): + raise ModelBehaviorError( + "Chat Completions tool call index identified both a buffered " + "function call and a passthrough call." + ) passthrough_tool_call_indexes.add(tool_call_delta.index) saw_passthrough_tool_call = True remaining_tool_calls.append(tool_call_delta) @@ -458,7 +563,11 @@ async def buffer_tool_call_stream( if buffered_calls: if last_chunk is None: return - yield cls._buffered_tool_calls_chunk(last_chunk, buffered_calls) + yield cls._buffered_tool_calls_chunk( + last_chunk, + buffered_calls, + passthrough_tool_call_indexes, + ) @staticmethod def _merged_provider_data( diff --git a/tests/models/test_openai_chatcompletions_stream.py b/tests/models/test_openai_chatcompletions_stream.py index 5c79493769..75993695eb 100644 --- a/tests/models/test_openai_chatcompletions_stream.py +++ b/tests/models/test_openai_chatcompletions_stream.py @@ -5,6 +5,7 @@ import httpx2 import pytest +from openai._models import construct_type from openai.types.chat.chat_completion import ChatCompletion, Choice as ChatCompletionChoice from openai.types.chat.chat_completion_chunk import ( ChatCompletionChunk, @@ -107,6 +108,49 @@ async def _collect_buffered_tool_call_chunks( ] +async def _collect_buffered_handler_events(*chunks: ChatCompletionChunk) -> list[Any]: + return [ + event + async for event in ChatCmplStreamHandler.handle_stream( + _empty_response(), + cast(Any, ChatCmplStreamHandler.buffer_tool_call_stream(_completion_stream(*chunks))), + ) + ] + + +def _lenient_chunk( + delta: dict[str, Any], + finish_reason: str | None = None, + *, + choice_index: int = 0, +) -> ChatCompletionChunk: + """Construct a chunk as AsyncStream does, without validating a missing tool-call index.""" + return cast( + ChatCompletionChunk, + construct_type( + type_=ChatCompletionChunk, + value={ + "id": "chunk-id", + "object": "chat.completion.chunk", + "created": 1, + "model": "fake", + "choices": [ + {"index": choice_index, "delta": delta, "finish_reason": finish_reason} + ], + }, + ), + ) + + +def _completed_function_calls(events: list[Any]) -> list[tuple[str, str, str]]: + completed = cast(ResponseCompletedEvent, events[-1]) + return [ + (item.call_id, item.name, item.arguments) + for item in completed.response.output + if isinstance(item, ResponseFunctionToolCall) + ] + + def _url_citation( url: str = "https://example.com/weather", title: str = "Weather", @@ -858,6 +902,396 @@ async def test_buffer_tool_call_stream_keeps_passthrough_index_passthrough() -> assert buffered_chunks[1].choices[0].delta.tool_calls == [function_tool_call_delta] +@pytest.mark.asyncio +async def test_buffer_tool_call_stream_tolerates_missing_tool_call_index() -> None: + chunks = ( + _lenient_chunk( + { + "role": "assistant", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "my_func", "arguments": '{"a":'}, + } + ], + } + ), + _lenient_chunk({"tool_calls": [{"function": {"arguments": "1}"}}]}), + _lenient_chunk({}, finish_reason="tool_calls"), + ) + first_delta = chunks[0].choices[0].delta + assert first_delta.tool_calls and first_delta.tool_calls[0].index is None + + unbuffered_events = await _collect_handler_events(*chunks) + buffered_events = await _collect_buffered_handler_events(*chunks) + + expected_calls = [("call_1", "my_func", '{"a":1}')] + assert _completed_function_calls(unbuffered_events) == expected_calls + assert _completed_function_calls(buffered_events) == expected_calls + + +@pytest.mark.asyncio +async def test_buffer_tool_call_stream_orders_missing_index_after_indexed_calls() -> None: + chunks = ( + _lenient_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "first_func", "arguments": "{}"}, + } + ] + } + ), + _lenient_chunk( + { + "tool_calls": [ + { + "id": "call_2", + "type": "function", + "function": {"name": "second_func", "arguments": "{}"}, + } + ] + } + ), + _lenient_chunk({}, finish_reason="tool_calls"), + ) + + buffered_events = await _collect_buffered_handler_events(*chunks) + + assert _completed_function_calls(buffered_events) == [ + ("call_1", "first_func", "{}"), + ("call_2", "second_func", "{}"), + ] + + +@pytest.mark.asyncio +async def test_missing_index_replay_avoids_passthrough_index_collision() -> None: + chunks = ( + _lenient_chunk({"tool_calls": [{"index": 0, "id": "custom-id", "type": "custom"}]}), + _lenient_chunk( + { + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "my_func", "arguments": "{}"}, + } + ] + } + ), + ) + + buffered_chunks = await _collect_buffered_tool_call_chunks(*chunks) + + assert buffered_chunks[0].choices[0].delta.tool_calls == chunks[0].choices[0].delta.tool_calls + replayed_calls = buffered_chunks[-1].choices[0].delta.tool_calls + assert replayed_calls and replayed_calls[0].index == 1 + assert replayed_calls[0].id == "call_1" + + +@pytest.mark.asyncio +async def test_unindexed_passthrough_continuation_remains_passthrough() -> None: + opening = _lenient_chunk( + {"tool_calls": [{"id": "custom-id", "type": "custom", "custom": {"input": "start"}}]} + ) + continuation = _lenient_chunk({"tool_calls": [{"custom": {"input": "-end"}}]}) + + buffered_chunks = await _collect_buffered_tool_call_chunks(opening, continuation) + + assert [chunk.choices[0].delta.tool_calls for chunk in buffered_chunks] == [ + opening.choices[0].delta.tool_calls, + continuation.choices[0].delta.tool_calls, + ] + + +@pytest.mark.asyncio +async def test_indexed_function_does_not_block_unindexed_passthrough_continuation() -> None: + indexed_function = _lenient_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "my_func", "arguments": "{}"}, + } + ] + } + ) + custom_opening = _lenient_chunk( + {"tool_calls": [{"id": "custom-id", "type": "custom", "custom": {"input": "start"}}]} + ) + custom_continuation = _lenient_chunk({"tool_calls": [{"custom": {"input": "-end"}}]}) + + buffered_chunks = await _collect_buffered_tool_call_chunks( + indexed_function, + custom_opening, + custom_continuation, + ) + + assert [chunk.choices[0].delta.tool_calls for chunk in buffered_chunks[:-1]] == [ + custom_opening.choices[0].delta.tool_calls, + custom_continuation.choices[0].delta.tool_calls, + ] + replayed_calls = buffered_chunks[-1].choices[0].delta.tool_calls + assert replayed_calls and replayed_calls[0].id == "call_1" + + +@pytest.mark.asyncio +async def test_nonzero_choice_tool_call_suppresses_choice_zero_finish_error() -> None: + nonzero_choice = _lenient_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "my_func", "arguments": "{}"}, + } + ] + }, + choice_index=1, + ) + choice_zero_finish = _lenient_chunk({}, finish_reason="tool_calls") + + buffered_chunks = await _collect_buffered_tool_call_chunks( + nonzero_choice, + choice_zero_finish, + ) + + assert buffered_chunks == [nonzero_choice] + + +@pytest.mark.asyncio +async def test_missing_index_continuation_merges_into_sole_buffered_function() -> None: + chunks = ( + _lenient_chunk( + { + "tool_calls": [ + { + "index": 2, + "id": "call_1", + "type": "function", + "function": {"name": "my_func", "arguments": '{"a":'}, + } + ] + } + ), + _lenient_chunk({"tool_calls": [{"function": {"arguments": "1}"}}]}), + ) + + buffered_chunks = await _collect_buffered_tool_call_chunks(*chunks) + replayed_calls = buffered_chunks[-1].choices[0].delta.tool_calls + + assert replayed_calls and replayed_calls[0].index == 2 + assert replayed_calls[0].function + assert replayed_calls[0].function.arguments == '{"a":1}' + + +@pytest.mark.asyncio +async def test_missing_index_continuation_uses_a_unique_function_call_id() -> None: + opening = _lenient_chunk( + { + "tool_calls": [ + { + "index": index, + "id": f"call_{index}", + "type": "function", + "function": {"name": f"func_{index}", "arguments": "start"}, + } + for index in range(2) + ] + } + ) + continuation = _lenient_chunk( + {"tool_calls": [{"id": "call_1", "function": {"arguments": "-end"}}]} + ) + + buffered_chunks = await _collect_buffered_tool_call_chunks(opening, continuation) + replayed_calls = buffered_chunks[-1].choices[0].delta.tool_calls + + assert replayed_calls and [ + call.function.arguments for call in replayed_calls if call.function + ] == [ + "start", + "start-end", + ] + + +@pytest.mark.asyncio +async def test_late_index_for_unindexed_function_is_rejected() -> None: + opening = _lenient_chunk( + { + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "my_func", "arguments": '{"a":'}, + } + ] + } + ) + late_index = _lenient_chunk( + {"tool_calls": [{"index": 0, "id": "call_1", "function": {"arguments": "1}"}}]} + ) + + with pytest.raises(ModelBehaviorError, match="index and ID identified different"): + await _collect_buffered_tool_call_chunks(opening, late_index) + + +@pytest.mark.asyncio +async def test_missing_index_continuation_rejects_multiple_buffered_functions() -> None: + chunks = ( + _lenient_chunk( + { + "tool_calls": [ + { + "index": index, + "id": f"call_{index}", + "type": "function", + "function": {"name": f"func_{index}", "arguments": "{}"}, + } + ] + } + ) + for index in range(2) + ) + continuation = _lenient_chunk({"tool_calls": [{"function": {"arguments": "tail"}}]}) + + with pytest.raises(ModelBehaviorError, match="multiple function calls"): + await _collect_buffered_tool_call_chunks(*chunks, continuation) + + +@pytest.mark.asyncio +async def test_parallel_unindexed_function_openings_are_rejected() -> None: + first = _lenient_chunk( + { + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "first_func", "arguments": "{}"}, + } + ] + } + ) + second = _lenient_chunk( + { + "tool_calls": [ + { + "id": "call_2", + "type": "function", + "function": {"name": "second_func", "arguments": "{}"}, + } + ] + } + ) + + with pytest.raises(ModelBehaviorError, match="omitted an index for a new call"): + await _collect_buffered_tool_call_chunks(first, second) + + +@pytest.mark.asyncio +async def test_passthrough_call_rejects_a_buffered_function_index() -> None: + function = _lenient_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "my_func", "arguments": "{}"}, + } + ] + } + ) + custom = _lenient_chunk({"tool_calls": [{"index": 0, "id": "custom-id", "type": "custom"}]}) + + with pytest.raises(ModelBehaviorError, match="identified both a buffered function call"): + await _collect_buffered_tool_call_chunks(function, custom) + + +@pytest.mark.asyncio +async def test_anonymous_delta_rejects_mixed_unindexed_call_owners() -> None: + custom = _lenient_chunk( + {"tool_calls": [{"id": "custom-id", "type": "custom", "custom": {"input": "x"}}]} + ) + function = _lenient_chunk( + { + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "my_func", "arguments": "{}"}, + } + ] + } + ) + anonymous = _lenient_chunk( + {"tool_calls": [{"extra_content": {"google": {"thought_signature": "sig"}}}]} + ) + + with pytest.raises(ModelBehaviorError, match="could not be attributed safely"): + await _collect_buffered_tool_call_chunks(custom, function, anonymous) + + +@pytest.mark.asyncio +async def test_tool_call_index_and_id_owner_conflict_fails_closed() -> None: + custom = _lenient_chunk({"tool_calls": [{"index": 0, "id": "custom-id", "type": "custom"}]}) + function = _lenient_chunk( + { + "tool_calls": [ + { + "index": 1, + "id": "call_1", + "type": "function", + "function": {"name": "my_func", "arguments": "{}"}, + } + ] + } + ) + conflicting = _lenient_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "call_1", + "extra_content": {"google": {"thought_signature": "sig"}}, + } + ] + } + ) + + with pytest.raises(ModelBehaviorError, match="different tool call owners"): + await _collect_buffered_tool_call_chunks(custom, function, conflicting) + + +@pytest.mark.asyncio +async def test_same_name_is_not_used_to_infer_a_missing_index_owner() -> None: + indexed = _lenient_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "same_func", "arguments": "{}"}, + } + ] + } + ) + unindexed = _lenient_chunk( + {"tool_calls": [{"type": "function", "function": {"name": "same_func"}}]} + ) + + with pytest.raises(ModelBehaviorError, match="could not be attributed safely"): + await _collect_buffered_tool_call_chunks(indexed, unindexed) + + @pytest.mark.parametrize( ("delta", "expected"), [