scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-05 02:41:56 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 48

fix(responses): tolerate dict terminal responses when estimating usage

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kerry 2026-09-16 02:05:44 +00:00
parent 7ed20406d7
commit 7680d3de86
2 changed files with 60 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

35
litellm/responses/streaming_iterator.py
View file

@ -428,24 +428,25 @@ class BaseResponsesAPIStreamingIterator:
openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED,
):
self.completed_response = openai_responses_api_chunk
_response_obj: Final[ResponsesAPIResponse | None] = getattr(
openai_responses_api_chunk, "response", None
)
if (
_chunk_type
in (
openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
)
and _response_obj is not None
and _response_obj.usage is None
_response_obj: Final[object] = getattr(openai_responses_api_chunk, "response", None)
if _chunk_type in (
openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
):
_response_obj.usage = _estimate_usage_safely(
self.model or "",
self.request_data.get("input"),
self.request_data,
self._generated_content + self._generated_tool_arguments,
)
if isinstance(_response_obj, ResponsesAPIResponse) and _response_obj.usage is None:
_response_obj.usage = _estimate_usage_safely(
self.model or "",
self.request_data.get("input"),
self.request_data,
self._generated_content + self._generated_tool_arguments,
)
elif isinstance(_response_obj, dict) and _response_obj.get("usage") is None: # pyright: ignore[reportUnknownMemberType] # the model_constructed terminal event leaves response as an untyped dict
_response_obj["usage"] = _estimate_usage_safely(
self.model or "",
self.request_data.get("input"),
self.request_data,
self._generated_content + self._generated_tool_arguments,
)
_stamp_responses_usage_cost(getattr(openai_responses_api_chunk, "response", None), self.logging_obj)
if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED:

42
tests/test_litellm/responses/test_streaming_iterator.py
View file

@ -849,3 +849,45 @@ async def test_completed_event_without_usage_counts_tool_input_deltas(tool_delta
assert usage is not None
assert usage.output_tokens > 0
assert usage.total_tokens == usage.input_tokens + usage.output_tokens
@pytest.mark.asyncio
async def test_completed_event_with_a_dict_response_still_gets_the_usage_estimate():
"""transform_streaming_response can model_construct a terminal event whose
response stays a plain dict; the estimate must fill it without raising."""
dict_response: Final = {
"id": "resp_dict",
"model": "gpt-4o-mini",
"object": "response",
"output": [],
"usage": None,
}
def _transform(model, parsed_chunk, logging_obj):
if parsed_chunk.get("type") == "response.completed":
return ResponseCompletedEvent.model_construct(type="response.completed", response=dict_response)
stub: Final = Mock()
stub.type = parsed_chunk.get("type")
if "delta" in parsed_chunk:
stub.delta = parsed_chunk.get("delta")
return stub
config: Final = Mock(spec=BaseResponsesAPIConfig)
config.transform_streaming_response.side_effect = _transform
iterator: Final = _make_iterator(
sse_events=[
_sse_event({"type": "response.output_text.delta", "delta": "hello world"}),
_sse_event({"type": "response.completed", "response": {}}),
],
logging_obj=_logging_obj_stub(),
config=config,
request_data={"input": "count these input tokens please"},
)
async for _ in iterator:
pass
usage: Final = iterator.completed_response.response["usage"]
assert usage is not None
assert usage.input_tokens > 0
assert usage.output_tokens > 0