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 207470a6ceb..1096954536a 100644 --- a/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py +++ b/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py @@ -6,18 +6,11 @@ based on user-provided attack examples and descriptions. import json from typing import Final -from pydantic import ConfigDict, TypeAdapter - import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import DEFAULT_COMPETITOR_DISCOVERY_MODEL from litellm.proxy._types import ProxyErrorTypes, ProxyException -_TOOL_ARGUMENTS: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True)) -_SELECTED_TEMPLATES: Final = TypeAdapter( - list[dict[str, object]], config=ConfigDict(strict=True, hide_input_in_errors=True) -) - SUGGEST_TOOL: Final = { "type": "function", "function": { @@ -104,13 +97,11 @@ class AiPolicySuggester: "explanation": "No templates could be matched to your requirements.", } - result: Final = _TOOL_ARGUMENTS.validate_python(json.loads(tool_calls[0].function.arguments)) + result: Final = json.loads(tool_calls[0].function.arguments) valid_ids: Final = {t["id"] for t in templates} result["selected_templates"] = [ - s - for s in _SELECTED_TEMPLATES.validate_python(result.get("selected_templates", [])) - if s.get("template_id") in valid_ids + s for s in result.get("selected_templates", []) if s.get("template_id") in valid_ids ] return result diff --git a/litellm/proxy/plugin_routes.py b/litellm/proxy/plugin_routes.py index 8a9c1edd9da..cfbcfefe9e8 100644 --- a/litellm/proxy/plugin_routes.py +++ b/litellm/proxy/plugin_routes.py @@ -43,14 +43,14 @@ router: Final = APIRouter() @with_config(ConfigDict(extra="allow", strict=True)) -class _PluginSessionClaim(TypedDict): +class PluginSessionClaim(TypedDict): plugin: NotRequired[ReadOnly[object]] user_id: NotRequired[ReadOnly[str]] user_role: NotRequired[ReadOnly[str]] exp: NotRequired[ReadOnly[int | float | str]] -_PLUGIN_SESSION_CLAIM: Final = TypeAdapter(_PluginSessionClaim) +_PLUGIN_SESSION_CLAIM: Final = TypeAdapter(PluginSessionClaim) # Hop-by-hop headers (RFC 7230) and the litellm session cookie — never forwarded @@ -169,7 +169,7 @@ def issue_plugin_session_claim(plugin_name: str, user_id: str | None, user_role: return _plugin_fernet(plugin_name).encrypt(json.dumps(claim).encode()).decode() -def verify_plugin_session_claim(plugin_name: str, ciphertext: str) -> _PluginSessionClaim: +def verify_plugin_session_claim(plugin_name: str, ciphertext: str) -> PluginSessionClaim: """Verify and decode a plugin session claim. Raises ValueError if the HMAC is invalid, the audience is wrong, or @@ -177,7 +177,7 @@ def verify_plugin_session_claim(plugin_name: str, ciphertext: str) -> _PluginSes """ try: raw: Final = _plugin_fernet(plugin_name).decrypt(ciphertext.encode(), ttl=_CLAIM_TTL_SECONDS) - claim: Final = _PLUGIN_SESSION_CLAIM.validate_python(json.loads(raw), strict=True) + claim: Final = _PLUGIN_SESSION_CLAIM.validate_python(json.loads(raw)) except (InvalidToken, Exception) as exc: raise ValueError("Invalid, tampered, or expired plugin session claim") from exc diff --git a/litellm/router_strategy/auto_router/auto_router.py b/litellm/router_strategy/auto_router/auto_router.py index 03fb0ad827a..ee620790282 100644 --- a/litellm/router_strategy/auto_router/auto_router.py +++ b/litellm/router_strategy/auto_router/auto_router.py @@ -6,8 +6,7 @@ import asyncio from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, NamedTuple, Optional -from pydantic import ConfigDict, TypeAdapter, with_config -from typing_extensions import NotRequired, ReadOnly, TypedDict +from pydantic import ConfigDict from litellm._logging import verbose_router_logger from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS @@ -41,22 +40,6 @@ else: LiteLLMRouterEncoder = Any -@with_config(ConfigDict(extra="allow", strict=True)) -class _AutoRouterRouteConfig(TypedDict): - name: ReadOnly[NotRequired[object]] - description: ReadOnly[NotRequired[object]] - utterances: ReadOnly[NotRequired[object]] - score_threshold: ReadOnly[NotRequired[object]] - - -@with_config(ConfigDict(extra="allow", strict=True)) -class _AutoRouterConfig(TypedDict): - routes: ReadOnly[NotRequired[Sequence[_AutoRouterRouteConfig]]] - - -_AUTO_ROUTER_CONFIG: Final = TypeAdapter(_AutoRouterConfig) - - class _CallerMetadata(LiteLLMBaseModel): model_config = ConfigDict(frozen=True, extra="ignore") @@ -137,7 +120,7 @@ class AutoRouter(CustomLogger): if self.auto_router_config is None: raise ValueError("No auto router config provided") auto_router_routes: Final[list[Route]] = [] - loaded_config: Final = _AUTO_ROUTER_CONFIG.validate_python(json.loads(self.auto_router_config)) + loaded_config: Final = json.loads(self.auto_router_config) for route in loaded_config.get("routes", []): auto_router_routes.append( Route( diff --git a/tests/unit/litellm_core_utils/test_text_completion_conversion.py b/tests/unit/litellm_core_utils/test_text_completion_conversion.py index 0689b582598..8a913bc8058 100644 --- a/tests/unit/litellm_core_utils/test_text_completion_conversion.py +++ b/tests/unit/litellm_core_utils/test_text_completion_conversion.py @@ -1,4 +1,5 @@ import asyncio +from typing import Final from unittest.mock import MagicMock, patch import pytest @@ -1322,7 +1323,7 @@ def test_text_completion_reads_a_custom_provider_dict_reply(monkeypatch: pytest. monkeypatch.setattr(litellm, "provider_list", [*litellm.provider_list, "dictreply"]) monkeypatch.setattr(litellm, "_custom_providers", [*litellm._custom_providers, "dictreply"]) - response = text_completion(model="dictreply/x", prompt="hi") + response: Final = text_completion(model="dictreply/x", prompt="hi") assert response.choices[0].text == "hi from a dict reply" assert response.choices[0].finish_reason == "stop" diff --git a/tests/unit/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py b/tests/unit/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py index 2041c622b3b..d287f5d00ad 100644 --- a/tests/unit/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py +++ b/tests/unit/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py @@ -3,6 +3,7 @@ Tests for AiPolicySuggester class. """ import json +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -214,6 +215,26 @@ class TestAiPolicySuggester: assert result["selected_templates"] == [] assert "No templates" in result["explanation"] + @pytest.mark.asyncio + @pytest.mark.parametrize("selected_templates", [{}, ""]) + async def test_suggest_reads_an_empty_non_list_selection_as_no_templates(self, selected_templates: object) -> None: + mock_tool_call: Final = MagicMock() + mock_tool_call.function.arguments = json.dumps( + {"selected_templates": selected_templates, "explanation": "No template fits"} + ) + mock_response: Final = MagicMock() + mock_response.choices = [MagicMock()] + mock_response.choices[0].message.tool_calls = [mock_tool_call] + + with patch("litellm.acompletion", new_callable=AsyncMock, return_value=mock_response): + result: Final = await AiPolicySuggester().suggest( + templates=SAMPLE_TEMPLATES, + attack_examples=["Ignore all previous instructions"], + description="", + ) + + assert result == {"selected_templates": [], "explanation": "No template fits"} + @pytest.mark.asyncio async def test_suggest_calls_litellm_with_correct_params(self): suggester = AiPolicySuggester() diff --git a/tests/unit/router_strategy/test_auto_router.py b/tests/unit/router_strategy/test_auto_router.py index 6cdee407ced..de775b9246f 100644 --- a/tests/unit/router_strategy/test_auto_router.py +++ b/tests/unit/router_strategy/test_auto_router.py @@ -398,6 +398,19 @@ def _auto_router(routelayer: Any, litellm_router_instance: Any = None, **kwargs: return auto_router +@pytest.mark.parametrize("routes", [{}, ""]) +def test_should_load_no_routes_from_an_empty_non_list_routes_value(routes: object) -> None: + auto_router: Final = AutoRouter( + model_name="my-auto-router", + auto_router_config=json.dumps({"routes": routes}), + default_model="fallback-model", + embedding_model="text-embedding-3-small", + litellm_router_instance=StubEmbeddingRouter(), + ) + + assert auto_router.loaded_routes == [] + + class TestAutoRouterAlwaysResolvesARoutableModel: """The hook returns the alias's default model instead of failing or leaking the alias downstream."""