diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 1d860b02875..7e161756f0d 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1370,28 +1370,45 @@ def _get_cost_breakdown_from_logging_obj( ) -def _classifier_cost_from_request_data(request_data: Mapping[str, object] | None) -> float | None: - """Cost of the auto-router's LLM classifier call, read from the request's routing_decision. +def _routing_decision_from_request_data(request_data: Mapping[str, object] | None) -> Mapping[str, object] | None: + """The auto-router's routing_decision recorded on the request, or None when no + pre-routing strategy ran. The pre-routing hook records the decision in `litellm_metadata` on messages/batch-style routes and in `metadata` on chat-style routes, so both buckets are consulted, in the same precedence `get_or_create_metadata_bucket` writes them. """ + return next( + ( + decision + for metadata_key in ("litellm_metadata", "metadata") + if request_data is not None + and isinstance(metadata := request_data.get(metadata_key), dict) + and isinstance(decision := metadata.get("routing_decision"), dict) + ), + None, + ) + + +def _classifier_cost_from_request_data(request_data: Mapping[str, object] | None) -> float | None: + """Cost of the auto-router's LLM classifier call, read from the request's routing_decision.""" from litellm.proxy.spend_tracking.savings import classifier_cost_from_decision - data: Final = request_data or {} - for metadata_key in ("litellm_metadata", "metadata"): - metadata = data.get(metadata_key) - if not isinstance(metadata, dict): - continue - decision = metadata.get("routing_decision") - if not isinstance(decision, dict): - continue - cost = classifier_cost_from_decision(decision) - if cost is None: - continue - return cost - return None + return classifier_cost_from_decision(_routing_decision_from_request_data(request_data)) + + +def _classifier_model_from_request_data(request_data: Mapping[str, object] | None) -> str | None: + """The model the auto-router's LLM classifier call used, or None when no classifier ran. + + `classifier_model` is stamped on the routing decision only when an LLM classifier call + decided the route, so presence is the fact that a second billable LLM call happened; + `classifier_cost` cannot carry that fact because it is also None for an unpriced + classifier model. The billable-request middleware counts the extra gateway request + off the header this feeds. + """ + decision: Final = _routing_decision_from_request_data(request_data) + model: Final = decision.get("classifier_model") if decision is not None else None + return model if isinstance(model, str) and model else None def _has_attribute_error_in_chain(exc: Exception) -> bool: @@ -1570,6 +1587,7 @@ class ProxyBaseLLMRequestProcessing: model_name: Final = ProxyBaseLLMRequestProcessing._get_deployment_model_name(litellm_logging_obj) classifier_cost: Final = _classifier_cost_from_request_data(request_data) + classifier_model: Final = _classifier_model_from_request_data(request_data) headers: Final = { "x-litellm-call-id": call_id, @@ -1613,6 +1631,7 @@ class ProxyBaseLLMRequestProcessing: str(cost_breakdown.tool_usage_cost) if cost_breakdown.tool_usage_cost is not None else None ), "x-litellm-classifier-cost": (str(classifier_cost) if classifier_cost is not None else None), + "x-litellm-classifier-model": classifier_model, "x-litellm-key-tpm-limit": str(user_api_key_dict.tpm_limit), "x-litellm-key-rpm-limit": str(user_api_key_dict.rpm_limit), "x-litellm-key-max-budget": str(user_api_key_dict.max_budget), diff --git a/litellm/proxy/middleware/billable_request_metrics_middleware.py b/litellm/proxy/middleware/billable_request_metrics_middleware.py index ac119e81d9c..d0ffaa8c32c 100644 --- a/litellm/proxy/middleware/billable_request_metrics_middleware.py +++ b/litellm/proxy/middleware/billable_request_metrics_middleware.py @@ -55,6 +55,15 @@ class GatewayRequestSink(Protocol): _MODEL_ID_HEADER: Final = b"x-litellm-model-id" +# An auto-routed request whose LLM classifier ran made two billable LLM calls behind +# one HTTP request. The classifier-model header is the edge-visible fact that the +# classifier call happened, and it is stamped only when that call succeeded, so the +# extra record is a 200 by construction. Heuristic-routed requests carry no header and +# keep counting as one. The dedicated route keeps the by-route breakdown explainable +# instead of silently inflating the serving route's count. +_CLASSIFIER_MODEL_HEADER: Final = b"x-litellm-classifier-model" +_LLM_CLASSIFIER_ROUTE: Final = "llm_classifier" + # Ordered: a longer suffix that shares an ending with a shorter one must come # first, e.g. "/chat/completions" before "/completions". This is the POST # inference surface that writes a SpendLogs row on success, so the exported @@ -179,9 +188,9 @@ def classify_billable_request(path: str, method: str = "POST") -> tuple[Billable return None -def _extract_model_id(headers: Sequence[tuple[bytes, bytes]]) -> str | None: +def _extract_header(headers: Sequence[tuple[bytes, bytes]], header_name: bytes) -> str | None: return next( - (value.decode("latin-1") for name, value in headers if name.lower() == _MODEL_ID_HEADER and value), + (value.decode("latin-1") for name, value in headers if name.lower() == header_name and value), None, ) @@ -261,24 +270,35 @@ class BillableRequestMetricsMiddleware: category, route = classification status_code = 0 model_id: str | None = None + classifier_model: str | None = None async def send_wrapper(message: Message) -> None: - nonlocal status_code, model_id + nonlocal status_code, model_id, classifier_model if message["type"] == "http.response.start": status_code = message["status"] - model_id = _extract_model_id(message.get("headers", [])) + headers: Final = message.get("headers", ()) + model_id = _extract_header(headers, _MODEL_ID_HEADER) + classifier_model = _extract_header(headers, _CLASSIFIER_MODEL_HEADER) await send(message) await self.app(scope, receive, send_wrapper) + classifier_ran: Final = classifier_model is not None + if sink is not None: try: sink.record(category=category, route=route, status_code=status_code) + if classifier_ran: + sink.record(category=BillableCategory.LLM, route=_LLM_CLASSIFIER_ROUTE, status_code=200) except Exception: # noqa: BLE001 -- metering must never fail a request that was already served verbose_proxy_logger.warning("gateway request metering failed for %s", route, exc_info=True) if recorder is not None and 200 <= status_code < 300: try: recorder.record(category=category, route=route, status_code=status_code, model_id=model_id) + if classifier_ran: + recorder.record( + category=BillableCategory.LLM, route=_LLM_CLASSIFIER_ROUTE, status_code=200, model_id=None + ) except Exception: # noqa: BLE001 -- metering must never fail a request that was already served verbose_proxy_logger.warning("billable request metering failed for %s", route, exc_info=True) diff --git a/tests/test_litellm/proxy/middleware/test_billable_request_metrics_middleware.py b/tests/test_litellm/proxy/middleware/test_billable_request_metrics_middleware.py index 9c61412bd6e..7e1dbcec216 100644 --- a/tests/test_litellm/proxy/middleware/test_billable_request_metrics_middleware.py +++ b/tests/test_litellm/proxy/middleware/test_billable_request_metrics_middleware.py @@ -20,9 +20,11 @@ from starlette.testclient import TestClient from litellm.proxy.db.gateway_request_tracking import GatewayRequestAccumulator from litellm.proxy.middleware.billable_request_metrics_middleware import ( + _CLASSIFIER_MODEL_HEADER, + _MODEL_ID_HEADER, BillableCategory, BillableRequestMetricsMiddleware, - _extract_model_id, + _extract_header, classify_billable_request, ) from litellm.proxy.middleware.in_flight_requests_middleware import ( @@ -40,9 +42,21 @@ class FakeRecorder: ) -def _make_app(recorder: Optional[FakeRecorder], status_code: int = 200, model_id: Optional[str] = None) -> Starlette: +def _make_app( + recorder: Optional[FakeRecorder], + status_code: int = 200, + model_id: Optional[str] = None, + classifier_model: Optional[str] = None, +) -> Starlette: async def handler(request: Request) -> Response: - headers = {"x-litellm-model-id": model_id} if model_id else {} + headers = { + key: value + for key, value in { + "x-litellm-model-id": model_id, + "x-litellm-classifier-model": classifier_model, + }.items() + if value + } return JSONResponse({}, status_code=status_code, headers=headers) paths = [ @@ -232,17 +246,17 @@ def test_chat_completions_not_misclassified_as_plain_completions(): # ── _extract_model_id ───────────────────────────────────────────────────────── -def test_extract_model_id_present(): +def test_extract_header_present(): headers = [(b"content-type", b"application/json"), (b"x-litellm-model-id", b"deploy-123")] - assert _extract_model_id(headers) == "deploy-123" + assert _extract_header(headers, _MODEL_ID_HEADER) == "deploy-123" -def test_extract_model_id_case_insensitive(): - assert _extract_model_id([(b"X-LiteLLM-Model-Id", b"deploy-9")]) == "deploy-9" +def test_extract_header_case_insensitive(): + assert _extract_header([(b"X-LiteLLM-Model-Id", b"deploy-9")], _MODEL_ID_HEADER) == "deploy-9" -def test_extract_model_id_absent(): - assert _extract_model_id([(b"content-type", b"application/json")]) is None +def test_extract_header_absent(): + assert _extract_header([(b"content-type", b"application/json")], _CLASSIFIER_MODEL_HEADER) is None # ── Middleware recording behaviour ──────────────────────────────────────────── @@ -478,8 +492,9 @@ def _make_sink_app( sink: Optional[FakeSink], status_code: int = 200, model_id: Optional[str] = None, + classifier_model: Optional[str] = None, ) -> Starlette: - app = _make_app(None, status_code=status_code, model_id=model_id) + app = _make_app(None, status_code=status_code, model_id=model_id, classifier_model=classifier_model) app.user_middleware.clear() app.add_middleware(BillableRequestMetricsMiddleware, recorder=recorder, sink=sink) return app @@ -585,3 +600,45 @@ def test_sink_factory_resolved_once_across_requests(): client.post("/v1/chat/completions") assert calls == [1] assert len(sink.calls) == 2 + + +# ── LLM classifier double-count ─────────────────────────────────────────────── + + +def test_classifier_header_counts_a_second_successful_llm_request(): + """An auto-routed request whose LLM classifier ran made two billable LLM calls, + so both sinks get a second record under the dedicated classifier route.""" + sink, recorder = FakeSink(), FakeRecorder() + TestClient( + _make_sink_app(recorder, sink, status_code=200, model_id="deploy-7", classifier_model="gpt-5-mini") + ).post("/v1/chat/completions") + assert sink.calls == [ + {"category": BillableCategory.LLM, "route": "/chat/completions", "status_code": 200}, + {"category": BillableCategory.LLM, "route": "llm_classifier", "status_code": 200}, + ] + assert recorder.calls == [ + {"category": BillableCategory.LLM, "route": "/chat/completions", "status_code": 200, "model_id": "deploy-7"}, + {"category": BillableCategory.LLM, "route": "llm_classifier", "status_code": 200, "model_id": None}, + ] + + +def test_classifier_success_counts_even_when_the_routed_call_fails(): + """The classifier call succeeded and billed before the upstream failed: SGR keeps + both facts, while enterprise billing stays gated on the served response's 2xx.""" + sink, recorder = FakeSink(), FakeRecorder() + TestClient(_make_sink_app(recorder, sink, status_code=503, classifier_model="gpt-5-mini")).post( + "/v1/chat/completions" + ) + assert [call["status_code"] for call in sink.calls] == [503, 200] + assert recorder.calls == [] + + +def test_classifier_request_folds_into_its_own_route_key(): + accumulator = GatewayRequestAccumulator() + TestClient(_make_sink_app(None, accumulator, status_code=200, classifier_model="gpt-5-mini")).post( + "/v1/chat/completions" + ) + snapshot = accumulator.drain() + assert {key.route for key in snapshot} == {"/chat/completions", "llm_classifier"} + assert sum(counts.successful_requests for counts in snapshot.values()) == 2 + assert sum(counts.failed_requests for counts in snapshot.values()) == 0 diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index df14224af5c..73cf12149b2 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -1124,6 +1124,56 @@ class TestProxyBaseLLMRequestProcessing: assert "x-litellm-classifier-cost" not in headers + @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) + def test_get_custom_headers_classifier_model_from_routing_decision(self, metadata_key): + """classifier_model is stamped only when an LLM classifier call decided the route, + so its header is the edge-visible fact the billable-request middleware counts the + second gateway request off. It must surface even when the classifier model is + unpriced (classifier_cost absent).""" + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict.tpm_limit = None + mock_user_api_key_dict.rpm_limit = None + mock_user_api_key_dict.max_budget = None + mock_user_api_key_dict.spend = 0 + + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + request_data={ + metadata_key: { + "routing_decision": {"cause": "llm_classifier", "classifier_model": "gpt-5-mini"}, + } + }, + ) + + assert headers["x-litellm-classifier-model"] == "gpt-5-mini" + assert "x-litellm-classifier-cost" not in headers + + @pytest.mark.parametrize( + "request_data", + [ + None, + {}, + {"metadata": {"routing_decision": {"cause": "heuristic_scorer"}}}, + {"metadata": {"routing_decision": {"cause": "llm_classifier", "classifier_model": ""}}}, + {"metadata": {"routing_decision": {"cause": "llm_classifier", "classifier_model": 7}}}, + ], + ) + def test_get_custom_headers_omits_classifier_model_without_a_classifier_call(self, request_data): + """A heuristic decision, a malformed value, or no decision at all must omit the + header entirely: its presence is what makes the request count twice.""" + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict.tpm_limit = None + mock_user_api_key_dict.rpm_limit = None + mock_user_api_key_dict.max_budget = None + mock_user_api_key_dict.spend = 0 + + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + request_data=request_data, + ) + + assert "x-litellm-classifier-model" not in headers + def test_get_cost_breakdown_from_logging_obj_helper(self): """ Test the helper function that extracts cost breakdown information.