diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 6cf36e3e660..df3c26c1a4f 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -356,7 +356,7 @@ class BaseResponsesAPIStreamingIterator: _delta: Final = getattr(openai_responses_api_chunk, "delta", None) if isinstance(_delta, str): self._generated_content += _delta - elif _event_type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA: + elif _event_type in _TOOL_ARGUMENTS_DELTA_EVENTS: _args_delta: Final = getattr(openai_responses_api_chunk, "delta", None) if isinstance(_args_delta, str): self._generated_tool_arguments += _args_delta @@ -1358,6 +1358,15 @@ def _usage_as_model(usage: object) -> ResponseAPIUsage | None: return None +_TOOL_ARGUMENTS_DELTA_EVENTS: Final = frozenset( + { + ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA, + ResponsesAPIStreamEvents.CUSTOM_TOOL_CALL_INPUT_DELTA, + ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DELTA, + } +) + + def _estimate_usage_from_text( model: str, request_input: object, diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index 234e63ca3be..7bd03863717 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -816,3 +816,31 @@ async def test_completed_event_survives_a_failing_usage_estimate(): assert yielded assert iterator.completed_response.response.usage is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tool_delta_event_type", + ["response.custom_tool_call_input.delta", "response.mcp_call_arguments.delta"], +) +async def test_completed_event_without_usage_counts_tool_input_deltas(tool_delta_event_type): + """Custom-tool and MCP argument deltas feed the streamed usage fallback the + same way function_call_arguments deltas do.""" + response = _responses_api_response_without_usage() + iterator = _make_iterator( + sse_events=[ + _sse_event({"type": tool_delta_event_type, "delta": '{"query": "weather in sf"}'}), + _sse_event({"type": "response.completed", "response": {}}), + ], + logging_obj=_logging_obj_stub(), + config=_mock_config_with_completed_response(response), + request_data={"input": "what is the weather in san francisco"}, + ) + + async for _ in iterator: + pass + + usage = iterator.completed_response.response.usage + assert usage is not None + assert usage.output_tokens > 0 + assert usage.total_tokens == usage.input_tokens + usage.output_tokens