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