diff --git a/src/agents/models/chatcmpl_stream_handler.py b/src/agents/models/chatcmpl_stream_handler.py index 289e116b06..937ab5e795 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 = "" @@ -157,6 +157,27 @@ class _BufferedToolCall: extra_content: dict[str, Any] | None = None +def _buffered_tool_calls_in_replay_order( + buffered_calls: dict[int | None, _BufferedToolCall], +) -> list[_BufferedToolCall]: + """Sort indexed calls while preserving where the index-less call first appeared.""" + indexed_calls = sorted( + (call for call in buffered_calls.values() if isinstance(call.index, int)), + key=lambda call: cast(int, call.index), + ) + if None not in buffered_calls: + return indexed_calls + + unindexed_position = 0 + for index in buffered_calls: + if index is None: + break + unindexed_position += 1 + + indexed_calls.insert(unindexed_position, buffered_calls[None]) + return indexed_calls + + def _merge_buffered_metadata( current: dict[str, Any] | None, incoming: dict[str, Any], @@ -307,12 +328,94 @@ 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 None in matching_indexes: + if len(matching_indexes) > 1: + raise ModelBehaviorError( + "Chat Completions tool call delta supplied an index and the same ID " + "matched multiple buffered calls." + ) + if tool_call_index in buffered_calls: + raise ModelBehaviorError( + "Chat Completions tool call delta supplied an index already used by " + "another buffered call." + ) + + buffered_call = buffered_calls.pop(None) + buffered_call.index = tool_call_index + buffered_calls[tool_call_index] = buffered_call + elif ( + not tool_call_delta.id + and None in buffered_calls + and tool_call_index not in buffered_calls + ): + buffered_name = buffered_calls[None].name + if not function_name or not buffered_name or function_name == buffered_name: + if len(buffered_calls) > 1: + raise ModelBehaviorError( + "Chat Completions tool call delta supplied a new index without an ID " + "while multiple function calls were being buffered." + ) + buffered_call = buffered_calls.pop(None) + buffered_call.index = tool_call_index + buffered_calls[tool_call_index] = buffered_call + else: + if len(matching_indexes) == 1: + tool_call_index = matching_indexes[0] + elif len(matching_indexes) > 1: + raise ModelBehaviorError( + "Chat Completions tool call delta omitted an index and the same ID " + "matched multiple buffered calls." + ) + elif tool_call_delta.id and None in buffered_calls and buffered_calls[None].call_id: + raise ModelBehaviorError( + "Chat Completions tool call delta omitted an index with a new ID while " + "another index-less call was being buffered." + ) + else: + if function_name and None in buffered_calls: + buffered_name = buffered_calls[None].name + if buffered_name and buffered_name != function_name: + raise ModelBehaviorError( + "Chat Completions tool call delta omitted an index and used a " + "different function name from the buffered index-less call." + ) + + if not tool_call_delta.id and function_name: + if None in buffered_calls: + tool_call_index = None + elif len(buffered_calls) == 1: + sole_index = next(iter(buffered_calls)) + sole_name = buffered_calls[sole_index].name + if not sole_name or sole_name == function_name: + tool_call_index = sole_index + else: + tool_call_index = None + else: + tool_call_index = None + elif not tool_call_delta.id and None in buffered_calls: + tool_call_index = None + elif not tool_call_delta.id and len(buffered_calls) == 1: + tool_call_index = next(iter(buffered_calls)) + elif not tool_call_delta.id and len(buffered_calls) > 1: + raise ModelBehaviorError( + "Chat Completions tool call delta omitted an index while multiple " + "function tool 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: @@ -341,6 +444,7 @@ 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 +457,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 +478,17 @@ 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], ) -> ChatCompletionChunk: + ordered_calls = _buffered_tool_calls_in_replay_order(buffered_calls) + occupied_indexes = passthrough_tool_call_indexes | { + call.index for call in ordered_calls if isinstance(call.index, int) + } + fallback_index = max(occupied_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 +503,10 @@ 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] = {} + buffered_calls: dict[int | None, _BufferedToolCall] = {} passthrough_tool_call_indexes: set[int] = set() + passthrough_tool_call_indexes_by_id: dict[str, int | None] = {} + saw_unindexed_passthrough_tool_call = False saw_passthrough_tool_call = False last_chunk: ChatCompletionChunk | None = None @@ -418,13 +530,137 @@ 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: + is_untyped_continuation = ( + getattr(tool_call_delta, "type", None) is None + and tool_call_delta.function is None + ) + is_unindexed_untyped_continuation = ( + not isinstance(tool_call_delta.index, int) and is_untyped_continuation + ) + buffered_id_matches = [ + buffered_call + for buffered_call in buffered_calls.values() + if tool_call_delta.id and buffered_call.call_id == tool_call_delta.id + ] + is_passthrough_continuation_by_id = ( + is_untyped_continuation + and bool(tool_call_delta.id) + and tool_call_delta.id in passthrough_tool_call_indexes_by_id + ) + is_unindexed_passthrough_continuation = ( + is_unindexed_untyped_continuation + and ( + is_passthrough_continuation_by_id + or (saw_unindexed_passthrough_tool_call and not tool_call_delta.id) + ) + ) + if is_passthrough_continuation_by_id and buffered_id_matches: + raise ModelBehaviorError( + "Chat Completions tool call delta ID matched both a buffered " + "function call and a passthrough call." + ) + + unindexed_buffered_call = buffered_calls.get(None) + if ( + isinstance(tool_call_delta.index, int) + and tool_call_delta.index in passthrough_tool_call_indexes + and unindexed_buffered_call is not None + and ( + ( + tool_call_delta.id + and unindexed_buffered_call.call_id == tool_call_delta.id + ) + or ( + not tool_call_delta.id + and len(buffered_calls) == 1 + and tool_call_delta.function is not None + ) + ) + ): + raise ModelBehaviorError( + "Chat Completions tool call delta supplied an index already used " + "by a passthrough call." + ) + + if ( + tool_call_delta.index in passthrough_tool_call_indexes + or is_passthrough_continuation_by_id + or is_unindexed_passthrough_continuation + ): + if passthrough_id := tool_call_delta.id: + owner_index = passthrough_tool_call_indexes_by_id.get( + passthrough_id + ) + if isinstance(tool_call_delta.index, int): + promotes_unindexed_passthrough_owner = ( + is_passthrough_continuation_by_id + and not isinstance(owner_index, int) + ) + if isinstance(owner_index, int) and ( + owner_index != tool_call_delta.index + ): + raise ModelBehaviorError( + "Chat Completions passthrough tool call delta supplied " + "a different index from its buffered ID owner." + ) + if tool_call_delta.index in buffered_calls: + raise ModelBehaviorError( + "Chat Completions passthrough tool call delta supplied " + "an index already used by a buffered function call." + ) + passthrough_tool_call_indexes.add(tool_call_delta.index) + passthrough_tool_call_indexes_by_id[passthrough_id] = ( + tool_call_delta.index + ) + if promotes_unindexed_passthrough_owner: + tool_call_delta = tool_call_delta.model_copy( + update={"type": "custom"} + ) + elif isinstance(owner_index, int): + tool_call_delta = tool_call_delta.model_copy( + update={"index": owner_index} + ) + passthrough_tool_call_indexes_by_id.setdefault( + passthrough_id, + tool_call_delta.index + if isinstance(tool_call_delta.index, int) + else None, + ) saw_passthrough_tool_call = True remaining_tool_calls.append(tool_call_delta) + elif ( + is_unindexed_untyped_continuation + and saw_passthrough_tool_call + and ( + not tool_call_delta.id + or ( + bool( + set(tool_call_delta.model_extra or {}) + - {"provider_specific_fields", "extra_content"} + ) + and not buffered_id_matches + ) + ) + ): + raise ModelBehaviorError( + "Chat Completions tool call delta omitted an index, type, and " + "function payload after a passthrough call, so it could not be " + "attributed safely." + ) elif cls._should_buffer_tool_call_delta(tool_call_delta): cls._accumulate_tool_call_delta(buffered_calls, tool_call_delta) else: - passthrough_tool_call_indexes.add(tool_call_delta.index) + if isinstance(tool_call_delta.index, int): + passthrough_tool_call_indexes.add(tool_call_delta.index) + else: + saw_unindexed_passthrough_tool_call = True + if tool_call_delta.id: + passthrough_tool_call_indexes_by_id.setdefault( + tool_call_delta.id, + tool_call_delta.index + if isinstance(tool_call_delta.index, int) + else None, + ) saw_passthrough_tool_call = True remaining_tool_calls.append(tool_call_delta) @@ -458,7 +694,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..9f85db983f 100644 --- a/tests/models/test_openai_chatcompletions_stream.py +++ b/tests/models/test_openai_chatcompletions_stream.py @@ -1,10 +1,12 @@ import asyncio import logging from collections.abc import AsyncIterator +from dataclasses import replace from typing import Any, cast 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 +109,44 @@ 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) -> ChatCompletionChunk: + # Build the chunk the way ``AsyncStream`` does: ``construct_type`` does not validate the + # provider payload, so a tool call delta that omits ``index`` keeps ``index=None`` instead + # of failing validation. That is the shape OpenAI-compatible providers can produce. + return cast( + ChatCompletionChunk, + construct_type( + type_=ChatCompletionChunk, + value={ + "id": "chunk-id", + "object": "chat.completion.chunk", + "created": 1, + "model": "fake", + "choices": [{"index": 0, "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 +898,1135 @@ 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( + { + "role": "assistant", + "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"), + ) + + unbuffered_events = await _collect_handler_events(*chunks) + buffered_events = await _collect_buffered_handler_events(*chunks) + + expected_calls = [("call_1", "first_func", "{}"), ("call_2", "second_func", "{}")] + 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_preserves_missing_index_call_arrival_order() -> None: + chunks = ( + _lenient_chunk( + { + "role": "assistant", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "first_func", "arguments": "{}"}, + } + ], + } + ), + _lenient_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "call_2", + "type": "function", + "function": {"name": "second_func", "arguments": "{}"}, + } + ] + } + ), + _lenient_chunk({}, finish_reason="tool_calls"), + ) + + unbuffered_events = await _collect_handler_events(*chunks) + buffered_events = await _collect_buffered_handler_events(*chunks) + + expected_calls = [("call_1", "first_func", "{}"), ("call_2", "second_func", "{}")] + 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_sorts_indexed_calls_by_index() -> None: + chunks = ( + _lenient_chunk( + { + "tool_calls": [ + { + "index": 1, + "id": "call_2", + "type": "function", + "function": {"name": "second_func", "arguments": "{}"}, + } + ] + } + ), + _lenient_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "first_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_buffer_tool_call_stream_avoids_passthrough_index_for_missing_index() -> None: + custom_tool_call_delta = ChoiceDeltaToolCall.model_construct( + index=0, + id="custom-id", + type="custom", + ) + custom_chunk = ChatCompletionChunk( + id="chunk-id", + created=1, + model="fake", + object="chat.completion.chunk", + choices=[Choice(index=0, delta=ChoiceDelta(tool_calls=[custom_tool_call_delta]))], + ) + chunks = ( + custom_chunk, + _lenient_chunk( + { + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "my_func", "arguments": "{}"}, + } + ] + } + ), + _lenient_chunk({}, finish_reason="tool_calls"), + ) + + buffered_events = await _collect_buffered_handler_events(*chunks) + + assert _completed_function_calls(buffered_events) == [("call_1", "my_func", "{}")] + + +@pytest.mark.parametrize("function_tool_call_index", [0, None]) +@pytest.mark.asyncio +async def test_buffer_tool_call_stream_ignores_missing_passthrough_index( + function_tool_call_index: int | None, +) -> None: + custom_tool_call_delta = ChoiceDeltaToolCall.model_construct( + index=None, + id="custom-id", + type="custom", + ) + custom_chunk = ChatCompletionChunk( + id="chunk-id", + created=1, + model="fake", + object="chat.completion.chunk", + choices=[Choice(index=0, delta=ChoiceDelta(tool_calls=[custom_tool_call_delta]))], + ) + function_tool_call: dict[str, Any] = { + "id": "call_1", + "type": "function", + "function": {"name": "my_func", "arguments": "{}"}, + } + if function_tool_call_index is not None: + function_tool_call["index"] = function_tool_call_index + chunks = ( + custom_chunk, + _lenient_chunk({"tool_calls": [function_tool_call]}), + _lenient_chunk({}, finish_reason="tool_calls"), + ) + + buffered_events = await _collect_buffered_handler_events(*chunks) + + assert _completed_function_calls(buffered_events) == [("call_1", "my_func", "{}")] + + +@pytest.mark.parametrize("continuation_id", [None, "custom-id"]) +@pytest.mark.asyncio +async def test_buffer_tool_call_stream_forwards_missing_index_passthrough_continuation( + continuation_id: str | None, +) -> None: + continuation: dict[str, Any] = {"custom": {"input": "nt(1)"}} + if continuation_id is not None: + continuation["id"] = continuation_id + chunks = ( + _lenient_chunk( + { + "tool_calls": [ + { + "id": "custom-id", + "type": "custom", + "custom": {"name": "code_exec", "input": "pri"}, + } + ] + } + ), + _lenient_chunk({"tool_calls": [continuation]}), + _lenient_chunk( + { + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "my_func", "arguments": "{}"}, + } + ] + } + ), + _lenient_chunk({}, finish_reason="tool_calls"), + ) + + buffered_chunks = await _collect_buffered_tool_call_chunks(*chunks) + buffered_events = await _collect_buffered_handler_events(*chunks) + + assert buffered_chunks[0].choices[0].delta.tool_calls == chunks[0].choices[0].delta.tool_calls + assert buffered_chunks[1].choices[0].delta.tool_calls == chunks[1].choices[0].delta.tool_calls + assert _completed_function_calls(buffered_events) == [("call_1", "my_func", "{}")] + + +@pytest.mark.asyncio +async def test_buffer_tool_call_stream_promotes_indexed_passthrough_continuation_by_id() -> None: + chunks = ( + _lenient_chunk( + { + "tool_calls": [ + { + "id": "custom-id", + "type": "custom", + "custom": {"name": "code_exec", "input": "pri"}, + } + ] + } + ), + _lenient_chunk( + { + "tool_calls": [ + { + "index": 2, + "id": "custom-id", + "custom": {"input": "nt("}, + } + ] + } + ), + _lenient_chunk({"tool_calls": [{"id": "custom-id", "custom": {"input": "1)"}}]}), + _lenient_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "my_func", "arguments": "{}"}, + } + ] + } + ), + _lenient_chunk({}, finish_reason="tool_calls"), + ) + + buffered_chunks = await _collect_buffered_tool_call_chunks(*chunks) + buffered_events = await _collect_buffered_handler_events(*chunks) + + indexed_continuation = buffered_chunks[1].choices[0].delta.tool_calls + unindexed_continuation = buffered_chunks[2].choices[0].delta.tool_calls + assert indexed_continuation and indexed_continuation[0].index == 2 + assert unindexed_continuation and unindexed_continuation[0].index == 2 + assert _completed_function_calls(buffered_events) == [("call_1", "my_func", "{}")] + + +@pytest.mark.asyncio +async def test_buffer_tool_call_stream_rejects_indexed_passthrough_continuation_collision() -> None: + chunks = ( + _lenient_chunk( + { + "tool_calls": [ + { + "id": "custom-id", + "type": "custom", + "custom": {"name": "code_exec", "input": "pri"}, + } + ] + } + ), + _lenient_chunk( + { + "tool_calls": [ + { + "index": 2, + "id": "call_1", + "type": "function", + "function": {"name": "my_func", "arguments": "{}"}, + } + ] + } + ), + _lenient_chunk( + { + "tool_calls": [ + { + "index": 2, + "id": "custom-id", + "custom": {"input": "nt(1)"}, + } + ] + } + ), + ) + + with pytest.raises(ModelBehaviorError, match="index already used by a buffered function"): + await _collect_buffered_tool_call_chunks(*chunks) + + +@pytest.mark.asyncio +async def test_buffer_tool_call_stream_rejects_passthrough_continuation_index_change() -> None: + chunks = ( + _lenient_chunk( + { + "tool_calls": [ + { + "index": 1, + "id": "custom-id", + "type": "custom", + "custom": {"name": "code_exec", "input": "pri"}, + } + ] + } + ), + _lenient_chunk( + { + "tool_calls": [ + { + "index": 2, + "id": "custom-id", + "custom": {"input": "nt(1)"}, + } + ] + } + ), + ) + + with pytest.raises(ModelBehaviorError, match="different index from its buffered ID owner"): + await _collect_buffered_tool_call_chunks(*chunks) + + +@pytest.mark.parametrize("function_tool_call_index", [1, None], ids=["indexed", "unindexed"]) +@pytest.mark.asyncio +async def test_buffer_tool_call_stream_forwards_indexed_passthrough_continuation_by_id( + function_tool_call_index: int | None, +) -> None: + function_tool_call: dict[str, Any] = { + "type": "function", + "function": {"name": "my_func", "arguments": "{}"}, + } + if function_tool_call_index is not None: + function_tool_call["index"] = function_tool_call_index + function_tool_call["id"] = "call_1" + + chunks = [ + _lenient_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "custom-id", + "type": "custom", + "custom": {"name": "code_exec", "input": "pri"}, + } + ] + } + ), + _lenient_chunk({"tool_calls": [function_tool_call]}), + _lenient_chunk({"tool_calls": [{"id": "custom-id", "custom": {"input": "nt(1)"}}]}), + ] + if function_tool_call_index is None: + chunks.append(_lenient_chunk({"tool_calls": [{"id": "call_1"}]})) + chunks.append(_lenient_chunk({}, finish_reason="tool_calls")) + + buffered_chunks = await _collect_buffered_tool_call_chunks(*chunks) + buffered_events = await _collect_buffered_handler_events(*chunks) + + continuation_tool_calls = buffered_chunks[1].choices[0].delta.tool_calls + assert continuation_tool_calls + assert continuation_tool_calls[0].index == 0 + assert continuation_tool_calls[0].id == "custom-id" + assert continuation_tool_calls[0].model_extra == {"custom": {"input": "nt(1)"}} + assert _completed_function_calls(buffered_events) == [("call_1", "my_func", "{}")] + + +@pytest.mark.asyncio +async def test_buffer_tool_call_stream_records_passthrough_id_from_indexed_continuation() -> None: + chunks = ( + _lenient_chunk( + { + "tool_calls": [ + { + "index": 0, + "type": "custom", + "custom": {"name": "code_exec", "input": "pri"}, + } + ] + } + ), + _lenient_chunk( + {"tool_calls": [{"index": 0, "id": "custom-id", "custom": {"input": "nt("}}]} + ), + _lenient_chunk( + { + "tool_calls": [ + { + "type": "function", + "function": {"name": "my_func", "arguments": "{}"}, + } + ] + } + ), + _lenient_chunk({"tool_calls": [{"id": "custom-id", "custom": {"input": "1)"}}]}), + _lenient_chunk({"tool_calls": [{"id": "call_1"}]}), + _lenient_chunk({}, finish_reason="tool_calls"), + ) + + buffered_chunks = await _collect_buffered_tool_call_chunks(*chunks) + buffered_events = await _collect_buffered_handler_events(*chunks) + + continuation_tool_calls = buffered_chunks[2].choices[0].delta.tool_calls + assert continuation_tool_calls + assert continuation_tool_calls[0].index == 0 + assert continuation_tool_calls[0].id == "custom-id" + assert continuation_tool_calls[0].model_extra == {"custom": {"input": "1)"}} + assert _completed_function_calls(buffered_events) == [("call_1", "my_func", "{}")] + + +@pytest.mark.asyncio +async def test_buffer_tool_call_stream_rejects_cross_domain_passthrough_id() -> None: + chunks = ( + _lenient_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "shared-id", + "type": "function", + "function": {"name": "my_func", "arguments": "{}"}, + } + ] + } + ), + _lenient_chunk( + { + "tool_calls": [ + { + "index": 1, + "id": "shared-id", + "type": "custom", + "custom": {"name": "code_exec", "input": "print(1)"}, + } + ] + } + ), + _lenient_chunk( + { + "tool_calls": [ + { + "id": "shared-id", + "extra_content": {"google": {"thought_signature": "sig"}}, + } + ] + } + ), + ) + + with pytest.raises(ModelBehaviorError, match="matched both"): + await _collect_buffered_tool_call_chunks(*chunks) + + +@pytest.mark.asyncio +async def test_buffer_tool_call_stream_allows_function_metadata_on_late_id() -> None: + chunks = ( + _lenient_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "custom-id", + "type": "custom", + "custom": {"name": "code_exec", "input": "print(1)"}, + } + ] + } + ), + _lenient_chunk( + { + "tool_calls": [ + { + "type": "function", + "function": {"name": "my_func", "arguments": "{}"}, + } + ] + } + ), + _lenient_chunk( + { + "tool_calls": [ + { + "id": "call_1", + "extra_content": {"google": {"thought_signature": "sig"}}, + } + ] + } + ), + _lenient_chunk({}, finish_reason="tool_calls"), + ) + + buffered_chunks = await _collect_buffered_tool_call_chunks(*chunks) + unbuffered_events = await _collect_handler_events(*chunks) + buffered_events = await _collect_buffered_handler_events(*chunks) + + replayed_tool_calls = buffered_chunks[-1].choices[0].delta.tool_calls + assert replayed_tool_calls + assert cast(Any, replayed_tool_calls[0]).extra_content == { + "google": {"thought_signature": "sig"} + } + expected_calls = [("call_1", "my_func", "{}")] + assert _completed_function_calls(unbuffered_events) == expected_calls + assert _completed_function_calls(buffered_events) == expected_calls + + +@pytest.mark.parametrize("continuation_id", [None, "unknown-id"], ids=["without-id", "unknown-id"]) +@pytest.mark.asyncio +async def test_buffer_tool_call_stream_rejects_ambiguous_unindexed_passthrough_continuation( + continuation_id: str | None, +) -> None: + continuation: dict[str, Any] = {"custom": {"input": "nt(1)"}} + if continuation_id is not None: + continuation["id"] = continuation_id + + chunks = ( + _lenient_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "custom-id", + "type": "custom", + "custom": {"name": "code_exec", "input": "pri"}, + } + ] + } + ), + _lenient_chunk( + { + "tool_calls": [ + { + "index": 1, + "id": "call_1", + "type": "function", + "function": {"name": "my_func", "arguments": "{}"}, + } + ] + } + ), + _lenient_chunk({"tool_calls": [continuation]}), + ) + + with pytest.raises(ModelBehaviorError, match="could not be attributed"): + await _collect_buffered_tool_call_chunks(*chunks) + + +@pytest.mark.parametrize("continuation_name", [None, "my_func"]) +@pytest.mark.asyncio +async def test_buffer_tool_call_stream_merges_missing_index_continuation( + continuation_name: str | None, +) -> None: + continuation_function = {"arguments": "1}"} + if continuation_name is not None: + continuation_function["name"] = continuation_name + chunks = ( + _lenient_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "my_func", "arguments": '{"a":'}, + } + ] + } + ), + _lenient_chunk({"tool_calls": [{"function": continuation_function}]}), + _lenient_chunk({}, finish_reason="tool_calls"), + ) + + buffered_events = await _collect_buffered_handler_events(*chunks) + + assert _completed_function_calls(buffered_events) == [("call_1", "my_func", '{"a":1}')] + + +@pytest.mark.asyncio +async def test_buffer_tool_call_stream_rejects_late_index_used_by_passthrough() -> None: + chunks = ( + _lenient_chunk( + { + "tool_calls": [ + { + "index": 2, + "id": "custom-id", + "type": "custom", + "custom": {"name": "code_exec", "input": "print(1)"}, + } + ] + } + ), + _lenient_chunk( + { + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "my_func", "arguments": '{"a":'}, + } + ] + } + ), + _lenient_chunk( + { + "tool_calls": [ + { + "index": 2, + "id": "call_1", + "type": "function", + "function": {"arguments": "1}"}, + } + ] + } + ), + ) + + with pytest.raises(ModelBehaviorError, match="index already used by a passthrough call"): + await _collect_buffered_tool_call_chunks(*chunks) + + +@pytest.mark.asyncio +async def test_buffer_tool_call_stream_merges_late_id_into_missing_index_call() -> None: + chunks = ( + _lenient_chunk( + { + "tool_calls": [ + { + "type": "function", + "function": {"name": "my_func", "arguments": '{"a":'}, + } + ] + } + ), + _lenient_chunk( + { + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"arguments": "1}"}, + } + ] + } + ), + _lenient_chunk({}, finish_reason="tool_calls"), + ) + + 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.parametrize("continuation_name", [None, "my_func"]) +@pytest.mark.asyncio +async def test_buffer_tool_call_stream_reconciles_late_index_by_id( + continuation_name: str | None, +) -> None: + continuation_function = {"arguments": "1}"} + if continuation_name is not None: + continuation_function["name"] = continuation_name + + chunks = ( + _lenient_chunk( + { + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "my_func", "arguments": '{"a":'}, + } + ] + } + ), + _lenient_chunk( + { + "tool_calls": [ + { + "index": 2, + "id": "call_1", + "type": "function", + "function": continuation_function, + } + ] + } + ), + _lenient_chunk({}, finish_reason="tool_calls"), + ) + + buffered_chunks = await _collect_buffered_tool_call_chunks(*chunks) + buffered_events = await _collect_buffered_handler_events(*chunks) + + replayed_tool_calls = buffered_chunks[-1].choices[0].delta.tool_calls + assert replayed_tool_calls + assert [tool_call.index for tool_call in replayed_tool_calls] == [2] + assert _completed_function_calls(buffered_events) == [("call_1", "my_func", '{"a":1}')] + + +@pytest.mark.asyncio +async def test_buffer_tool_call_stream_reconciles_late_index_without_id_for_sole_call() -> None: + chunks = ( + _lenient_chunk( + { + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "my_func", "arguments": '{"a":'}, + } + ] + } + ), + _lenient_chunk( + { + "tool_calls": [ + { + "index": 2, + "type": "function", + "function": {"arguments": "1}"}, + } + ] + } + ), + _lenient_chunk({}, finish_reason="tool_calls"), + ) + + buffered_chunks = await _collect_buffered_tool_call_chunks(*chunks) + buffered_events = await _collect_buffered_handler_events(*chunks) + + replayed_tool_calls = buffered_chunks[-1].choices[0].delta.tool_calls + assert replayed_tool_calls + assert [tool_call.index for tool_call in replayed_tool_calls] == [2] + assert _completed_function_calls(buffered_events) == [("call_1", "my_func", '{"a":1}')] + + +@pytest.mark.asyncio +async def test_buffer_tool_call_stream_keeps_different_named_idless_late_index_distinct() -> None: + chunks = ( + _lenient_chunk( + { + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "first_func", "arguments": '{"a":1}'}, + } + ] + } + ), + _lenient_chunk( + { + "tool_calls": [ + { + "index": 2, + "type": "function", + "function": {"name": "second_func", "arguments": '{"b":'}, + } + ] + } + ), + _lenient_chunk( + { + "tool_calls": [ + { + "index": 2, + "id": "call_2", + "type": "function", + "function": {"arguments": "1}"}, + } + ] + } + ), + _lenient_chunk({}, finish_reason="tool_calls"), + ) + + buffered_events = await _collect_buffered_handler_events(*chunks) + + assert _completed_function_calls(buffered_events) == [ + ("call_1", "first_func", '{"a":1}'), + ("call_2", "second_func", '{"b":1}'), + ] + + +@pytest.mark.parametrize( + "opening_name", + ["my_func", None], + ids=["same-name", "fills-missing-name"], +) +@pytest.mark.asyncio +async def test_buffer_tool_call_stream_merges_named_continuation_into_missing_index_slot( + opening_name: str | None, +) -> None: + opening_function = {"arguments": '{"a":'} + if opening_name is not None: + opening_function["name"] = opening_name + + chunks = ( + _lenient_chunk( + { + "tool_calls": [ + { + "type": "function", + "function": opening_function, + } + ] + } + ), + _lenient_chunk( + { + "tool_calls": [ + { + "type": "function", + "function": {"name": "my_func", "arguments": "1}"}, + } + ] + } + ), + _lenient_chunk({"tool_calls": [{"id": "call_1"}]}), + _lenient_chunk({}, finish_reason="tool_calls"), + ) + + 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_opens_named_missing_index_call_beside_indexed_call() -> None: + chunks = ( + _lenient_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "first_func", "arguments": "{}"}, + } + ] + } + ), + _lenient_chunk( + { + "tool_calls": [ + { + "type": "function", + "function": {"name": "second_func", "arguments": '{"b":'}, + } + ] + } + ), + _lenient_chunk( + { + "tool_calls": [ + { + "id": "call_2", + "type": "function", + "function": {"arguments": "1}"}, + } + ] + } + ), + _lenient_chunk({}, finish_reason="tool_calls"), + ) + + unbuffered_events = await _collect_handler_events(*chunks) + buffered_events = await _collect_buffered_handler_events(*chunks) + + expected_calls = [ + ("call_1", "first_func", "{}"), + ("call_2", "second_func", '{"b":1}'), + ] + assert _completed_function_calls(unbuffered_events) == expected_calls + assert _completed_function_calls(buffered_events) == expected_calls + + +@pytest.mark.parametrize("continuation_id", [None, "call_2"], ids=["without-id", "with-new-id"]) +@pytest.mark.asyncio +async def test_buffer_tool_call_stream_rejects_different_name_for_missing_index_slot( + continuation_id: str | None, +) -> None: + continuation: dict[str, Any] = { + "type": "function", + "function": {"name": "second_func", "arguments": "{}"}, + } + if continuation_id is not None: + continuation["id"] = continuation_id + + chunks = ( + _lenient_chunk( + { + "tool_calls": [ + { + "type": "function", + "function": {"name": "first_func", "arguments": "{}"}, + } + ] + } + ), + _lenient_chunk({"tool_calls": [continuation]}), + ) + + with pytest.raises(ModelBehaviorError, match="different function name"): + await _collect_buffered_handler_events(*chunks) + + +@pytest.mark.parametrize( + "continuation_function", + [ + {"arguments": "1}"}, + {"name": "my_func", "arguments": "1}"}, + ], + ids=["arguments-only", "repeated-name"], +) +@pytest.mark.asyncio +async def test_buffer_tool_call_stream_merges_missing_index_continuation_by_id( + continuation_function: dict[str, str], +) -> None: + chunks = ( + _lenient_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "my_func", "arguments": '{"a":'}, + } + ] + } + ), + _lenient_chunk( + { + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": continuation_function, + } + ] + } + ), + _lenient_chunk({}, finish_reason="tool_calls"), + ) + + buffered_events = await _collect_buffered_handler_events(*chunks) + + assert _completed_function_calls(buffered_events) == [("call_1", "my_func", '{"a":1}')] + + +def test_accumulate_tool_call_delta_rejects_ambiguous_repeated_id() -> None: + buffered_calls = { + 0: _BufferedToolCall(index=0, call_id="call_1", name="first_func"), + 1: _BufferedToolCall(index=1, call_id="call_1", name="second_func"), + } + continuation = ChoiceDeltaToolCall.model_construct( + index=None, + id="call_1", + type="function", + function=ChoiceDeltaToolCallFunction(arguments="{}"), + ) + + with pytest.raises(ModelBehaviorError, match="same ID matched multiple buffered calls"): + ChatCmplStreamHandler._accumulate_tool_call_delta(buffered_calls, continuation) + + +def test_accumulate_tool_call_delta_rejects_new_id_for_occupied_missing_index() -> None: + buffered_calls = { + None: _BufferedToolCall(index=None, call_id="call_1", name="first_func"), + } + new_call = ChoiceDeltaToolCall.model_construct( + index=None, + id="call_2", + type="function", + function=ChoiceDeltaToolCallFunction(name="second_func", arguments="{}"), + ) + + with pytest.raises(ModelBehaviorError, match="new ID while another index-less call"): + ChatCmplStreamHandler._accumulate_tool_call_delta(buffered_calls, new_call) + + +def test_accumulate_tool_call_delta_rejects_late_index_collision() -> None: + unindexed_call = _BufferedToolCall(index=None, call_id="call_1", name="first_func") + indexed_call = _BufferedToolCall(index=2, call_id="call_2", name="second_func") + buffered_calls = {None: unindexed_call, 2: indexed_call} + expected_calls = {None: replace(unindexed_call), 2: replace(indexed_call)} + continuation = ChoiceDeltaToolCall.model_construct( + index=2, + id="call_1", + type="function", + function=ChoiceDeltaToolCallFunction(arguments="{}"), + ) + + with pytest.raises(ModelBehaviorError, match="index already used"): + ChatCmplStreamHandler._accumulate_tool_call_delta(buffered_calls, continuation) + + assert buffered_calls == expected_calls + + +def test_accumulate_tool_call_delta_rejects_ambiguous_late_index() -> None: + unindexed_call = _BufferedToolCall(index=None, call_id="call_1", name="first_func") + indexed_call = _BufferedToolCall(index=1, call_id="call_1", name="second_func") + buffered_calls = {None: unindexed_call, 1: indexed_call} + expected_calls = {None: replace(unindexed_call), 1: replace(indexed_call)} + continuation = ChoiceDeltaToolCall.model_construct( + index=2, + id="call_1", + type="function", + function=ChoiceDeltaToolCallFunction(arguments="{}"), + ) + + with pytest.raises(ModelBehaviorError, match="same ID matched multiple buffered calls"): + ChatCmplStreamHandler._accumulate_tool_call_delta(buffered_calls, continuation) + + assert buffered_calls == expected_calls + + +def test_accumulate_tool_call_delta_rejects_ambiguous_idless_late_index() -> None: + unindexed_call = _BufferedToolCall(index=None, call_id="call_1", name="first_func") + indexed_call = _BufferedToolCall(index=0, call_id="call_2", name="second_func") + buffered_calls = {None: unindexed_call, 0: indexed_call} + expected_calls = {None: replace(unindexed_call), 0: replace(indexed_call)} + continuation = ChoiceDeltaToolCall.model_construct( + index=2, + id=None, + type="function", + function=ChoiceDeltaToolCallFunction(arguments="{}"), + ) + + with pytest.raises(ModelBehaviorError, match="new index without an ID"): + ChatCmplStreamHandler._accumulate_tool_call_delta(buffered_calls, continuation) + + assert buffered_calls == expected_calls + + +@pytest.mark.asyncio +async def test_buffer_tool_call_stream_rejects_ambiguous_missing_index_continuation() -> None: + chunks = ( + _lenient_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "first_func", "arguments": '{"a":'}, + }, + { + "index": 1, + "id": "call_2", + "type": "function", + "function": {"name": "second_func", "arguments": '{"b":'}, + }, + ] + } + ), + _lenient_chunk({"tool_calls": [{"function": {"arguments": "1}"}}]}), + ) + + with pytest.raises( + ModelBehaviorError, match="omitted an index while multiple function tool calls" + ): + await _collect_buffered_handler_events(*chunks) + + @pytest.mark.parametrize( ("delta", "expected"), [