Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py: 81%

1""" 

2AI Policy Suggester - uses LLM tool calling to suggest policy templates 

3based on user-provided attack examples and descriptions. 

4""" 

5 

6import json 

7from typing import Final 

8 

9import litellm 

10from litellm._logging import verbose_proxy_logger 

11from litellm.constants import DEFAULT_COMPETITOR_DISCOVERY_MODEL 

12from litellm.proxy._types import ProxyErrorTypes, ProxyException 

13 

14SUGGEST_TOOL: Final = { 

15 "type": "function", 

16 "function": { 

17 "name": "select_policy_templates", 

18 "description": "Select one or more policy templates that best match the user's security requirements", 

19 "parameters": { 

20 "type": "object", 

21 "properties": { 

22 "selected_templates": { 

23 "type": "array", 

24 "items": { 

25 "type": "object", 

26 "properties": { 

27 "template_id": { 

28 "type": "string", 

29 "description": "The ID of the selected template", 

30 }, 

31 "reason": { 

32 "type": "string", 

33 "description": "Brief reason why this template matches", 

34 }, 

35 }, 

36 "required": ["template_id", "reason"], 

37 }, 

38 "description": "List of templates that match the user's requirements", 

39 }, 

40 "explanation": { 

41 "type": "string", 

42 "description": "Overall explanation of why these templates were suggested", 

43 }, 

44 }, 

45 "required": ["selected_templates", "explanation"], 

46 }, 

47 }, 

48} 

49 

50 

51class AiPolicySuggester: 

52 """Suggests policy templates using LLM tool calling.""" 

53 

54 async def suggest( 

55 self, 

56 templates: list, 

57 attack_examples: list[str], 

58 description: str, 

59 model: str | None = None, 

60 ) -> dict: 

61 system_prompt: Final = self._build_system_prompt(templates) 

62 user_prompt: Final = self._build_user_prompt(attack_examples, description) 

63 model = model or DEFAULT_COMPETITOR_DISCOVERY_MODEL 

64 custom_llm_provider: Final = model.split("/", 1)[0] if "/" in model else None 

65 supported_params: Final = litellm.get_supported_openai_params( 

66 model=model, 

67 custom_llm_provider=custom_llm_provider, 

68 ) 

69 if supported_params is not None and "tools" not in supported_params: 69 ↛ 70line 69 didn't jump to line 70 because the condition on line 69 was never true

70 raise ProxyException( 

71 message=(f"AI policy suggestion requires tool calling; model '{model}' does not support it"), 

72 type=ProxyErrorTypes.validation_error.value, 

73 param="model", 

74 code=400, 

75 ) 

76 

77 try: 

78 response: Final = await litellm.acompletion( 

79 model=model, 

80 messages=[ 

81 {"role": "system", "content": system_prompt}, 

82 {"role": "user", "content": user_prompt}, 

83 ], 

84 tools=[SUGGEST_TOOL], 

85 tool_choice={ 

86 "type": "function", 

87 "function": {"name": "select_policy_templates"}, 

88 }, 

89 temperature=0.2, 

90 drop_params=True, 

91 ) 

92 

93 tool_calls: Final = response.choices[0].message.tool_calls 

94 if not tool_calls: 

95 return { 

96 "selected_templates": [], 

97 "explanation": "No templates could be matched to your requirements.", 

98 } 

99 

100 result: Final = json.loads(tool_calls[0].function.arguments) 

101 

102 valid_ids: Final = {t["id"] for t in templates} 

103 result["selected_templates"] = [ 

104 s for s in result.get("selected_templates", []) if s.get("template_id") in valid_ids 

105 ] 

106 

107 return result 

108 except Exception as e: 

109 verbose_proxy_logger.error("AI policy suggestion failed: %s", e) 

110 raise 

111 

112 def _build_system_prompt(self, templates: list) -> str: 

113 template_descriptions: Final = [] 

114 for t in templates: 

115 examples = t.get("example_sentences", []) 

116 examples_str = ", ".join(f'"{e}"' for e in examples) if examples else "none" 

117 entry = ( 

118 f"- ID: {t['id']}\n" 

119 f" Title: {t['title']}\n" 

120 f" Description: {t['description']}\n" 

121 f" Example attacks it protects against: {examples_str}" 

122 ) 

123 template_descriptions.append(entry) 

124 

125 return ( 

126 "You are a security policy advisor. The user will describe attacks or content " 

127 "they want to block. Your job is to select the most relevant policy templates " 

128 "from the available set. Use the select_policy_templates tool to return your " 

129 "selections. Only select templates that are clearly relevant to what the user " 

130 "wants to block.\n\n" 

131 "Available templates:\n\n" + "\n\n".join(template_descriptions) 

132 ) 

133 

134 def _build_user_prompt(self, attack_examples: list[str], description: str) -> str: 

135 parts: Final = [] 

136 filtered_examples: Final = [e for e in attack_examples if e.strip()] 

137 if filtered_examples: 

138 parts.append("Example attack prompts I want to block:") 

139 for i, ex in enumerate(filtered_examples, 1): 

140 parts.append(f" {i}. {ex}") 

141 if description.strip(): 

142 parts.append(f"\nDescription of what I want to block: {description}") 

143 return "\n".join(parts)