diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index a294bade654..b7af8b6e9ec 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -61,6 +61,8 @@ - {id: quota_management.spend_tracking.messages_bridge.keeps_cache_tokens, module: quota_management, tier: P1, behavior: spend_tracking, variant: messages_bridge, assertions: [keeps_cache_tokens], exercised_on: [messages], source: "llms/anthropic/experimental_pass_through/responses_adapters/handler.py", rationale: "A /v1/messages request served by a Responses-only OpenAI model keeps its cache-read tokens and their discounted billing across the bridge (#34957)"} - {id: quota_management.spend_tracking.service_tier.bills_tier_rates, module: quota_management, tier: P1, behavior: spend_tracking, variant: service_tier, assertions: [bills_tier_rates], exercised_on: [chat_completions], source: "cost_calculator.py", rationale: "A priority service_tier call bills input, output, and reasoning at the deployment's *_priority rates and records the tier on the row (#35923, #35925)"} - {id: quota_management.spend_tracking.service_tier_stream.records_served_tier, module: quota_management, tier: P1, behavior: spend_tracking, variant: service_tier_stream, assertions: [records_served_tier], exercised_on: [chat_completions], source: "litellm_core_utils/streaming_chunk_builder_utils.py", fail_before_fix: proven, rationale: "A streamed call with no service_tier requested bills at the rates of the tier OpenAI stamps on its chunks and records that served tier on the row; the reassembled stream dropped the provider tier so the row recorded none and priced at the default rates"} +- {id: quota_management.spend_tracking.service_tier_stream.responses_records_served_tier, module: quota_management, tier: P1, behavior: spend_tracking, variant: service_tier_stream, assertions: [records_served_tier], exercised_on: [responses], source: "responses/streaming_iterator.py", rationale: "A streamed /v1/responses call bills at the tier carried on the response.completed event's inner response and records that served tier on the spend row"} +- {id: quota_management.spend_tracking.service_tier_stream.messages_records_served_tier, module: quota_management, tier: P1, behavior: spend_tracking, variant: service_tier_stream, assertions: [records_served_tier], exercised_on: [messages], source: "llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py", rationale: "A streamed /v1/messages call on an OpenAI-backed deployment bills at the tier OpenAI served; the Anthropic wire format has no tier field, so the spend row is the only record of it"} - {id: quota_management.spend_tracking.cost_headers.additive_components, module: quota_management, tier: P1, behavior: spend_tracking, variant: cost_headers, assertions: [additive_components], exercised_on: [chat_completions], source: "proxy/common_request_processing.py", rationale: "The x-litellm-response-cost-* component headers sum to the total, input covers only fresh tokens, and reasoning stays a subset of output (#36965)"} - {id: quota_management.spend_tracking.passthrough_stream.injects_usage_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: passthrough_stream, assertions: [injects_usage_cost], exercised_on: [openai_passthrough], source: "proxy/pass_through_endpoints/streaming_handler.py", rationale: "With include_cost_in_streaming_usage on, the /openai passthrough's final streaming usage frame carries the proxy-computed cost (#36503). Uncovered: the flag is only settable in litellm_settings, and the shared e2e stack does not turn it on yet"} - {id: quota_management.spend_tracking.websearch_interception.bills_under_request_session, module: quota_management, tier: P1, behavior: spend_tracking, variant: websearch_interception, assertions: [bills_under_request_session], exercised_on: [messages], source: "integrations/websearch_interception/handler.py", fail_before_fix: proven, rationale: "A web_search server tool the proxy intercepts into litellm.asearch writes its own asearch spend row, and that row carries the parent request's session_id so the session view counts the search and its cost next to the turn that triggered it (LIT-8063)"} diff --git a/tests/e2e/models.py b/tests/e2e/models.py index bd7f5171172..6a4312d5279 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -568,6 +568,17 @@ class AnthropicMessagesBody(BaseModel): cache: dict[str, bool] | None = {"no-cache": True} +class ResponsesStreamBody(BaseModel): + """POST /v1/responses body in the subset the spend tests stream with. + `input` stays a plain string: the tests only drive single-turn prompts.""" + + model: str + input: str + stream: bool = True + max_output_tokens: int | None = None + cache: dict[str, bool] | None = {"no-cache": True} + + class CountTokensBody(BaseModel): """POST /v1/messages/count_tokens body: the /v1/messages shape minus max_tokens (the endpoint only counts the prompt).""" diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 51ea9fbe7bd..408a35318aa 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -83,6 +83,7 @@ from models import ( OcrResponse, RerankBody, RerankResponse, + ResponsesStreamBody, RouterCurrentValues, RouterSettingsResponse, SearchToolCreateBody, @@ -970,6 +971,9 @@ class ProxyClient: def messages_stream(self, key: str, body: AnthropicMessagesBody) -> StreamingResponse: return self.transport.stream("/v1/messages", headers=self.transport.bearer(key), json=body) + def responses_stream(self, key: str, body: ResponsesStreamBody) -> StreamingResponse: + return self.transport.stream("/v1/responses", headers=self.transport.bearer(key), json=body) + def embed(self, key: str, body: EmbedBody) -> Result[EmbedResponse]: return self.transport.post( "/embeddings", diff --git a/tests/e2e/quota_management/spend_tracking/test_service_tier_pricing_e2e.py b/tests/e2e/quota_management/spend_tracking/test_service_tier_pricing_e2e.py index 008f2d8d350..2502c3cad89 100644 --- a/tests/e2e/quota_management/spend_tracking/test_service_tier_pricing_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_service_tier_pricing_e2e.py @@ -30,12 +30,20 @@ from cost_rows import ( assert_fresh_tokens_billed_at, assert_total_is_sum_of_components, poll_cost_row, + poll_cost_row_where, register_priced_model, ) from e2e_config import CHEAP_OPENAI_MODEL, unique_marker from e2e_http import unwrap from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, ChatStreamOptions, LiteLLMParamsBody +from models import ( + AnthropicMessagesBody, + ChatBody, + ChatMessage, + ChatStreamOptions, + LiteLLMParamsBody, + ResponsesStreamBody, +) from pydantic import BaseModel from spend_e2e_client import SpendClient @@ -60,6 +68,20 @@ class _StreamChunk(BaseModel): service_tier: str | None = None +class _CompletedResponseObject(BaseModel): + id: str | None = None + service_tier: str | None = None + + +class _ResponsesStreamEvent(BaseModel): + type: str | None = None + response: _CompletedResponseObject | None = None + + +class _MessagesStreamEvent(BaseModel): + type: str | None = None + + def _stream_chunks(events: list[str]) -> list[_StreamChunk]: return [_StreamChunk.model_validate_json(event) for event in events if event.strip() != "[DONE]"] @@ -223,3 +245,76 @@ class TestServiceTierPricing: f"{len(missing)} of {len(chunks)} relayed chunks dropped the provider's service_tier " f"{served_tier!r}: {missing}" ) + + @pytest.mark.covers("quota_management.spend_tracking.service_tier_stream.responses_records_served_tier") + def test_responses_stream_records_the_served_tier( + self, client: SpendClient, resources: ResourceManager, scoped_key: str + ) -> None: + model = register_priced_model( + client.proxy, + resources, + "tier-responses-stream", + LiteLLMParamsBody(model=STREAM_BACKEND, api_key=OPENAI_API_KEY), + ) + + result = client.proxy.responses_stream( + scoped_key, + ResponsesStreamBody(model=model, input=f"{unique_marker()} reply with one word"), + ) + assert result.ok and result.stream_events, ( + f"streamed responses call failed (status {result.status_code}): {result.body[:300]}" + ) + + events = [_ResponsesStreamEvent.model_validate_json(event) for event in result.stream_events] + completed = next((event for event in reversed(events) if event.type == "response.completed"), None) + assert completed is not None and completed.response is not None, ( + f"no response.completed event in the stream: {[e.type for e in events]}" + ) + served_tier = completed.response.service_tier + assert served_tier, f"response.completed carried no service_tier: {completed.response}" + assert served_tier in TIER_INPUT_RATES, f"no custom rate registered for served tier {served_tier!r}" + assert completed.response.id, f"response.completed carried no id: {completed.response}" + + row = poll_cost_row(client.proxy, completed.response.id) + assert row is not None, f"no spend row with a cost breakdown landed for {completed.response.id}" + assert row.breakdown.service_tier == served_tier, ( + f"response.completed served tier {served_tier!r} but the bill records " + f"pricing basis {row.breakdown.service_tier!r}" + ) + + @pytest.mark.covers("quota_management.spend_tracking.service_tier_stream.messages_records_served_tier") + def test_messages_stream_records_the_served_tier( + self, client: SpendClient, resources: ResourceManager, scoped_key: str + ) -> None: + model = register_priced_model( + client.proxy, + resources, + "tier-messages-stream", + LiteLLMParamsBody(model=STREAM_BACKEND, api_key=OPENAI_API_KEY), + ) + + result = client.proxy.messages_stream( + scoped_key, + AnthropicMessagesBody( + model=model, + messages=[ChatMessage(role="user", content=f"{unique_marker()} reply with one word")], + max_tokens=64, + stream=True, + ), + ) + assert result.ok and result.stream_events, ( + f"streamed messages call failed (status {result.status_code}): {result.body[:300]}" + ) + + events = [_MessagesStreamEvent.model_validate_json(event) for event in result.stream_events] + assert any(event.type == "message_delta" for event in events), ( + f"the anthropic stream emitted no message_delta: {[e.type for e in events]}" + ) + + row = poll_cost_row_where(client.proxy, scoped_key, lambda r: r.spend is not None and r.spend > 0) + assert row is not None, f"no spend row with a cost breakdown landed for the streamed messages call on {model}" + served_tier = row.breakdown.service_tier + assert served_tier in TIER_INPUT_RATES and served_tier is not None, ( + "the anthropic wire format carries no service_tier, so the bill is the only record of " + f"the tier OpenAI served; the row recorded pricing basis {served_tier!r}" + ) diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 23c01841b1b..c33b01be09e 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -8694,3 +8694,55 @@ async def test_async_failure_handler_delivers_failure_payload_to_custom_logger() assert "smoke-failure" in payload["error_str"] assert payload["model"] == "openai/gpt-5.6" assert events.empty() + + +def test_responses_completed_event_bills_the_served_service_tier(): + """The served service_tier on response.completed's inner ResponsesAPIResponse + must reach the cost calculator, so a priority-served stream prices at the + priority rates instead of the default tier's.""" + logging_obj: Final = LitellmLogging( + model="openai/gpt-5.1", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="aresponses", + start_time=time.time(), + litellm_call_id="resp-served-tier", + function_id="resp-served-tier", + ) + logging_obj.update_environment_variables( + model="openai/gpt-5.1", + user="", + optional_params={}, + litellm_params={}, + custom_llm_provider="openai", + ) + inner: Final = ResponsesAPIResponse( + id="resp-served-tier", + created_at=1, + object="response", + status="completed", + model="gpt-5.1", + output=[], + usage=ResponseAPIUsage(input_tokens=10, output_tokens=20, total_tokens=30), + service_tier="priority", + ) + event: Final = ResponseCompletedEvent(type="response.completed", response=inner) + + cost: Final = logging_obj._response_cost_calculator(result=event) # pyright: ignore[reportPrivateUsage] # parity with the suite's own direct calls + + billed_response: Final = ModelResponse( + model="gpt-5.1", + usage=litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30), + ) + tier_cost: Final = litellm.completion_cost( + completion_response=billed_response, + model="openai/gpt-5.1", + service_tier="priority", + ) + default_cost: Final = litellm.completion_cost( + completion_response=billed_response, + model="openai/gpt-5.1", + ) + + assert cost == pytest.approx(tier_cost) + assert cost > default_cost diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 7b74e69685c..9030b0fef8f 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -7277,6 +7277,48 @@ class TestStreamingClientDisconnectBilling: assert standard_logging_object["total_tokens"] > 0 assert standard_logging_object["response_cost"] >= 0.002 + @pytest.mark.asyncio + async def test_disconnect_bills_partial_spend_for_anthropic_adapter_stream(self): + """ + /v1/messages wraps the chat stream in AnthropicStreamWrapper, which + hides the CustomStreamWrapper's collected chunks behind + .completion_stream; the partial-billing helper reads response.chunks, + so the wrapper must delegate inward or a disconnect bills nothing. + """ + from litellm.llms.anthropic.experimental_pass_through.adapters.streaming_iterator import ( + AnthropicStreamWrapper, + ) + + recorder = _RecordingSuccessLogger() + original_callbacks = litellm.callbacks + litellm.callbacks = [recorder] + try: + response = await self._start_partial_stream() + setattr(response.chunks[-1], "service_tier", "priority") # noqa: B010 # pydantic extra, not a declared field + wrapped: Final = AnthropicStreamWrapper( + completion_stream=response, + model=response.model or "gpt-4o-mini", + ) + + billed: Final = await _bill_partial_streamed_spend_on_disconnect( + {"litellm_logging_obj": response.logging_obj}, + wrapped, + ) + + for _ in range(50): + if recorder.success_events: + break + await asyncio.sleep(0.1) + await asyncio.sleep(0.5) + finally: + litellm.callbacks = original_callbacks + + assert billed is True + assert len(recorder.success_events) == 1 + partial_response: Final = recorder.success_events[0]["response_obj"] + assert getattr(partial_response, "service_tier") == "priority" + assert partial_response.usage.total_tokens > 0 + @pytest.mark.asyncio async def test_completed_stream_does_not_double_bill_on_late_disconnect(self): recorder = _RecordingSuccessLogger()