From 2a84c397620d0eea027440950ed33925a1dce41e Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Wed, 29 Jul 2026 17:13:50 -0700 Subject: [PATCH] fix(complexity_router): capture the classifier request body in spend logs --- .../complexity_router/complexity_router.py | 17 +++++- .../router_strategy/test_complexity_router.py | 56 +++++++++++++++++++ 2 files changed, 72 insertions(+), 1 deletion(-) diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index e5268b5107b..836c2e9d4c0 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -429,12 +429,21 @@ class ComplexityRouter(CustomLogger): # internal classifier call) is responsible for reconciling. metadata = _classifier_call_metadata((request_kwargs or {}).get("litellm_metadata")) + proxy_server_request = { + "body": { + "model": llm_config.model, + "messages": [{"role": "user", "content": classification_prompt}], + "response_format": TierClassification.model_json_schema(), + } + } + response: ModelResponse = await self.litellm_router_instance.acompletion( model=llm_config.model, messages=[{"role": "user", "content": classification_prompt}], response_format=TierClassification, timeout=llm_config.timeout_ms / 1000, metadata=metadata, + proxy_server_request=proxy_server_request, ) content = response.choices[0].message.content if not content: @@ -821,8 +830,14 @@ class ComplexityRouter(CustomLogger): # key/team budget. Key/team attribution fields are preserved for spend logging. metadata = _classifier_call_metadata(request_kwargs.get("metadata")) litellm_metadata = _classifier_call_metadata(request_kwargs.get("litellm_metadata")) + proxy_server_request = {"body": {"model": self.config.embedding_model, "input": [user_message]}} query_vector = ( - await encoder.aencode_queries([user_message], metadata=metadata, litellm_metadata=litellm_metadata) + await encoder.aencode_queries( + [user_message], + metadata=metadata, + litellm_metadata=litellm_metadata, + proxy_server_request=proxy_server_request, + ) )[0] route_choice = await routelayer.acall(vector=query_vector) diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index ef70687bd97..a9b86c16c33 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -1417,6 +1417,29 @@ class TestLLMClassifier: call_kwargs = mock_router_instance.acompletion.call_args.kwargs assert call_kwargs["metadata"] == request_metadata + @pytest.mark.asyncio + async def test_aclassify_captures_request_body_in_proxy_server_request( + self, llm_complexity_router, mock_router_instance + ): + """The classifier call must supply proxy_server_request so its request body is logged. + + proxy_server_request["body"] is populated only by the proxy's HTTP ingress + middleware, which never runs for this internally-initiated router.acompletion + call. Without it _get_proxy_server_request_for_spend_logs_payload reads nothing + and stores "{}" for the request, so the classifier's spend-log row shows a + populated response but an empty request and the log cannot show which prompt + drove the tier decision. The captured body must carry the classification prompt + actually sent, so the classifier model, the classification prompt, and the user + text are all asserted here. + """ + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "COMPLEX"}')) + await llm_complexity_router.aclassify("explain quantum tunneling in depth") + call_kwargs = mock_router_instance.acompletion.call_args.kwargs + body = call_kwargs["proxy_server_request"]["body"] + assert body["model"] == "haiku-classifier" + assert body["messages"] == call_kwargs["messages"] + assert "explain quantum tunneling in depth" in body["messages"][0]["content"] + @pytest.mark.asyncio async def test_aclassify_strips_budget_reservation_from_classifier_metadata( self, llm_complexity_router, mock_router_instance @@ -2169,6 +2192,39 @@ class TestSemanticKeywordTierRules: assert fake_router.async_embedding_kwargs[0]["metadata"] == caller_metadata assert fake_router.async_embedding_kwargs[0]["litellm_metadata"] == caller_litellm_metadata + @pytest.mark.asyncio + async def test_semantic_embedding_call_captures_request_body_in_proxy_server_request(self, basic_config): + """The query embedding call must supply proxy_server_request so its request is logged. + + Like the LLM classifier, this embedding is fired internally and never passes + through the proxy's HTTP ingress middleware, so proxy_server_request is unset and + the embedding's spend-log row stores "{}" for the request while its response is + captured. The captured body must carry the embedded input so the log shows what + was classified. + """ + fake_router = FakeEmbeddingRouter() + config = { + **basic_config, + "keyword_tier_rules": [{"keywords": ["kubernetes deployment"], "tier": "REASONING"}], + "semantic_keyword_matching": True, + "embedding_model": "fake-embed", + "match_threshold": 0.5, + } + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=fake_router, + complexity_router_config=config, + ) + await router.async_pre_routing_hook( + model="test-model", + request_kwargs={}, + messages=[{"role": "user", "content": "roll out my k8s cluster"}], + ) + assert fake_router.async_embedding_kwargs, "expected an embedding call for the prompt" + body = fake_router.async_embedding_kwargs[0]["proxy_server_request"]["body"] + assert body["model"] == "fake-embed" + assert body["input"] == ["roll out my k8s cluster"] + @pytest.mark.asyncio async def test_semantic_embedding_call_strips_budget_reservation(self, basic_config): """The embedding call must not carry the parent request's budget reservation.