diff --git a/src/ai/models/ai_gateway/stream.py b/src/ai/models/ai_gateway/stream.py index ce143806..aecb3c53 100644 --- a/src/ai/models/ai_gateway/stream.py +++ b/src/ai/models/ai_gateway/stream.py @@ -158,9 +158,16 @@ async def _build_request_body( # --------------------------------------------------------------------------- -def _expand_tool_call(data: dict[str, Any]) -> list[events_.Event]: - """Expand a complete ``tool-call`` part into Start + Delta + End.""" +def _expand_tool_call( + data: dict[str, Any], streamed_tool_ids: set[str] +) -> list[events_.Event]: + """Expand a complete ``tool-call`` part into Start + Delta + End. + + Returns empty when the tool was already streamed via ``tool-input-*``. + """ tc_id = data.get("toolCallId", "") + if tc_id in streamed_tool_ids: + return [] tool_name = data.get("toolName", "") tool_input = data.get("input", "") args_str = tool_input if isinstance(tool_input, str) else json.dumps(tool_input) @@ -198,7 +205,9 @@ def _parse_usage(data: Any) -> usage_.Usage: ) -def _parse_stream_part(data: dict[str, Any]) -> list[events_.Event]: +def _parse_stream_part( + data: dict[str, Any], streamed_tool_ids: set[str] +) -> list[events_.Event]: """Convert a ``LanguageModelV3StreamPart`` to public events.""" match data.get("type", ""): case "text-start": @@ -230,9 +239,11 @@ def _parse_stream_part(data: dict[str, Any]) -> list[events_.Event]: return [events_.ReasoningEnd(block_id=data.get("id", "reasoning"))] case "tool-input-start": + tcid = data.get("id", "") + streamed_tool_ids.add(tcid) return [ events_.ToolStart( - tool_call_id=data.get("id", ""), + tool_call_id=tcid, tool_name=data.get("toolName", ""), ) ] @@ -249,7 +260,7 @@ def _parse_stream_part(data: dict[str, Any]) -> list[events_.Event]: return [events_.ToolEnd(tool_call_id=data.get("id", ""))] case "tool-call": - return _expand_tool_call(data) + return _expand_tool_call(data, streamed_tool_ids) case "file": return [ @@ -313,8 +324,9 @@ async def stream( ) yield events_.StreamStart() + streamed_tool_ids: set[str] = set() async for data in _common.parse_sse_lines(response): - for event in _parse_stream_part(data): + for event in _parse_stream_part(data, streamed_tool_ids): yield event except errors.GatewayError: raise diff --git a/tests/models/ai_gateway/test_protocol.py b/tests/models/ai_gateway/test_protocol.py index 38928f4e..e16a05d8 100644 --- a/tests/models/ai_gateway/test_protocol.py +++ b/tests/models/ai_gateway/test_protocol.py @@ -235,7 +235,7 @@ class TestParseStreamPartComplex: def test_text_delta_uses_textDelta_key(self) -> None: """The gateway sends ``textDelta`` (camelCase), not ``delta``.""" events = stream_mod._parse_stream_part( - {"type": "text-delta", "id": "t1", "textDelta": "Hello"} + {"type": "text-delta", "id": "t1", "textDelta": "Hello"}, set() ) assert isinstance(events[0], events_.TextDelta) assert events[0].chunk == "Hello" @@ -249,7 +249,8 @@ def test_tool_call_expands_to_three_events(self) -> None: "toolCallId": "tc-1", "toolName": "get_weather", "input": {"city": "SF"}, - } + }, + set(), ) assert len(events) == 3 assert isinstance(events[0], events_.ToolStart) @@ -258,6 +259,24 @@ def test_tool_call_expands_to_three_events(self) -> None: assert json.loads(events[1].chunk) == {"city": "SF"} assert isinstance(events[2], events_.ToolEnd) + def test_tool_call_skipped_when_already_streamed(self) -> None: + """A ``tool-call`` that duplicates a streamed tool is dropped.""" + seen: set[str] = set() + stream_mod._parse_stream_part( + {"type": "tool-input-start", "id": "tc-1", "toolName": "get_weather"}, + seen, + ) + events = stream_mod._parse_stream_part( + { + "type": "tool-call", + "toolCallId": "tc-1", + "toolName": "get_weather", + "input": {"city": "SF"}, + }, + seen, + ) + assert events == [] + def test_finish_flat_usage(self) -> None: events = stream_mod._parse_stream_part( { @@ -267,7 +286,8 @@ def test_finish_flat_usage(self) -> None: "prompt_tokens": 10, "completion_tokens": 20, }, - } + }, + set(), ) done = events[0] assert isinstance(done, events_.StreamEnd) @@ -293,7 +313,8 @@ def test_finish_v3_nested_usage(self) -> None: "reasoning": 30, }, }, - } + }, + set(), ) done = events[0] assert isinstance(done, events_.StreamEnd) @@ -311,7 +332,8 @@ def test_file_part(self) -> None: "id": "f1", "mediaType": "image/png", "data": "iVBORw0KGgo=", - } + }, + set(), ) assert len(events) == 1 assert isinstance(events[0], events_.FileEvent) @@ -321,14 +343,16 @@ def test_file_part(self) -> None: def test_file_part_defaults(self) -> None: """A minimal ``file`` part uses sensible defaults.""" - events = stream_mod._parse_stream_part({"type": "file", "data": "somedata"}) + events = stream_mod._parse_stream_part( + {"type": "file", "data": "somedata"}, set() + ) assert len(events) == 1 assert isinstance(events[0], events_.FileEvent) assert events[0].media_type == "application/octet-stream" def test_unknown_types_produce_no_events(self) -> None: for t in ("stream-start", "raw", "response-metadata", "banana"): - assert stream_mod._parse_stream_part({"type": t}) == [] + assert stream_mod._parse_stream_part({"type": t}, set()) == [] # ---------------------------------------------------------------------------