From 2a5d09ee87cdb4989e8c8b164450e969d5bf64a8 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Sat, 29 Aug 2026 10:58:27 -0700 Subject: [PATCH] fix(policy): let the AI policy suggester drop sampling params its model refuses (#38594) The suggester pins temperature=0.2 for tool-selection determinism and passed no drop_params, so an operator-supplied reasoning model whose only accepted temperature is 1 made litellm raise UnsupportedParamsError and the whole suggestion fail. The default gpt-4o-mini is unaffected; the failure needs the caller to name a model. Every other internal LLM call the proxy makes on a user's behalf already opts in through judge_acompletion, which sets drop_params=True on both dispatch paths. This was the one caller outside that contract, so the sampling preference is now advisory here too and the call degrades instead of dying. Resolves LIT-6352 --- .../policy_endpoints/ai_policy_suggester.py | 14 ++++ .../test_ai_policy_suggester.py | 77 +++++++++++++++++++ 2 files changed, 91 insertions(+) diff --git a/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py b/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py index 4bc53678c23..1096954536a 100644 --- a/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py +++ b/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py @@ -9,6 +9,7 @@ from typing import Final import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import DEFAULT_COMPETITOR_DISCOVERY_MODEL +from litellm.proxy._types import ProxyErrorTypes, ProxyException SUGGEST_TOOL: Final = { "type": "function", @@ -60,6 +61,18 @@ class AiPolicySuggester: system_prompt: Final = self._build_system_prompt(templates) user_prompt: Final = self._build_user_prompt(attack_examples, description) model = model or DEFAULT_COMPETITOR_DISCOVERY_MODEL + custom_llm_provider: Final = model.split("/", 1)[0] if "/" in model else None + supported_params: Final = litellm.get_supported_openai_params( + model=model, + custom_llm_provider=custom_llm_provider, + ) + if supported_params is not None and "tools" not in supported_params: + raise ProxyException( + message=(f"AI policy suggestion requires tool calling; model '{model}' does not support it"), + type=ProxyErrorTypes.validation_error.value, + param="model", + code=400, + ) try: response: Final = await litellm.acompletion( @@ -74,6 +87,7 @@ class AiPolicySuggester: "function": {"name": "select_policy_templates"}, }, temperature=0.2, + drop_params=True, ) tool_calls: Final = response.choices[0].message.tool_calls diff --git a/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py b/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py index e3893a66094..93dc429168f 100644 --- a/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py +++ b/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py @@ -7,6 +7,9 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +import litellm + +from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints.policy_endpoints.ai_policy_suggester import ( SUGGEST_TOOL, AiPolicySuggester, @@ -234,6 +237,7 @@ class TestAiPolicySuggester: call_kwargs = mock_acompletion.call_args.kwargs assert call_kwargs["model"] == "gpt-4o-mini" assert call_kwargs["temperature"] == 0.2 + assert call_kwargs["drop_params"] is True assert len(call_kwargs["tools"]) == 1 assert call_kwargs["tools"][0]["function"]["name"] == "select_policy_templates" assert ( @@ -242,3 +246,76 @@ class TestAiPolicySuggester: assert len(call_kwargs["messages"]) == 2 assert call_kwargs["messages"][0]["role"] == "system" assert call_kwargs["messages"][1]["role"] == "user" + + +class TestSuggesterRejectsModelsWithoutToolCalling: + @pytest.mark.asyncio + async def test_a_tools_less_model_is_rejected(self, local_model_cost_map): + with pytest.raises(ProxyException) as exc: + await AiPolicySuggester().suggest( + templates=SAMPLE_TEMPLATES, + attack_examples=["Ignore all previous instructions"], + description="Block prompt injection attempts", + model="perplexity/sonar", + ) + + assert int(exc.value.code) == 400 + assert exc.value.param == "model" + assert "tool calling" in exc.value.message + + def test_a_model_without_forced_tool_choice_support_remains_eligible(self, local_model_cost_map): + supported_params = litellm.get_supported_openai_params( + model="amazon.nova-pro-v1:0", + custom_llm_provider="bedrock", + ) + + assert supported_params is not None + assert "tools" in supported_params + assert "tool_choice" not in supported_params + + +class TestSuggesterToleratesAModelThatRefusesItsSamplingParams: + """The model is operator-supplied, so it can be a reasoning model whose only accepted + temperature is 1. This call pins temperature=0.2 for tool-selection determinism, which such + a model rejects outright: without drop_params litellm raises UnsupportedParamsError and the + whole suggestion fails rather than degrading. Every other internal LLM call in the proxy + already opts in through judge_acompletion; this one was the exception. + """ + + @pytest.mark.asyncio + async def test_a_reasoning_model_gets_past_param_mapping(self, monkeypatch, local_model_cost_map): + """Drives the real entry point with no patching and no network. Which exception escapes is + the discriminator: param mapping runs before any credential check, so UnsupportedParamsError + means the call died on the pinned temperature, while AuthenticationError means it survived + that and got as far as needing a key. Asserting the latter is what the caller observes. + """ + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + + with pytest.raises(litellm.AuthenticationError): + await AiPolicySuggester().suggest( + templates=SAMPLE_TEMPLATES, + attack_examples=["My SSN is 123-45-6789"], + description="", + model="gpt-5.6-terra", + ) + + def test_the_pinned_temperature_is_what_such_a_model_refuses(self, local_model_cost_map): + """The other half of the discriminator above: the same temperature this call pins is + exactly what the model rejects, and drop_params is what removes it.""" + from litellm.utils import get_optional_params + + optional_params = get_optional_params( + model="gpt-5.6-terra", + custom_llm_provider="openai", + temperature=0.2, + tools=[SUGGEST_TOOL], + tool_choice={"type": "function", "function": {"name": "select_policy_templates"}}, + drop_params=True, + ) + + assert "temperature" not in optional_params + assert optional_params["tools"] == [SUGGEST_TOOL] + assert optional_params["tool_choice"] == { + "type": "function", + "function": {"name": "select_policy_templates"}, + }