diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 4bc847c923e..bbeeeb0be64 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -113,6 +113,14 @@ def _classifier_call_metadata(metadata: dict[str, Any] | None) -> dict[str, Any] } +def _effective_turn_off_message_logging(request_kwargs: dict[str, Any] | None) -> bool | None: + from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( + initialize_standard_callback_dynamic_params, + ) + + return initialize_standard_callback_dynamic_params(request_kwargs or {}).get("turn_off_message_logging") + + class DimensionScore: """Represents a score for a single dimension with optional signal.""" @@ -430,6 +438,7 @@ class ComplexityRouter(CustomLogger): # internal classifier call) is responsible for reconciling. request_metadata = (request_kwargs or {}).get("litellm_metadata") or (request_kwargs or {}).get("metadata") metadata = _classifier_call_metadata(request_metadata) + turn_off_message_logging = _effective_turn_off_message_logging(request_kwargs) proxy_server_request = { "body": { @@ -446,6 +455,7 @@ class ComplexityRouter(CustomLogger): timeout=llm_config.timeout_ms / 1000, metadata=metadata, proxy_server_request=proxy_server_request, + turn_off_message_logging=turn_off_message_logging, ) content = response.choices[0].message.content if not content: @@ -832,6 +842,7 @@ 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")) + turn_off_message_logging = _effective_turn_off_message_logging(request_kwargs) proxy_server_request = {"body": {"model": self.config.embedding_model, "input": [user_message]}} query_vector = ( await encoder.aencode_queries( @@ -839,6 +850,7 @@ class ComplexityRouter(CustomLogger): metadata=metadata, litellm_metadata=litellm_metadata, proxy_server_request=proxy_server_request, + turn_off_message_logging=turn_off_message_logging, ) )[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 1feeb150c87..f31ca32f4c5 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -1465,6 +1465,54 @@ class TestLLMClassifier: "REASONING", ] + @pytest.mark.asyncio + async def test_aclassify_propagates_top_level_turn_off_message_logging( + self, llm_complexity_router, mock_router_instance + ): + """A caller's top-level turn_off_message_logging must reach the classifier call. + + Without this, a caller who opts a request out of message logging still has their + prompt captured in full by the classifier's proxy_server_request: the spend-log + redaction gate (should_redact_message_logging) reads turn_off_message_logging off + the classifier call's own kwargs, and this internal call is not the caller's + request, so it never inherits the opt-out unless it's forwarded explicitly. + """ + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "SIMPLE"}')) + await llm_complexity_router.aclassify("secret prompt", request_kwargs={"turn_off_message_logging": True}) + call_kwargs = mock_router_instance.acompletion.call_args.kwargs + assert call_kwargs["turn_off_message_logging"] is True + + @pytest.mark.asyncio + async def test_aclassify_propagates_metadata_slot_turn_off_message_logging( + self, llm_complexity_router, mock_router_instance + ): + """turn_off_message_logging set inside metadata/litellm_metadata must also propagate. + + initialize_standard_callback_dynamic_params reads this flag from either the + top-level request kwargs or the metadata/litellm_metadata dicts (the same slots a + real HTTP request populates), so the classifier call must resolve it from there too. + """ + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "SIMPLE"}')) + await llm_complexity_router.aclassify( + "secret prompt", request_kwargs={"litellm_metadata": {"turn_off_message_logging": True}} + ) + call_kwargs = mock_router_instance.acompletion.call_args.kwargs + assert call_kwargs["turn_off_message_logging"] is True + + @pytest.mark.asyncio + async def test_aclassify_defaults_turn_off_message_logging_to_none( + self, llm_complexity_router, mock_router_instance + ): + """With no caller opt-out, the classifier call must not force redaction on or off. + + Passing None (rather than omitting the kwarg or defaulting to False) preserves the + existing header- and global-setting fallbacks in should_redact_message_logging. + """ + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "SIMPLE"}')) + await llm_complexity_router.aclassify("hi") + call_kwargs = mock_router_instance.acompletion.call_args.kwargs + assert call_kwargs["turn_off_message_logging"] is None + @pytest.mark.asyncio async def test_aclassify_strips_budget_reservation_from_classifier_metadata( self, llm_complexity_router, mock_router_instance @@ -2250,6 +2298,36 @@ class TestSemanticKeywordTierRules: assert body["model"] == "fake-embed" assert body["input"] == ["roll out my k8s cluster"] + @pytest.mark.asyncio + async def test_semantic_embedding_call_propagates_turn_off_message_logging(self, basic_config): + """A caller's turn_off_message_logging must reach the query embedding call. + + The embedding now captures the user's prompt in proxy_server_request, so a caller + who opts out of message logging must have that opt-out forwarded; otherwise the + embedding's spend-log row stores the prompt in the clear despite the parent request + being redacted, exposing it to anyone authorized to read the team's spend logs. + """ + 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={"turn_off_message_logging": True}, + messages=[{"role": "user", "content": "roll out my k8s cluster"}], + ) + assert fake_router.async_embedding_kwargs, "expected an embedding call for the prompt" + assert fake_router.async_embedding_kwargs[0]["turn_off_message_logging"] is True + @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.