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"}, + }