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

fix: compare stream event types by equality so typed completed events keep their usage

Browse source
This commit is contained in:
mateo-berri 2026-09-01 17:37:07 -07:00
parent a38dfecd96
commit 8c7fe00d80
2 changed files with 33 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/llms/base_llm/guardrail_translation/utils.py
View file

@ -171,7 +171,7 @@ def blocked_responses_stream_usage(original_response: object) -> ResponseAPIUsag
(
response
for item in reversed(original_response)
if str(stream_item_field(item, "type") or "") == "response.completed"
if stream_item_field(item, "type") == "response.completed"
and (response := stream_item_field(item, "response")) is not None
),
None,

32
tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py
View file

@ -1389,6 +1389,38 @@ class TestBuildBlockSseChunks:
assert completed["output"][0]["content"][0]["text"] == "Blocked by policy."
assert completed["usage"] == {"input_tokens": 7, "output_tokens": 21, "total_tokens": 28}
def test_continuation_reads_usage_from_typed_completed_event(self):
from litellm.types.llms.openai import (
ResponseCompletedEvent,
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
)
handler = OpenAIResponsesHandler()
original = [
ResponseCompletedEvent(
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
response=ResponsesAPIResponse.model_validate(
{
"id": "resp_live",
"created_at": 1,
"model": "gpt-5.4-mini",
"output": [],
"usage": {"input_tokens": 7, "output_tokens": 21, "total_tokens": 28},
}
),
)
]
payloads = self._payloads(
handler.build_block_sse_chunks(
self._exc(original_response=original), stream_started=True, responses_so_far=[]
)
)
completed = payloads[-1]["response"]
assert completed["usage"]["input_tokens"] == 7
assert completed["usage"]["output_tokens"] == 21
assert completed["usage"]["total_tokens"] == 28
def test_continuation_closes_open_item_given_pydantic_events_with_enum_types(self):
from litellm.types.llms.openai import (
BaseLiteLLMOpenAIResponseObject,