Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py: 14%
177 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1# +-------------------------------------------------------------+
2#
3# Use DynamoAI Guardrails for your LLM calls
4# https://dynamo.ai
5#
6# +-------------------------------------------------------------+
8import os
9from collections.abc import AsyncGenerator, AsyncIterable
10from datetime import datetime
11from typing import Any, Final
13import httpx
15import litellm
16from litellm._logging import verbose_proxy_logger
17from litellm.caching.caching import DualCache
18from litellm.integrations.custom_guardrail import CustomGuardrail
19from litellm.llms.custom_httpx.http_handler import (
20 get_async_httpx_client,
21 httpxSpecialProvider,
22)
23from litellm.proxy._types import UserAPIKeyAuth
24from litellm.types.guardrails import GuardrailEventHooks
25from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
26from litellm.types.proxy.guardrails.guardrail_hooks.dynamoai import (
27 DynamoAIProcessedResult,
28 DynamoAIRequest,
29 DynamoAIResponse,
30)
31from litellm.types.utils import CallTypesLiteral, GuardrailStatus, ModelResponseStream
33GUARDRAIL_NAME: Final = "dynamoai"
36class DynamoAIGuardrails(CustomGuardrail):
37 """
38 DynamoAI Guardrails integration for LiteLLM.
40 Provides content moderation and policy enforcement using DynamoAI's guardrail API.
41 """
43 def __init__(
44 self,
45 guardrail_name: str = "litellm_test",
46 api_key: str | None = None,
47 api_base: str | None = None,
48 model_id: str = "",
49 policy_ids: list[str] = [],
50 **kwargs,
51 ):
52 self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
54 # Set API configuration
55 self.api_key = api_key or os.getenv("DYNAMOAI_API_KEY")
56 if not self.api_key:
57 raise ValueError(
58 "DynamoAI API key is required. Set DYNAMOAI_API_KEY environment variable or pass api_key parameter."
59 )
61 self.api_base = api_base or os.getenv("DYNAMOAI_API_BASE", "https://api.dynamo.ai")
62 self.api_url = f"{self.api_base}/v1/moderation/analyze/"
64 # Model ID for tracking/logging purposes
65 self.model_id = model_id or os.getenv("DYNAMOAI_MODEL_ID", "")
67 # Policy IDs - get from parameter, env var, or use empty list
68 env_policy_ids: Final = os.getenv("DYNAMOAI_POLICY_IDS", "")
69 self.policy_ids = policy_ids or (env_policy_ids.split(",") if env_policy_ids else [])
70 self.guardrail_name = guardrail_name
71 self.guardrail_provider = "dynamoai"
73 # store kwargs as optional_params
74 self.optional_params = kwargs
76 # Set supported event hooks
77 kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks()))
79 super().__init__(guardrail_name=guardrail_name, **kwargs)
81 verbose_proxy_logger.debug(
82 "DynamoAI Guardrail initialized with guardrail_name=%s, model_id=%s",
83 self.guardrail_name,
84 self.model_id,
85 )
87 async def _call_dynamoai_guardrails(
88 self,
89 messages: list[dict[str, Any]],
90 event_type: GuardrailEventHooks,
91 text_type: str = "input",
92 request_data: dict | None = None,
93 ) -> DynamoAIResponse:
94 """
95 Call DynamoAI Guardrails API to analyze messages for policy violations.
97 Args:
98 messages: List of messages to analyze
99 text_type: Type of text being analyzed ("input" or "output")
100 request_data: Optional request data for logging purposes
102 Returns:
103 DynamoAIResponse: Response from the DynamoAI Guardrails API
104 """
105 start_time: Final = datetime.now()
107 payload: Final[DynamoAIRequest] = {
108 "messages": messages,
109 }
111 # Add optional fields if provided
112 if self.policy_ids:
113 payload["policyIds"] = self.policy_ids
114 if self.model_id:
115 payload["modelId"] = self.model_id
117 headers: Final = {
118 "Content-Type": "application/json",
119 "Authorization": f"Bearer {self.api_key}",
120 }
122 verbose_proxy_logger.debug(
123 "DynamoAI request to %s with payload=%s",
124 self.api_url,
125 payload,
126 )
128 try:
129 response: Final = await self.async_handler.post(
130 url=self.api_url,
131 json=dict(payload),
132 headers=headers,
133 )
134 response.raise_for_status()
135 response_json: Final = response.json()
137 end_time = datetime.now()
138 duration = (end_time - start_time).total_seconds()
140 # Add guardrail information to request trace
141 if request_data:
142 guardrail_status: Final = self._determine_guardrail_status(response_json)
143 self.add_standard_logging_guardrail_information_to_request_data(
144 guardrail_provider=self.guardrail_provider,
145 guardrail_json_response=response_json,
146 request_data=request_data,
147 guardrail_status=guardrail_status,
148 start_time=start_time.timestamp(),
149 end_time=end_time.timestamp(),
150 duration=duration,
151 event_type=event_type,
152 )
154 return response_json
156 except httpx.HTTPError as e:
157 end_time = datetime.now()
158 duration = (end_time - start_time).total_seconds()
160 verbose_proxy_logger.error("DynamoAI API request failed: %s", str(e))
162 # Add guardrail information with failure status
163 if request_data:
164 self.add_standard_logging_guardrail_information_to_request_data(
165 guardrail_provider=self.guardrail_provider,
166 guardrail_json_response={"error": str(e)},
167 request_data=request_data,
168 guardrail_status="guardrail_failed_to_respond",
169 start_time=start_time.timestamp(),
170 end_time=end_time.timestamp(),
171 duration=duration,
172 event_type=event_type,
173 )
175 raise
177 def _process_dynamoai_guardrails_response(self, response: DynamoAIResponse) -> DynamoAIProcessedResult:
178 """
179 Process the response from the DynamoAI Guardrails API
181 Args:
182 response: The response from the API with 'finalAction' and 'appliedPolicies' keys
184 Returns:
185 DynamoAIProcessedResult: Processed response with detected violations
186 """
187 final_action: Final = response.get("finalAction", "NONE")
188 applied_policies: Final = response.get("appliedPolicies", [])
190 violations_detected: Final[list[str]] = []
191 violation_details: Final[dict[str, object]] = {}
193 # For now, only handle BLOCK action
194 if final_action == "BLOCK":
195 for applied_policy in applied_policies:
196 policy_info = applied_policy.get("policy", {})
197 policy_outputs = applied_policy.get("outputs", {})
199 # Get policy name and action
200 policy_name = policy_info.get("name", "unknown")
202 # Check for action in multiple places
203 policy_action = (
204 applied_policy.get("action") or (policy_outputs.get("action") if policy_outputs else None) or "NONE"
205 )
207 # Only include policies with BLOCK action
208 if policy_action == "BLOCK":
209 violations_detected.append(policy_name)
210 violation_details[policy_name] = {
211 "policyId": policy_info.get("id"),
212 "action": policy_action,
213 "method": policy_info.get("method"),
214 "description": policy_info.get("description"),
215 "message": (policy_outputs.get("message") if policy_outputs else None),
216 }
218 return {
219 "violations_detected": violations_detected,
220 "violation_details": violation_details,
221 }
223 def _determine_guardrail_status(self, response_json: DynamoAIResponse) -> GuardrailStatus:
224 """
225 Determine the guardrail status based on DynamoAI API response.
227 Returns:
228 "success": Content allowed through with no violations (finalAction is NONE)
229 "guardrail_intervened": Content blocked (finalAction is BLOCK)
230 "guardrail_failed_to_respond": Technical error or API failure
231 """
232 try:
233 if not isinstance(response_json, dict):
234 return "guardrail_failed_to_respond"
236 # Check for error in response
237 if response_json.get("error"):
238 return "guardrail_failed_to_respond"
240 final_action: Final = response_json.get("finalAction", "NONE")
242 if final_action == "NONE":
243 return "success"
244 elif final_action == "BLOCK":
245 return "guardrail_intervened"
247 # For now, treat other actions as success (WARN, REDACT, SANITIZE not implemented yet)
248 return "success"
250 except Exception as e:
251 verbose_proxy_logger.error("Error determining DynamoAI guardrail status: %s", str(e))
252 return "guardrail_failed_to_respond"
254 def _create_error_message(self, processed_result: DynamoAIProcessedResult) -> str:
255 """
256 Create a detailed error message from processed guardrail results.
258 Args:
259 processed_result: Processed response with detected violations
261 Returns:
262 Formatted error message string
263 """
264 violations_detected: Final = processed_result["violations_detected"]
265 violation_details: Final = processed_result["violation_details"]
267 error_message = f"Guardrail failed: {len(violations_detected)} violation(s) detected\n\n"
269 for policy_name in violations_detected:
270 error_message += f"- {policy_name.upper()}:\n"
271 details = violation_details.get(policy_name, {})
273 # Format violation details
274 if details.get("action"):
275 error_message += f" Action: {details['action']}\n"
276 if details.get("method"):
277 error_message += f" Method: {details['method']}\n"
278 if details.get("description"):
279 error_message += f" Description: {details['description']}\n"
280 if details.get("message"):
281 error_message += f" Message: {details['message']}\n"
282 if details.get("policyId"):
283 error_message += f" Policy ID: {details['policyId']}\n"
284 error_message += "\n"
286 return error_message.strip()
288 async def async_pre_call_hook(
289 self,
290 user_api_key_dict: UserAPIKeyAuth,
291 cache: DualCache,
292 data: dict,
293 call_type: CallTypesLiteral,
294 ) -> Exception | str | dict | None:
295 """
296 Runs before the LLM API call
297 Runs on only Input
298 Use this if you want to MODIFY the input
299 """
300 verbose_proxy_logger.debug("Running DynamoAI pre-call hook")
302 from litellm.proxy.common_utils.callback_utils import (
303 add_guardrail_to_applied_guardrails_header,
304 )
306 event_type: Final[GuardrailEventHooks] = GuardrailEventHooks.pre_call
307 if self.should_run_guardrail(data=data, event_type=event_type) is not True:
308 return data
310 _messages: Final = data.get("messages")
311 if _messages:
312 result: Final = await self._call_dynamoai_guardrails(
313 messages=_messages,
314 text_type="input",
315 request_data=data,
316 event_type=GuardrailEventHooks.pre_call,
317 )
319 verbose_proxy_logger.debug("Guardrails async_pre_call_hook result=%s", result)
321 # Process the guardrails response
322 processed_result: Final = self._process_dynamoai_guardrails_response(result)
323 violations_detected: Final = processed_result["violations_detected"]
325 # If any violations are detected, raise an error
326 if violations_detected:
327 error_message: Final = self._create_error_message(processed_result)
328 raise ValueError(error_message)
330 # Add guardrail to applied guardrails header
331 add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
333 return data
335 async def async_moderation_hook(
336 self,
337 data: dict,
338 user_api_key_dict: UserAPIKeyAuth,
339 call_type: CallTypesLiteral,
340 ):
341 """
342 Runs in parallel to LLM API call
343 Runs on only Input
345 This can NOT modify the input, only used to reject or accept a call before going to LLM API
346 """
347 from litellm.proxy.common_utils.callback_utils import (
348 add_guardrail_to_applied_guardrails_header,
349 )
351 event_type: Final[GuardrailEventHooks] = GuardrailEventHooks.during_call
352 if self.should_run_guardrail(data=data, event_type=event_type) is not True:
353 return
355 _messages: Final = data.get("messages")
356 if _messages:
357 result: Final = await self._call_dynamoai_guardrails(
358 messages=_messages,
359 text_type="input",
360 request_data=data,
361 event_type=GuardrailEventHooks.during_call,
362 )
364 verbose_proxy_logger.debug("Guardrails async_moderation_hook result=%s", result)
366 # Process the guardrails response
367 processed_result: Final = self._process_dynamoai_guardrails_response(result)
368 violations_detected: Final = processed_result["violations_detected"]
370 # If any violations are detected, raise an error
371 if violations_detected:
372 error_message: Final = self._create_error_message(processed_result)
373 raise ValueError(error_message)
375 # Add guardrail to applied guardrails header
376 add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
378 return data
380 async def async_post_call_success_hook(
381 self,
382 data: dict,
383 user_api_key_dict: UserAPIKeyAuth,
384 response,
385 ):
386 """
387 Runs on response from LLM API call
389 It can be used to reject a response
391 Uses DynamoAI guardrails to check the response for policy violations
392 """
393 from litellm.proxy.common_utils.callback_utils import (
394 add_guardrail_to_applied_guardrails_header,
395 )
396 from litellm.types.guardrails import GuardrailEventHooks
398 if self.should_run_guardrail(data=data, event_type=GuardrailEventHooks.post_call) is not True:
399 return
401 verbose_proxy_logger.debug("async_post_call_success_hook response=%s", response)
403 # Check if the ModelResponse has text content in its choices
404 # to avoid sending empty content to DynamoAI (e.g., during tool calls)
405 if isinstance(response, litellm.ModelResponse):
406 has_text_content = False
407 dynamoai_messages: Final[list[dict[str, str]]] = []
409 for choice in response.choices:
410 if isinstance(choice, litellm.Choices):
411 if choice.message.content and isinstance(choice.message.content, str):
412 has_text_content = True
413 dynamoai_messages.append(
414 {
415 "role": choice.message.role or "assistant",
416 "content": choice.message.content,
417 }
418 )
420 if not has_text_content:
421 verbose_proxy_logger.warning("DynamoAI: not running guardrail. No output text in response")
422 return
424 if dynamoai_messages:
425 result: Final = await self._call_dynamoai_guardrails(
426 messages=dynamoai_messages,
427 text_type="output",
428 request_data=data,
429 event_type=GuardrailEventHooks.post_call,
430 )
432 verbose_proxy_logger.debug("Guardrails async_post_call_success_hook result=%s", result)
434 # Process the guardrails response
435 processed_result: Final = self._process_dynamoai_guardrails_response(result)
436 violations_detected: Final = processed_result["violations_detected"]
438 # If any violations are detected, raise an error
439 if violations_detected:
440 error_message: Final = self._create_error_message(processed_result)
441 raise ValueError(error_message)
443 # Add guardrail to applied guardrails header
444 add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
446 async def async_post_call_streaming_iterator_hook(
447 self,
448 user_api_key_dict: UserAPIKeyAuth,
449 response: AsyncIterable[ModelResponseStream],
450 request_data: dict,
451 ) -> AsyncGenerator[ModelResponseStream, None]:
452 """
453 Passes the entire stream to the guardrail
455 This is useful for guardrails that need to see the entire response, such as PII masking.
457 Triggered by mode: 'post_call'
458 """
459 async for item in response:
460 yield item
462 @staticmethod
463 def get_config_model() -> type[GuardrailConfigModel] | None:
464 from litellm.types.proxy.guardrails.guardrail_hooks.dynamoai import (
465 DynamoAIGuardrailConfigModel,
466 )
468 return DynamoAIGuardrailConfigModel
470 @classmethod
471 def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
472 return [
473 GuardrailEventHooks.pre_call,
474 GuardrailEventHooks.post_call,
475 GuardrailEventHooks.during_call,
476 ]