From 5a50522bc59e4eca48ed8b3389526f55c024e2bd Mon Sep 17 00:00:00 2001 From: narmstrong-berri Date: Fri, 16 Oct 2026 05:35:42 +0000 Subject: [PATCH] refactor(responses): log bridged Responses streams on chat providers as Responses rows The Responses stream that wraps the planned chat stream now logs the request, so the spend row carries the Responses object the client received instead of the chat completion. The chat stream hands its assembled response to the bridge instead of dispatching success logging itself, and deferred stream logging is parked through Logging.hold_stream_completion. --- litellm/litellm_core_utils/litellm_logging.py | 4 + .../litellm_core_utils/streaming_handler.py | 36 +++++--- .../handler.py | 2 + .../streaming_iterator.py | 27 +++++- litellm/responses/streaming_iterator.py | 71 +++++++++------ .../test_deferred_guardrail_logging.py | 4 + .../test_google_ai_studio_responses_bridge.py | 4 +- .../test_handler.py | 91 ++++++++++++++++++- .../test_image_generation_output.py | 46 +++------- .../test_reasoning_input_item_preservation.py | 12 +-- .../test_reasoning_items.py | 4 +- .../test_session_handler_with_cold_storage.py | 29 ++---- .../test_streaming_iterator_transformation.py | 83 ++++------------- ..._tool_output_order_preserved_for_gemini.py | 6 +- 14 files changed, 237 insertions(+), 182 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 78ff073fdcb..f05cf0cb3e4 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -600,6 +600,10 @@ class Logging(LiteLLMLoggingBaseClass): baseline_cache_context: "BaselineCacheContext | None" = None baseline_observation: "CapturedBaselineObservation | None" = None + def hold_stream_completion(self, completion: tuple[object, ...]) -> None: + """Parks stream-end logging for the proxy to release once its end-of-stream guardrails finish.""" + setattr(self, "_deferred_stream_complete_args", completion) # noqa: B010 # declaring the slot here turns the proxy's reads of it into private-usage findings + @property def _defer_async_logging(self) -> bool: return self.defer_async_logging diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 19bc4b66fd7..c5ffe74414f 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -259,6 +259,8 @@ def _provider_hidden_params( class CustomStreamWrapper: + _on_complete: Callable[[object], None] | None = None + @property def hidden_params(self) -> dict[str, object]: # mutable-ok: API requires mutation return self._hidden_params @@ -1854,11 +1856,15 @@ class CustomStreamWrapper: if not cache_hit and self.logging_obj.llm_caching_handler is not None: await self.logging_obj.llm_caching_handler.add_streaming_response_to_cache(processed_chunk) + def hand_completion_to(self, on_complete: Callable[[object], None]) -> None: + """A bridge whose own client stream logs the request takes the assembled response instead of the logger.""" + self._on_complete = on_complete + def run_success_logging_and_cache_storage(self, processed_chunk, cache_hit: bool): """ Runs success logging in a thread and adds the response to the cache """ - if litellm.disable_streaming_logging is True: + if litellm.disable_streaming_logging is True or self._on_complete is not None: """ [NOT RECOMMENDED] Set this via `litellm.disable_streaming_logging = True`. @@ -2090,13 +2096,18 @@ class CustomStreamWrapper: processed_chunk=_cache_copy, cache_hit=cache_hit, ) - executor.submit( - self.logging_obj.success_handler, - _log_copy, - None, - None, - cache_hit, - ) + if self._on_complete is not None: + self._on_complete(_log_copy) + else: + executor.submit( + self.logging_obj.success_handler, + _log_copy, + None, + None, + cache_hit, + ) + elif self._on_complete is not None: + self._on_complete(response) else: executor.submit( self.logging_obj.success_handler, @@ -2346,7 +2357,9 @@ class CustomStreamWrapper: "_on_deferred_stream_complete", None, ) - if _deferred_cb is not None: + if self._on_complete is not None: + self._on_complete(complete_streaming_response) + elif _deferred_cb is not None: # Proxy has post-call guardrails. Store the assembled # response so the outer streaming consumer # (ProxyLogging.async_post_call_streaming_iterator_hook) @@ -2354,10 +2367,7 @@ class CustomStreamWrapper: # end-of-stream blocks complete. Scheduling here via # create_task would race with unified_guardrail's # end-of-stream block for short-stream providers. - self.logging_obj._deferred_stream_complete_args = ( - complete_streaming_response, - cache_hit, - ) + self.logging_obj.hold_stream_completion((complete_streaming_response, cache_hit)) else: # prefer_async_handlers routes CustomLogger to async_success_handler # when consumers use ``async for`` on sync-SDK streams. Legacy string diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index f78b33537e2..9125879a0c3 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -110,6 +110,7 @@ class LiteLLMCompletionTransformationHandler: responses_api_request=bridged_request, custom_llm_provider=custom_llm_provider, litellm_metadata=kwargs.get("litellm_metadata", {}), + logs_completion=planned_chat is not None, ) raise ValueError(f"Unexpected response type: {type(litellm_completion_response)}") @@ -163,5 +164,6 @@ class LiteLLMCompletionTransformationHandler: responses_api_request=responses_api_request, custom_llm_provider=litellm_completion_request.get("custom_llm_provider"), litellm_metadata=kwargs.get("litellm_metadata", {}), + logs_completion=planned_chat is not None, ) raise ValueError(f"Unexpected response type: {type(litellm_completion_response)}") diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 84d4648f75f..758e531d27d 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -6,6 +6,7 @@ from typing import Any, Final, cast import litellm from litellm.litellm_core_utils.hidden_params import get_or_create_hidden_params +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.main import stream_chunk_builder from litellm.responses.litellm_completion_transformation.custom_tools import ( build_tool_call_item_kwargs, @@ -22,7 +23,11 @@ from litellm.responses.litellm_completion_transformation.tool_search import ( from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) -from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator, stamp_responses_usage_cost +from litellm.responses.streaming_iterator import ( + ResponsesAPIStreamingIterator, + dispatch_stream_success, + stamp_responses_usage_cost, +) from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( PART_UNION_TYPES, @@ -105,6 +110,10 @@ def _delta_carries_reasoning_and_answer(delta: ChatCompletionDelta) -> bool: return has_reasoning and bool(delta.content or delta.tool_calls) +def _logged_by_this_stream(_assembled: object) -> None: + return None + + class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): """ Async iterator for processing streaming responses from the Responses API. @@ -118,9 +127,14 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): responses_api_request: ResponsesAPIOptionalRequestParams, custom_llm_provider: str | None = None, litellm_metadata: dict | None = None, + logs_completion: bool = False, ): + """``logs_completion`` makes this stream, not the chat stream it wraps, log the request as a Responses row.""" self.model: str = model self.litellm_custom_stream_wrapper: litellm.CustomStreamWrapper = litellm_custom_stream_wrapper + self._logs_completion: Final = logs_completion + if logs_completion: + litellm_custom_stream_wrapper.hand_completion_to(_logged_by_this_stream) self.request_input: str | ResponseInputParam = request_input self.responses_api_request: ResponsesAPIOptionalRequestParams = responses_api_request self.custom_llm_provider: str | None = custom_llm_provider @@ -129,7 +143,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): # repeated Pydantic attribute access in end-of-stream assembly. self.collected_chat_completion_chunks: list[dict[str, object]] = [] self.finished: bool = False - self.litellm_logging_obj = litellm_custom_stream_wrapper.logging_obj + self.litellm_logging_obj: LiteLLMLoggingObj = litellm_custom_stream_wrapper.logging_obj self.sent_response_created_event: bool = False self.sent_response_in_progress_event: bool = False self.sent_output_item_added_event: bool = False @@ -1024,6 +1038,15 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): # Latch so wrappers (FallbackResponsesStreamWrapper) + proxy # container-ownership hook can read completed_response. self.completed_response = response_completed_event + if self._logs_completion: + dispatch_stream_success( + self.litellm_logging_obj, + response_completed_event, + start_time=None, + end_time=None, + cache_hit=None, + is_async=not sync_mode, + ) return response_completed_event else: if sync_mode: diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 5ca42c38622..71abef9e4b1 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -355,6 +355,46 @@ def _mid_stream_fallback_eligible(mapped_exception: Exception) -> bool: PRE_OUTPUT_LIFECYCLE_EVENT_TYPES: Final = frozenset({"response.created", "response.in_progress", "response.queued"}) +def dispatch_stream_success( + logging_obj: LiteLLMLoggingObj, + logging_response: object, + *, + start_time: datetime | None, + end_time: datetime | None, + cache_hit: bool | None, + is_async: bool, +) -> None: + if not is_async: + run_async_function( + async_function=logging_obj.async_success_handler, + result=logging_response, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + ) + executor.submit( + logging_obj.success_handler, + result=logging_response, + cache_hit=cache_hit, + start_time=start_time, + end_time=end_time, + ) + return + logging_coroutine: Final = logging_obj.dispatch_success_handlers( + logging_response, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + prefer_async_handlers=True, + ) + if getattr(logging_obj, "_on_deferred_stream_complete", None) is not None: + # End-of-stream guardrail scans write guardrail_information after + # the terminal event; dispatching now would snapshot metadata early. + logging_obj.hold_stream_completion((logging_coroutine,)) + return + spawn_logging_task(asyncio.get_running_loop(), logging_coroutine, task_name="Responses stream success logging") + + class BaseResponsesAPIStreamingIterator: """ Base class for streaming iterators that process responses from the Responses API. @@ -622,37 +662,14 @@ class BaseResponsesAPIStreamingIterator: end_time: Final = datetime.now() if self._on_complete is not None: self._on_complete(logging_response) - elif is_async: - logging_coroutine: Final = self.logging_obj.dispatch_success_handlers( + else: + dispatch_stream_success( + self.logging_obj, logging_response, start_time=self.start_time, end_time=end_time, cache_hit=self._completed_response_cache_hit, - prefer_async_handlers=True, - ) - deferred_dispatch_armed: Final = getattr(self.logging_obj, "_on_deferred_stream_complete", None) is not None - if deferred_dispatch_armed: - # End-of-stream guardrail scans write guardrail_information after - # the terminal event; dispatching now would snapshot metadata early. - self.logging_obj._deferred_stream_complete_args = (logging_coroutine,) - else: - spawn_logging_task( - asyncio.get_running_loop(), logging_coroutine, task_name="Responses stream success logging" - ) - else: - run_async_function( - async_function=self.logging_obj.async_success_handler, - result=logging_response, - start_time=self.start_time, - end_time=end_time, - cache_hit=self._completed_response_cache_hit, - ) - executor.submit( - self.logging_obj.success_handler, - result=logging_response, - cache_hit=self._completed_response_cache_hit, - start_time=self.start_time, - end_time=end_time, + is_async=is_async, ) self._run_post_success_hooks(end_time=end_time) diff --git a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py index 7e5d0f6342e..703041a9f81 100644 --- a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py @@ -18,6 +18,7 @@ import asyncio import logging from collections.abc import Callable, Coroutine, Mapping from datetime import datetime +from functools import partial from typing import Any, Final, cast from unittest.mock import AsyncMock, MagicMock, patch @@ -59,6 +60,7 @@ def _attach_mock_success_dispatch(mock_logging_obj, async_success_fn): mock_logging_obj.dispatch_success_handlers = dispatch_success_handlers mock_logging_obj.async_success_handler = async_success_fn + mock_logging_obj.hold_stream_completion = partial(Logging.hold_stream_completion, mock_logging_obj) async def _wait_until(condition: Callable[[], bool]) -> None: @@ -617,6 +619,7 @@ class TestDeferredStreamingClosure: response args on logging_obj instead of calling the closure directly.""" mock_logging_obj = MagicMock() mock_logging_obj._on_deferred_stream_complete = MagicMock() + mock_logging_obj.hold_stream_completion = partial(Logging.hold_stream_completion, mock_logging_obj) resp = await litellm.acompletion( model="gpt-3.5-turbo", @@ -1434,6 +1437,7 @@ class TestResponsesIteratorDeferredLogging: logging_obj = MagicMock() logging_obj.dispatch_success_handlers = dispatch_success_handlers + logging_obj.hold_stream_completion = partial(Logging.hold_stream_completion, logging_obj) return logging_obj, recorded @pytest.mark.asyncio diff --git a/tests/unit/responses/litellm_completion_transformation/test_google_ai_studio_responses_bridge.py b/tests/unit/responses/litellm_completion_transformation/test_google_ai_studio_responses_bridge.py index 8aa03d7629a..47af56a9f33 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_google_ai_studio_responses_bridge.py +++ b/tests/unit/responses/litellm_completion_transformation/test_google_ai_studio_responses_bridge.py @@ -23,9 +23,7 @@ async def test_mock_basic_google_ai_studio_responses_api_with_tools(): choices=[ litellm.utils.Choices( index=0, - message=litellm.utils.Message( - role="assistant", content="Test response" - ), + message=litellm.utils.Message(role="assistant", content="Test response"), finish_reason="stop", ) ], diff --git a/tests/unit/responses/litellm_completion_transformation/test_handler.py b/tests/unit/responses/litellm_completion_transformation/test_handler.py index 6897eadaa53..c64a2689f0b 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_handler.py +++ b/tests/unit/responses/litellm_completion_transformation/test_handler.py @@ -1,17 +1,24 @@ +import asyncio import json +import math from collections.abc import Mapping, Sequence +from datetime import datetime from types import SimpleNamespace from typing import Final from unittest.mock import patch import httpx import pytest +import respx +from pydantic import TypeAdapter import litellm +from litellm.integrations.custom_logger import CustomLogger from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.responses.litellm_completion_transformation.handler import ( LiteLLMCompletionTransformationHandler, ) +from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import ADDRESSED_RESPONSE_ID_FIELD @@ -245,10 +252,88 @@ async def test_bridged_next_turn_replays_a_turn_stored_under_the_id_the_client_h ) assert provider.request_body is not None, "the bridged turn never reached the provider" - assert [ - (message["role"], json.dumps(message["content"])) for message in provider.request_body["messages"] - ] == [ + assert [(message["role"], json.dumps(message["content"])) for message in provider.request_body["messages"]] == [ ("user", json.dumps([{"type": "text", "text": "My favorite number is 7."}])), ("assistant", json.dumps([{"type": "text", "text": "Noted."}])), ("user", json.dumps([{"type": "text", "text": "Double it"}])), ] + + +_LOGGED_ROW: Final = TypeAdapter(dict[str, object]) + + +class _SuccessRowRecorder(CustomLogger): + rows: tuple[dict[str, object], ...] = () + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.rows = (*self.rows, _LOGGED_ROW.validate_python(kwargs["standard_logging_object"])) + + +@pytest.mark.asyncio +async def test_bridged_responses_stream_logs_one_responses_row_billed_from_provider_usage( + monkeypatch: pytest.MonkeyPatch, +): + recorder: Final = _SuccessRowRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + prices: Final = {"input_cost_per_token": 1e-6, "output_cost_per_token": 4e-6, "cache_read_input_token_cost": 1e-7} + litellm.register_model( + {"deepseek/bridge-stream-model": {"litellm_provider": "deepseek", "mode": "chat", "max_tokens": 100, **prices}} + ) + chunk: Final[dict[str, object]] = {"id": "chatcmpl-stream", "object": "chat.completion.chunk", "created": 1} + chunks: Final[tuple[dict[str, object], ...]] = ( + {**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "streamed"}}]}, + {**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + { + **chunk, + "choices": [], + "usage": { + "prompt_tokens": 100, + "completion_tokens": 20, + "total_tokens": 120, + "prompt_tokens_details": {"cached_tokens": 40}, + }, + }, + ) + sse: Final = ("".join(f"data: {json.dumps(c)}\n\n" for c in chunks) + "data: [DONE]\n\n").encode() + + with respx.mock(assert_all_called=False) as upstream: + chat_route: Final = upstream.post("https://example.invalid/v1/chat/completions").mock( + return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, content=sse) + ) + stream: Final = await litellm.aresponses( + model="deepseek/bridge-stream-model", + api_base="https://example.invalid/v1", + api_key="x", + input="hi", + stream=True, + ) + assert isinstance(stream, ResponsesAPIStreamingIterator) + async for _ in stream: + pass + for _ in range(100): + if recorder.rows: + break + await asyncio.sleep(0.01) + await asyncio.sleep(0.1) + + assert stream.completed_response is not None + assert chat_route.call_count == 1 + assert len(recorder.rows) == 1 + row: Final = recorder.rows[0] + logged: Final = _LOGGED_ROW.validate_python(row["response"]) + assert logged["object"] == "response" + assert logged["status"] == "completed" + assert row["prompt_tokens"] == 100 + assert row["completion_tokens"] == 20 + cost: Final = row["response_cost"] + assert isinstance(cost, float) + assert math.isclose( + cost, + 60 * prices["input_cost_per_token"] + + 40 * prices["cache_read_input_token_cost"] + + 20 * prices["output_cost_per_token"], + ) diff --git a/tests/unit/responses/litellm_completion_transformation/test_image_generation_output.py b/tests/unit/responses/litellm_completion_transformation/test_image_generation_output.py index 80e5335b4d4..dc4122a5c7f 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_image_generation_output.py +++ b/tests/unit/responses/litellm_completion_transformation/test_image_generation_output.py @@ -23,33 +23,20 @@ class TestExtractBase64FromDataUrl: def test_extracts_base64_from_data_url(self): """Should extract pure base64 from data URL with prefix""" data_url = "data:image/png;base64,iVBORw0KGgoAAAANS" - result = LiteLLMCompletionResponsesConfig._extract_base64_from_data_url( - data_url - ) + result = LiteLLMCompletionResponsesConfig._extract_base64_from_data_url(data_url) assert result == "iVBORw0KGgoAAAANS" def test_returns_base64_as_is_if_no_prefix(self): """Should return base64 as-is if no data: prefix""" pure_base64 = "iVBORw0KGgoAAAANS" - result = LiteLLMCompletionResponsesConfig._extract_base64_from_data_url( - pure_base64 - ) + result = LiteLLMCompletionResponsesConfig._extract_base64_from_data_url(pure_base64) assert result == pure_base64 def test_handles_invalid_inputs(self): """Should return None for empty/None/malformed inputs""" - assert ( - LiteLLMCompletionResponsesConfig._extract_base64_from_data_url("") is None - ) - assert ( - LiteLLMCompletionResponsesConfig._extract_base64_from_data_url(None) is None - ) - assert ( - LiteLLMCompletionResponsesConfig._extract_base64_from_data_url( - "data:image/png;base64" - ) - is None - ) + assert LiteLLMCompletionResponsesConfig._extract_base64_from_data_url("") is None + assert LiteLLMCompletionResponsesConfig._extract_base64_from_data_url(None) is None + assert LiteLLMCompletionResponsesConfig._extract_base64_from_data_url("data:image/png;base64") is None class TestExtractImageGenerationOutputItems: @@ -75,10 +62,8 @@ class TestExtractImageGenerationOutputItems: mock_choice.message = mock_message mock_choice.finish_reason = "stop" - result = ( - LiteLLMCompletionResponsesConfig._extract_image_generation_output_items( - choice=mock_choice, - ) + result = LiteLLMCompletionResponsesConfig._extract_image_generation_output_items( + choice=mock_choice, ) assert len(result) == 2 @@ -99,10 +84,8 @@ class TestExtractImageGenerationOutputItems: mock_choice.message = mock_message mock_choice.finish_reason = "stop" - result = ( - LiteLLMCompletionResponsesConfig._extract_image_generation_output_items( - choice=mock_choice, - ) + result = LiteLLMCompletionResponsesConfig._extract_image_generation_output_items( + choice=mock_choice, ) assert result == [] @@ -122,10 +105,8 @@ class TestExtractImageGenerationOutputItems: mock_choice.message = mock_message mock_choice.finish_reason = "length" - result = ( - LiteLLMCompletionResponsesConfig._extract_image_generation_output_items( - choice=mock_choice, - ) + result = LiteLLMCompletionResponsesConfig._extract_image_generation_output_items( + choice=mock_choice, ) assert result[0].status == "incomplete" @@ -200,10 +181,7 @@ class TestImageGenerationOutputItemIds: def _choice_with_images(self, count): mock_message = Mock(spec=Message) - mock_message.images = [ - {"image_url": {"url": f"data:image/png;base64,IMG{idx}"}} - for idx in range(count) - ] + mock_message.images = [{"image_url": {"url": f"data:image/png;base64,IMG{idx}"}} for idx in range(count)] mock_choice = Mock(spec=Choices) mock_choice.message = mock_message mock_choice.finish_reason = "stop" diff --git a/tests/unit/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py b/tests/unit/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py index 5e001bdbbbb..cf3befe6e35 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py +++ b/tests/unit/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py @@ -178,16 +178,12 @@ class TestEncryptedReasoningRoundTrip: item = { "type": "reasoning", "id": "rs_1", - "encrypted_content": json.dumps( - [{"type": "thinking", "thinking": "hidden", "signature": "sig-one"}] - ), + "encrypted_content": json.dumps([{"type": "thinking", "thinking": "hidden", "signature": "sig-one"}]), } messages = _transform_item(item) assert len(messages) == 1 assert messages[0]["content"] is None - assert messages[0]["thinking_blocks"] == [ - {"type": "thinking", "thinking": "hidden", "signature": "sig-one"} - ] + assert messages[0]["thinking_blocks"] == [{"type": "thinking", "thinking": "hidden", "signature": "sig-one"}] def test_unsigned_blocks_dropped(self): """Blocks without a signature or redacted payload are not replayed.""" @@ -229,9 +225,7 @@ class TestEncryptedReasoningRoundTrip: ) assert len(messages) == 1 assert messages[0]["reasoning_content"] == "look it up" - assert messages[0]["thinking_blocks"] == [ - {"type": "thinking", "thinking": "hidden", "signature": "sig-one"} - ] + assert messages[0]["thinking_blocks"] == [{"type": "thinking", "thinking": "hidden", "signature": "sig-one"}] assert len(messages[0]["tool_calls"]) == 1 def test_replayed_blocks_precede_existing_blocks(self): diff --git a/tests/unit/responses/litellm_completion_transformation/test_reasoning_items.py b/tests/unit/responses/litellm_completion_transformation/test_reasoning_items.py index 093d1744418..d3229913cc0 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_reasoning_items.py +++ b/tests/unit/responses/litellm_completion_transformation/test_reasoning_items.py @@ -58,5 +58,7 @@ def test_a_reasoning_item_is_litellm_minted_by_its_id_or_by_its_encoded_thinking "encrypted_content": A_PROVIDER_OWNED_ENCRYPTED_BLOB, } ) - assert not is_litellm_minted_reasoning_item({"type": "message", "id": mint_reasoning_item_id(), "role": "assistant"}) + assert not is_litellm_minted_reasoning_item( + {"type": "message", "id": mint_reasoning_item_id(), "role": "assistant"} + ) assert not is_litellm_minted_reasoning_item("a bare string input") diff --git a/tests/unit/responses/litellm_completion_transformation/test_session_handler_with_cold_storage.py b/tests/unit/responses/litellm_completion_transformation/test_session_handler_with_cold_storage.py index 1565820cbbd..c19b96a35ed 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_session_handler_with_cold_storage.py +++ b/tests/unit/responses/litellm_completion_transformation/test_session_handler_with_cold_storage.py @@ -78,9 +78,7 @@ class TestColdStorageObjectKeyIntegration: metadata = {"user_api_key": "test_key", "user_api_key_team_id": "test_team"} # Call the function - result = _get_spend_logs_metadata( - metadata=metadata, cold_storage_object_key="test/path/to/object.json" - ) + result = _get_spend_logs_metadata(metadata=metadata, cold_storage_object_key="test/path/to/object.json") # Verify the object key is stored in the result assert result.get("cold_storage_object_key") == "test/path/to/object.json" @@ -104,11 +102,7 @@ class TestColdStorageObjectKeyIntegration: } # Test the extraction method - object_key = ( - ResponsesSessionHandler._get_cold_storage_object_key_from_spend_log( - spend_log - ) - ) + object_key = ResponsesSessionHandler._get_cold_storage_object_key_from_spend_log(spend_log) assert object_key == "test/path/to/object.json" @@ -128,11 +122,7 @@ class TestColdStorageObjectKeyIntegration: } # Test the extraction method - object_key = ( - ResponsesSessionHandler._get_cold_storage_object_key_from_spend_log( - spend_log - ) - ) + object_key = ResponsesSessionHandler._get_cold_storage_object_key_from_spend_log(spend_log) assert object_key == "test/path/to/object.json" @@ -148,20 +138,17 @@ class TestColdStorageObjectKeyIntegration: # Mock the custom logger mock_logger = AsyncMock() - mock_logger.get_proxy_server_request_from_cold_storage_with_object_key = ( - AsyncMock(return_value={"test": "data"}) + mock_logger.get_proxy_server_request_from_cold_storage_with_object_key = AsyncMock( + return_value={"test": "data"} ) with ( - patch.object( - handler, "_select_custom_logger_for_cold_storage", return_value="s3_v2" - ), + patch.object(handler, "_select_custom_logger_for_cold_storage", return_value="s3_v2"), patch( "litellm.logging_callback_manager.get_active_custom_logger_for_callback_name", return_value=mock_logger, ), ): - result = await handler.get_proxy_server_request_from_cold_storage_with_object_key( object_key="test/path/to/object.json" ) @@ -184,9 +171,7 @@ class TestColdStorageObjectKeyIntegration: s3_logger = S3Logger(s3_bucket_name="test-bucket") # Mock the _download_object_from_s3 method - with patch.object( - s3_logger, "_download_object_from_s3", return_value={"test": "data"} - ) as mock_download: + with patch.object(s3_logger, "_download_object_from_s3", return_value={"test": "data"}) as mock_download: result = await s3_logger.get_proxy_server_request_from_cold_storage_with_object_key( object_key="test/path/to/object.json" ) diff --git a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 1a3430224d1..d171b4d5b08 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -35,9 +35,7 @@ from litellm.types.utils import ( ) CHAT_COMPLETION_ID = "chatcmpl-77d33d09-effa-4cd2-9c0d-c742d4358256" -RESPONSE_ID_EVENT_TYPES = frozenset( - {"response.created", "response.in_progress", "response.completed"} -) +RESPONSE_ID_EVENT_TYPES = frozenset({"response.created", "response.in_progress", "response.completed"}) def _chunk(content: str, finish_reason: str | None = None) -> ModelResponseStream: @@ -90,11 +88,7 @@ def _build_iterator(chunks) -> LiteLLMCompletionStreamingIterator: def _response_ids(events) -> list[str]: - return [ - event.response.id - for event in events - if getattr(event, "type", None) in RESPONSE_ID_EVENT_TYPES - ] + return [event.response.id for event in events if getattr(event, "type", None) in RESPONSE_ID_EVENT_TYPES] def test_tool_call_delta_is_emitted_as_responses_events(): @@ -225,9 +219,7 @@ async def test_web_search_stream_preserves_hosted_and_client_calls(sync_mode, to content=None, provider_specific_fields={ "web_search_results": [first_result], - "web_search_calls": [ - build_web_search_call(call_ids[0], {"query": "one"}, first_result) - ] + "web_search_calls": [build_web_search_call(call_ids[0], {"query": "one"}, first_result)] if tool_type != "function" and first_result["type"] == "web_search_tool_result" else [], }, @@ -326,9 +318,7 @@ async def test_web_search_stream_preserves_hosted_and_client_calls(sync_mode, to assert {event["output_index"] for event in search_events} == {completed_indexes[item["id"]]} assert search_events[0]["item"]["status"] == "in_progress" assert search_events[-1]["item"] == item - assert item["status"] == ( - "failed" if result_kind == "error" and call_id.endswith("01Search") else "completed" - ) + assert item["status"] == ("failed" if result_kind == "error" and call_id.endswith("01Search") else "completed") assert item["action"]["type"] == "search" assert item["action"]["query"] == ("one" if call_id.endswith("01Search") else "two") assert item["action"]["queries"] == [item["action"]["query"]] @@ -437,9 +427,7 @@ def test_tool_call_arguments_are_chunked_to_match_openai_behavior(): ) # Create a chunk with a large arguments string that should be split - large_arguments = ( - '{"param1": "value1", "param2": "value2", "param3": "value3"}' # 67 chars - ) + large_arguments = '{"param1": "value1", "param2": "value2", "param3": "value3"}' # 67 chars chunk = ModelResponseStream( id="chunk-1", created=123, @@ -494,9 +482,7 @@ def test_tool_call_arguments_are_chunked_to_match_openai_behavior(): # Keep draining pending events (expected: ceil(67 / 10) = 7 delta events) while iterator._pending_tool_events: - evt = iterator._transform_chat_completion_chunk_to_response_api_chunk( - empty_chunk - ) + evt = iterator._transform_chat_completion_chunk_to_response_api_chunk(empty_chunk) if evt and evt.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA: delta_events.append(evt) @@ -549,20 +535,12 @@ def test_tool_call_delta_without_id_uses_index_mapping(): while iterator._pending_tool_events: all_events.append(iterator._pending_tool_events.pop(0)) - delta_events = [ - evt - for evt in all_events - if evt.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA - ] + delta_events = [evt for evt in all_events if evt.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA] streamed_arguments = "".join(evt.delta for evt in delta_events) assert streamed_arguments == '{"location": "New York"}' - output_item_added_events = [ - evt - for evt in all_events - if evt.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED - ] + output_item_added_events = [evt for evt in all_events if evt.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED] assert len(output_item_added_events) == 1 assert output_item_added_events[0].item.id == "fc_call_abc123" assert output_item_added_events[0].item.call_id == "call_abc123" @@ -603,18 +581,10 @@ def test_parallel_tool_calls_without_ids_use_index_mapping(): while iterator._pending_tool_events: all_events.append(iterator._pending_tool_events.pop(0)) - output_item_added_events = [ - evt - for evt in all_events - if evt.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED - ] + output_item_added_events = [evt for evt in all_events if evt.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED] assert len(output_item_added_events) == 2 - delta_events = [ - evt - for evt in all_events - if evt.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA - ] + delta_events = [evt for evt in all_events if evt.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA] arguments_by_call_id = {} for evt in delta_events: arguments_by_call_id.setdefault(evt.item_id, "") @@ -667,11 +637,7 @@ def test_reused_index_with_new_call_id_marks_fallback_ambiguous(): while iterator._pending_tool_events: all_events.append(iterator._pending_tool_events.pop(0)) - delta_events = [ - evt - for evt in all_events - if evt.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA - ] + delta_events = [evt for evt in all_events if evt.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA] arguments_by_call_id = {} for evt in delta_events: arguments_by_call_id.setdefault(evt.item_id, "") @@ -711,22 +677,17 @@ def test_sync_streaming_events_share_the_chat_completion_response_id(): assert len(response_ids) == 3 assert len(set(response_ids)) == 1 assert ( - ResponsesAPIRequestUtils.decode_responses_api_response_id(response_ids[0])["response_id"] - == CHAT_COMPLETION_ID + ResponsesAPIRequestUtils.decode_responses_api_response_id(response_ids[0])["response_id"] == CHAT_COMPLETION_ID ) @pytest.mark.asyncio async def test_streaming_emits_every_chunk_after_priming_the_response_id(): - iterator = _build_iterator( - [_chunk("Hel"), _chunk("lo"), _chunk("!", finish_reason="stop")] - ) + iterator = _build_iterator([_chunk("Hel"), _chunk("lo"), _chunk("!", finish_reason="stop")]) events = [event async for event in iterator] - deltas = "".join( - event.delta for event in events if getattr(event, "type", None) == "response.output_text.delta" - ) + deltas = "".join(event.delta for event in events if getattr(event, "type", None) == "response.output_text.delta") assert deltas == "Hello!" @@ -772,7 +733,9 @@ async def test_leading_empty_choices_chunk_does_not_kill_the_stream(): event_types = [getattr(event, "type", None) for event in events] assert event_types.count(ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED) == 1 - assert "".join(event.delta for event in events if event.type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA) == "Hello!" + assert ( + "".join(event.delta for event in events if event.type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA) == "Hello!" + ) assert event_types[-1] == ResponsesAPIStreamEvents.RESPONSE_COMPLETED @@ -956,9 +919,7 @@ def test_streamed_unrecognized_tool_choice_is_echoed_as_auto() -> None: litellm_metadata={}, ) - response_events: Final = [ - event for event in iterator if getattr(event, "type", None) in RESPONSE_ID_EVENT_TYPES - ] + response_events: Final = [event for event in iterator if getattr(event, "type", None) in RESPONSE_ID_EVENT_TYPES] assert [event.response.tool_choice for event in response_events] == ["auto", "auto", "auto"] @@ -999,9 +960,7 @@ def _signed_thinking_and_text_in_one_chunk() -> tuple[ModelResponseStream, ...]: chunk: Final = _chunk("Hello!", finish_reason="stop") signed_block: Final = {"type": "thinking", "thinking": "", "signature": "sig_only"} delta: Final = Delta(**{**chunk.choices[0].delta.model_dump(), "thinking_blocks": [signed_block]}) - return ( - chunk.model_copy(update={"choices": [StreamingChoices(index=0, delta=delta, finish_reason="stop")]}), - ) + return (chunk.model_copy(update={"choices": [StreamingChoices(index=0, delta=delta, finish_reason="stop")]}),) StreamChunks = Callable[[], tuple[ModelResponseStream, ...]] @@ -1421,9 +1380,7 @@ async def test_anthropic_tool_search_call_streams_as_tool_search_call_item(sync_ model="claude-sonnet-5-5", object="chat.completion.chunk", choices=[ - StreamingChoices( - index=0, delta=delta, finish_reason="tool_calls" if index == len(deltas) - 1 else None - ) + StreamingChoices(index=0, delta=delta, finish_reason="tool_calls" if index == len(deltas) - 1 else None) ], ) for index, delta in enumerate(deltas) diff --git a/tests/unit/responses/litellm_completion_transformation/test_tool_output_order_preserved_for_gemini.py b/tests/unit/responses/litellm_completion_transformation/test_tool_output_order_preserved_for_gemini.py index 3dec1d571c4..3853a47422b 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_tool_output_order_preserved_for_gemini.py +++ b/tests/unit/responses/litellm_completion_transformation/test_tool_output_order_preserved_for_gemini.py @@ -60,11 +60,7 @@ def test_function_call_output_stays_adjacent_to_tool_call(): assistant_ok_idx = i elif isinstance(content, list): for block in content: - if ( - isinstance(block, dict) - and block.get("type") == "text" - and block.get("text") == "Done." - ): + if isinstance(block, dict) and block.get("type") == "text" and block.get("text") == "Done.": assistant_ok_idx = i break