diff --git a/litellm/proxy/management_endpoints/policy_endpoints/__init__.py b/litellm/proxy/management_endpoints/policy_endpoints/__init__.py new file mode 100644 index 00000000000..8486a5660ac --- /dev/null +++ b/litellm/proxy/management_endpoints/policy_endpoints/__init__.py @@ -0,0 +1,13 @@ +""" +Policy endpoints package. + +Re-exports everything from endpoints module so existing imports +like `from litellm.proxy.management_endpoints.policy_endpoints import router` +continue to work. Patch targets also resolve correctly since names +are imported directly into this namespace. +""" + +from litellm.proxy.management_endpoints.policy_endpoints.endpoints import * # noqa: F401, F403 +from litellm.proxy.management_endpoints.policy_endpoints.endpoints import ( + router, +) diff --git a/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py b/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py new file mode 100644 index 00000000000..460df93b543 --- /dev/null +++ b/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py @@ -0,0 +1,133 @@ +""" +AI Policy Suggester - uses LLM tool calling to suggest policy templates +based on user-provided attack examples and descriptions. +""" + +import json +from typing import List + +from litellm._logging import verbose_proxy_logger + +SUGGEST_TOOL = { + "type": "function", + "function": { + "name": "select_policy_templates", + "description": "Select one or more policy templates that best match the user's security requirements", + "parameters": { + "type": "object", + "properties": { + "selected_templates": { + "type": "array", + "items": { + "type": "object", + "properties": { + "template_id": { + "type": "string", + "description": "The ID of the selected template", + }, + "reason": { + "type": "string", + "description": "Brief reason why this template matches", + }, + }, + "required": ["template_id", "reason"], + }, + "description": "List of templates that match the user's requirements", + }, + "explanation": { + "type": "string", + "description": "Overall explanation of why these templates were suggested", + }, + }, + "required": ["selected_templates", "explanation"], + }, + }, +} + + +class AiPolicySuggester: + """Suggests policy templates using LLM tool calling.""" + + async def suggest( + self, + templates: list, + attack_examples: List[str], + description: str, + ) -> dict: + import litellm + + system_prompt = self._build_system_prompt(templates) + user_prompt = self._build_user_prompt(attack_examples, description) + + try: + response = await litellm.acompletion( + model="gpt-4o-mini", + messages=[ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": user_prompt}, + ], + tools=[SUGGEST_TOOL], + tool_choice={ + "type": "function", + "function": {"name": "select_policy_templates"}, + }, + temperature=0.2, + ) + + tool_calls = response.choices[0].message.tool_calls # type: ignore + if not tool_calls: + return { + "selected_templates": [], + "explanation": "No templates could be matched to your requirements.", + } + + result = json.loads(tool_calls[0].function.arguments) + + valid_ids = {t["id"] for t in templates} + result["selected_templates"] = [ + s + for s in result.get("selected_templates", []) + if s.get("template_id") in valid_ids + ] + + return result + except Exception as e: + verbose_proxy_logger.error("AI policy suggestion failed: %s", e) + raise + + def _build_system_prompt(self, templates: list) -> str: + template_descriptions = [] + for t in templates: + examples = t.get("example_sentences", []) + examples_str = ( + ", ".join(f'"{e}"' for e in examples) if examples else "none" + ) + entry = ( + f"- ID: {t['id']}\n" + f" Title: {t['title']}\n" + f" Description: {t['description']}\n" + f" Example attacks it protects against: {examples_str}" + ) + template_descriptions.append(entry) + + return ( + "You are a security policy advisor. The user will describe attacks or content " + "they want to block. Your job is to select the most relevant policy templates " + "from the available set. Use the select_policy_templates tool to return your " + "selections. Only select templates that are clearly relevant to what the user " + "wants to block.\n\n" + "Available templates:\n\n" + "\n\n".join(template_descriptions) + ) + + def _build_user_prompt( + self, attack_examples: List[str], description: str + ) -> str: + parts = [] + filtered_examples = [e for e in attack_examples if e.strip()] + if filtered_examples: + parts.append("Example attack prompts I want to block:") + for i, ex in enumerate(filtered_examples, 1): + parts.append(f" {i}. {ex}") + if description.strip(): + parts.append(f"\nDescription of what I want to block: {description}") + return "\n".join(parts) diff --git a/litellm/proxy/management_endpoints/policy_endpoints.py b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py similarity index 95% rename from litellm/proxy/management_endpoints/policy_endpoints.py rename to litellm/proxy/management_endpoints/policy_endpoints/endpoints.py index 8eec697462c..cd432b38eca 100644 --- a/litellm/proxy/management_endpoints/policy_endpoints.py +++ b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py @@ -640,3 +640,37 @@ def _build_competitor_guardrail_definitions( defn["litellm_params"]["blocked_words"] = blocked_words_map[guardrail_name] return enriched + + +class SuggestTemplatesRequest(BaseModel): + attack_examples: List[str] = Field(default_factory=list) + description: str = Field(default="") + + +@router.post( + "/policy/templates/suggest", + tags=["policy management"], + dependencies=[Depends(user_api_key_auth)], +) +@management_endpoint_wrapper +async def suggest_policy_templates( + data: SuggestTemplatesRequest, + request: Request, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> dict: + """ + Use AI to suggest policy templates based on attack examples and descriptions. + + Calls an LLM with tool calling to match user requirements to available templates. + """ + from litellm.proxy.management_endpoints.policy_endpoints.ai_policy_suggester import ( + AiPolicySuggester, + ) + + templates = _load_policy_templates_from_local_backup() + suggester = AiPolicySuggester() + return await suggester.suggest( + templates=templates, + attack_examples=data.attack_examples, + description=data.description, + )