diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 22a97a07d3b..3ecfc31f632 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -41,7 +41,11 @@ from litellm.exceptions import MidStreamFallbackError, RateLimitError from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.exception_mapping_utils import with_deployment_api_base -from litellm.litellm_core_utils.internal_call_metadata import requests_background, without_budget_reservation +from litellm.litellm_core_utils.internal_call_metadata import ( + is_background_response, + requests_background, + without_budget_reservation, +) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base from litellm.litellm_core_utils.llm_response_utils.get_headers import provider_request_id_headers @@ -975,6 +979,7 @@ class BaseResponsesAPIStreamingIterator: if ( not generated_text or requests_background(self.request_data) + or is_background_response(self._created_response) or self.logging_obj.model_call_details.get("combined_usage_object") is not None ): return diff --git a/tests/unit/responses/test_responses_streaming_iterator.py b/tests/unit/responses/test_responses_streaming_iterator.py index c572b4151b2..f4d6240e593 100644 --- a/tests/unit/responses/test_responses_streaming_iterator.py +++ b/tests/unit/responses/test_responses_streaming_iterator.py @@ -336,17 +336,12 @@ async def test_mid_stream_read_timeout_bills_the_tool_call_arguments_already_str assert payload["response_cost"] > 0 -@pytest.mark.parametrize( - "background_fields", - [{"background": True}, {"extra_body": {"background": True}}], - ids=["top_level", "extra_body"], -) @pytest.mark.asyncio async def test_mid_stream_failure_of_a_background_response_bills_nothing_since_retrieving_it_bills_the_whole_run( - monkeypatch: pytest.MonkeyPatch, background_fields: Mapping[str, object] + monkeypatch: pytest.MonkeyPatch, ) -> None: _, payload = await _failure_payload_after_mid_stream_timeout( - monkeypatch, _SEA_PROMPT, background_fields=background_fields + monkeypatch, _SEA_PROMPT, background_fields={"background": True} ) assert payload["status"] == "failure" @@ -469,25 +464,25 @@ class _UpstreamThatTimesOutAfterRelayingText: raise httpx.ReadTimeout("Timeout on reading data from socket") -class _UpstreamServedAtTierThatTimesOut: +class _UpstreamThatTimesOutAfterCreated: headers: dict[str, str] = {} - def __init__(self, served_tier: str) -> None: - self.served_tier = served_tier + def __init__(self, **created_fields: object) -> None: + self.created_fields = created_fields async def aiter_bytes(self) -> AsyncIterator[bytes]: created: Final = { **_CREATED_WITH_PRIORITY_TIER, - "response": {**_CREATED_RESPONSE_WITH_PRIORITY_TIER, "service_tier": self.served_tier}, + "response": {**_CREATED_RESPONSE_WITH_PRIORITY_TIER, **self.created_fields}, } for event in (created, _text_delta("Once upon a time", 1)): yield f"data: {json.dumps(event)}\n\n".encode() raise httpx.ReadTimeout("Timeout on reading data from socket") -async def _mid_stream_failure_cost( - monkeypatch: pytest.MonkeyPatch, requested_tier: str | None, served_tier: str -) -> float: +async def _mid_stream_failure_payload( + monkeypatch: pytest.MonkeyPatch, requested_tier: str | None, **created_fields: object +) -> StandardLoggingPayload: recorder: Final = FailureRecorder() monkeypatch.setattr(litellm, "failure_callback", [recorder]) monkeypatch.setattr(litellm, "_async_failure_callback", [recorder]) @@ -497,25 +492,35 @@ async def _mid_stream_failure_cost( optional_params={"stream": True} if requested_tier is None else {"stream": True, "service_tier": requested_tier}, litellm_params={"custom_llm_provider": "openai", "aresponses": True}, ) - iterator: Final = _live_iterator(logging_obj, upstream=_UpstreamServedAtTierThatTimesOut(served_tier)) + iterator: Final = _live_iterator(logging_obj, upstream=_UpstreamThatTimesOutAfterCreated(**created_fields)) with pytest.raises(litellm.Timeout): await _ChunkRecorder().drain(iterator) await iterator._await_pending_logging() (payload,) = recorder.failure_payloads - return payload["response_cost"] + return payload @pytest.mark.asyncio async def test_mid_stream_failure_bills_the_relayed_text_at_the_tier_the_provider_served( monkeypatch: pytest.MonkeyPatch, ) -> None: - flex_requested_default_served: Final = await _mid_stream_failure_cost(monkeypatch, "flex", "default") - nothing_requested_default_served: Final = await _mid_stream_failure_cost(monkeypatch, None, "default") + flex_requested: Final = await _mid_stream_failure_payload(monkeypatch, "flex", service_tier="default") + nothing_requested: Final = await _mid_stream_failure_payload(monkeypatch, None, service_tier="default") - assert nothing_requested_default_served > 0 - assert flex_requested_default_served == nothing_requested_default_served + assert nothing_requested["response_cost"] > 0 + assert flex_requested["response_cost"] == nothing_requested["response_cost"] + + +@pytest.mark.asyncio +async def test_mid_stream_failure_bills_nothing_when_the_provider_runs_the_response_in_background( + monkeypatch: pytest.MonkeyPatch, +) -> None: + payload: Final = await _mid_stream_failure_payload(monkeypatch, None, background=True) + + assert payload["status"] == "failure" + assert (payload["completion_tokens"], payload["response_cost"]) == (0, 0) def test_partial_response_on_disconnect_bills_the_relayed_text_at_the_served_service_tier():