Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py: 13%
566 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1import time
2from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence
3from enum import Enum, auto
4from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal
6import httpx
7from fastapi import HTTPException
9if TYPE_CHECKING: 9 ↛ 10line 9 didn't jump to line 10 because the condition on line 9 was never true
10 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
11 from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
13import json
15import litellm
16from litellm._logging import verbose_proxy_logger
17from litellm.caching import DualCache
18from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
19from litellm.integrations.custom_guardrail import (
20 CustomGuardrail,
21 log_guardrail_information,
22)
23from litellm.litellm_core_utils.core_helpers import (
24 get_metadata_variable_name_from_kwargs,
25 get_or_create_metadata_bucket,
26)
27from litellm.llms.custom_httpx.http_handler import (
28 AsyncHTTPHandler,
29 get_async_httpx_client,
30 httpxSpecialProvider,
31)
32from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
33from litellm.proxy._types import UserAPIKeyAuth
34from litellm.proxy.guardrails.anthropic_sse import (
35 anthropic_sse_chunks_from_response,
36 anthropic_sse_error_frames,
37 assemble_anthropic_sse_stream,
38 is_anthropic_sse_stream,
39 is_raw_sse_stream,
40 is_sse_error_stream,
41)
42from litellm.proxy.guardrails.guardrail_hooks.model_armor.file_scanning import (
43 MODEL_ARMOR_MAX_FILE_SIZE_BYTES,
44 plan_file_scans,
45)
46from litellm.types.guardrails import GuardrailEventHooks, LitellmParams
47from litellm.types.llms.openai import (
48 AllMessageValues,
49 ChatCompletionToolCallChunk,
50 ResponsesAPIResponse,
51 ResponsesAPIStreamEvents,
52)
53from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES
54from litellm.types.utils import (
55 CallTypes,
56 CallTypesLiteral,
57 Choices,
58 GenericGuardrailAPIInputs,
59 GuardrailStatus,
60 ModelResponse,
61 ModelResponseStream,
62 StandardLoggingGuardrailInformation,
63 TextCompletionResponse,
64)
66GUARDRAIL_NAME: Final = "model_armor"
68# Only these carry the finished output; response.created carries an empty body
69_RESPONSES_TERMINAL_EVENT_TYPES: Final = frozenset({"response.completed", "response.incomplete", "response.failed"})
71# Every event whose ``delta`` is model output already on its way to the client. Read off the event
72# enum rather than listed, so an event added there cannot quietly fall out of the scan
73_RESPONSES_DELTA_EVENT_TYPES: Final = frozenset(
74 event.value for event in ResponsesAPIStreamEvents if event.value.endswith(".delta")
75)
77# What makes two delta events part of the same field of the turn, rather than two fields that merely
78# streamed next to each other
79_RESPONSES_DELTA_FIELD_ATTRS: Final = ("type", "item_id", "output_index", "content_index", "summary_index")
82class _StreamSurface(Enum):
83 """Wire format of a buffered streaming response, which decides how it is read and how it is refused."""
85 CHAT_COMPLETIONS = auto()
86 ANTHROPIC_MESSAGES = auto()
87 RESPONSES = auto()
88 OPAQUE_SSE = auto()
91class ModelArmorAPIError(Exception):
92 """Model Armor API failure (non-2xx), distinct from a content-block decision so
93 hooks can honor fail_on_error. The detail is already sanitized per configuration."""
95 def __init__(self, detail: str):
96 super().__init__(detail)
97 self.detail = detail
100_SCANNED_CONTENT_KEYS: Final = frozenset({"text", "sanitizedText", "findings", "maliciousUriMatchedItems"})
102RedactablePayload = dict | list | str | int | float | bool | None
105def _redact_scanned_content(payload: RedactablePayload, depth: int = 0) -> RedactablePayload:
106 if depth >= DEFAULT_MAX_RECURSE_DEPTH:
107 return "[REDACTED]"
108 if isinstance(payload, dict):
109 return {
110 key: "[REDACTED]" if key in _SCANNED_CONTENT_KEYS else _redact_scanned_content(value, depth + 1)
111 for key, value in payload.items()
112 }
113 if isinstance(payload, list):
114 return [_redact_scanned_content(item, depth + 1) for item in payload]
115 return payload
118class ModelArmorGuardrail(CustomGuardrail, VertexBase):
119 """
120 Google Cloud Model Armor Guardrail integration for LiteLLM.
122 Supports:
123 - Pre-call sanitization (sanitizeUserPrompt)
124 - Post-call sanitization (sanitizeModelResponse)
125 - logging_only: scans the completed response after it reaches the client and
126 records the verdict in spend logs without blocking
127 """
129 use_native_lifecycle_hooks: ClassVar[bool] = True
131 @classmethod
132 def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
133 return [
134 GuardrailEventHooks.pre_call,
135 GuardrailEventHooks.during_call,
136 GuardrailEventHooks.post_call,
137 GuardrailEventHooks.pre_mcp_call,
138 GuardrailEventHooks.during_mcp_call,
139 GuardrailEventHooks.logging_only,
140 ]
142 def __init__(
143 self,
144 template_id: str | None = None,
145 project_id: str | None = None,
146 location: str | None = None,
147 credentials: VERTEX_CREDENTIALS_TYPES | None = None,
148 api_endpoint: str | None = None,
149 sanitize_error_detail: "bool | None" = True,
150 async_handler: AsyncHTTPHandler | None = None,
151 access_token_provider: Callable[[], Awaitable[tuple[str, str]]] | None = None,
152 **kwargs,
153 ):
154 # Set supported event hooks if not already provided
155 if "event_hook" not in kwargs:
156 kwargs["event_hook"] = [
157 GuardrailEventHooks.pre_call,
158 GuardrailEventHooks.during_call,
159 GuardrailEventHooks.post_call,
160 ]
161 kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks()))
163 # Initialize parent classes first
164 super().__init__(**kwargs)
165 VertexBase.__init__(self)
167 # Then set our attributes (this ensures project_id is not overwritten)
168 self.async_handler = async_handler or get_async_httpx_client(
169 llm_provider=httpxSpecialProvider.GuardrailCallback
170 )
171 self.access_token_provider = access_token_provider
172 self.template_id = template_id
173 self.project_id = project_id
174 self.location = location or "us-central1"
175 self.credentials = credentials
176 self.api_endpoint = api_endpoint
177 self.sanitize_error_detail = sanitize_error_detail is not False
179 # Store optional params
180 self.optional_params = kwargs
182 verbose_proxy_logger.debug(
183 "Model Armor Guardrail initialized with template_id: %s, project_id: %s, location: %s",
184 self.template_id,
185 self.project_id,
186 self.location,
187 )
189 def _get_api_endpoint(self) -> str:
190 """Get the API endpoint for Model Armor."""
191 if self.api_endpoint:
192 return self.api_endpoint
193 return f"https://modelarmor.{self.location}.rep.googleapis.com"
195 def _create_sanitize_request(self, content: str, source: Literal["user_prompt", "model_response"]) -> dict:
196 """Create request body for Model Armor API with correct camelCase field names."""
197 if source == "user_prompt":
198 return {"userPromptData": {"text": content}}
199 else:
200 return {"modelResponseData": {"text": content}}
202 def _extract_content_from_response(self, response: object) -> str:
203 """
204 Extract text content from model response.
206 Returns empty string for non-text responses (TTS, images, etc.) to skip guardrail processing.
207 """
208 from litellm.litellm_core_utils.prompt_templates.common_utils import (
209 get_content_from_model_response,
210 )
212 # Handle ModelResponse objects
213 if isinstance(response, litellm.ModelResponse):
214 return get_content_from_model_response(response)
216 # For non-ModelResponse types (e.g., TTS, images), return empty string
217 # These response types are not text-based and shouldn't be processed by text guardrails
218 verbose_proxy_logger.debug("Model Armor: Skipping non-ModelResponse type: %s", type(response).__name__)
219 return ""
221 def _build_api_error_detail(self, status_code: int, response_text: str) -> str:
222 if self.sanitize_error_detail:
223 return f"Model Armor API error (upstream {status_code})"
224 return f"Model Armor API error (upstream {status_code}): {response_text}"
226 def _build_block_error_detail(self, message: str, armor_response: RedactablePayload) -> dict:
227 if self.sanitize_error_detail:
228 return {"error": message}
229 return {"error": message, "model_armor_response": armor_response}
231 def _build_logging_response(self, armor_response: RedactablePayload) -> RedactablePayload:
232 if self.sanitize_error_detail:
233 return _redact_scanned_content(armor_response)
234 return armor_response
236 def _raise_if_fail_closed(self, e: ModelArmorAPIError) -> None:
237 if self.optional_params.get("fail_on_error", True):
238 raise e from None
240 def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None:
241 super().update_in_memory_litellm_params(litellm_params)
242 self.sanitize_error_detail = self.sanitize_error_detail is not False
244 def _log_request_debug(
245 self,
246 url: str,
247 body: dict,
248 file_bytes: "bytes | None",
249 file_type: "str | None",
250 ) -> None:
251 # Never log byteData: it is the full base64 of the scanned document. Log only its
252 # type and size so debug deployments cannot leak the contents the guardrail inspects.
253 if file_bytes is not None and file_type is not None:
254 verbose_proxy_logger.debug(
255 "Model Armor file request - URL: %s, byteDataType: %s, bytes: %d",
256 url,
257 file_type,
258 len(file_bytes),
259 )
260 elif self.sanitize_error_detail:
261 verbose_proxy_logger.debug("Model Armor request - URL: %s", url)
262 else:
263 verbose_proxy_logger.debug(
264 "Model Armor request - URL: %s, Body: %s",
265 url,
266 body,
267 )
269 def _log_response_debug(self, status_code: int, response_text: str) -> None:
270 if self.sanitize_error_detail:
271 verbose_proxy_logger.debug(
272 "Model Armor response - Status: %s",
273 status_code,
274 )
275 else:
276 verbose_proxy_logger.debug(
277 "Model Armor response - Status: %s, Body: %s",
278 status_code,
279 response_text,
280 )
282 async def make_model_armor_request(
283 self,
284 content: str | None = None,
285 source: Literal["user_prompt", "model_response"] = "user_prompt",
286 request_data: dict | None = None,
287 file_bytes: bytes | None = None,
288 file_type: str | None = None,
289 ) -> dict:
290 """
291 Make request to Model Armor API. Supports both text and file prompt sanitization.
292 If file_bytes and file_type are provided, file prompt sanitization is performed.
293 """
294 # Get access token using VertexBase auth
295 if self.access_token_provider is not None:
296 access_token, resolved_project_id = await self.access_token_provider()
297 else:
298 access_token, resolved_project_id = await self._ensure_access_token_async(
299 credentials=self.credentials,
300 project_id=self.project_id,
301 custom_llm_provider="vertex_ai",
302 )
304 # Use resolved project ID if not explicitly set
305 if not self.project_id and resolved_project_id:
306 self.project_id = resolved_project_id
308 # Construct URL
309 endpoint: Final = self._get_api_endpoint()
310 if source == "user_prompt":
311 url = f"{endpoint}/v1/projects/{self.project_id}/locations/{self.location}/templates/{self.template_id}:sanitizeUserPrompt"
312 else:
313 url = f"{endpoint}/v1/projects/{self.project_id}/locations/{self.location}/templates/{self.template_id}:sanitizeModelResponse"
315 # Create request body
316 if file_bytes is not None and file_type is not None:
317 body = self.sanitize_file_prompt(file_bytes, file_type, source)
318 elif content is not None:
319 body = self._create_sanitize_request(content, source)
320 else:
321 raise ValueError("Either content or file_bytes and file_type must be provided.")
323 # Set headers
324 headers: Final = {
325 "Content-Type": "application/json",
326 "Authorization": f"Bearer {access_token}",
327 }
329 self._log_request_debug(url=url, body=body, file_bytes=file_bytes, file_type=file_type)
331 # Make request
332 if self.async_handler is None:
333 raise ValueError("Async handler not initialized")
335 try:
336 response: Final = await self.async_handler.post(
337 url=url,
338 json=body,
339 headers=headers,
340 )
341 except httpx.HTTPStatusError as e:
342 detail = self._build_api_error_detail(e.response.status_code, e.response.text)
343 verbose_proxy_logger.error(
344 "Model Armor API error - Status: %s, Detail: %s",
345 e.response.status_code,
346 detail,
347 )
348 raise ModelArmorAPIError(detail) from None
350 self._log_response_debug(status_code=response.status_code, response_text=response.text)
352 if response.status_code != 200:
353 detail = self._build_api_error_detail(response.status_code, response.text)
354 verbose_proxy_logger.error(
355 "Model Armor API error - Status: %s, Detail: %s",
356 response.status_code,
357 detail,
358 )
359 raise ModelArmorAPIError(detail)
361 json_response: Final = response.json()
362 if hasattr(json_response, "__await__"):
363 return await json_response
364 return json_response
366 def sanitize_file_prompt(self, file_bytes: bytes, file_type: str, source: str = "user_prompt") -> dict:
367 """
368 Helper to build the request body for file prompt sanitization for Model Armor.
369 file_type should be one of: PLAINTEXT_UTF8, PDF, WORD_DOCUMENT, EXCEL_DOCUMENT, POWERPOINT_DOCUMENT, TXT, CSV
370 Returns the request body dict.
371 """
372 import base64
374 base64_data: Final = base64.b64encode(file_bytes).decode("utf-8")
375 if source == "user_prompt":
376 return {"userPromptData": {"byteItem": {"byteDataType": file_type, "byteData": base64_data}}}
377 else:
378 return {"modelResponseData": {"byteItem": {"byteDataType": file_type, "byteData": base64_data}}}
380 def _should_block_content(self, armor_response: Mapping[str, object], allow_sanitization: bool = False) -> bool:
381 """Check if Model Armor response indicates content should be blocked, including both inspectResult and deidentifyResult."""
382 for filt in self._filter_result_items(armor_response):
383 # Check RAI, PI/Jailbreak, Malicious URI, CSAM, Virus scan as before
384 if filt.get("raiFilterResult", {}).get("matchState") == "MATCH_FOUND":
385 return True
386 if filt.get("piAndJailbreakFilterResult", {}).get("matchState") == "MATCH_FOUND":
387 return True
388 if filt.get("maliciousUriFilterResult", {}).get("matchState") == "MATCH_FOUND":
389 return True
390 if filt.get("csamFilterFilterResult", {}).get("matchState") == "MATCH_FOUND":
391 return True
392 if filt.get("virusScanFilterResult", {}).get("matchState") == "MATCH_FOUND":
393 return True
394 # Check sdpFilterResult for both inspectResult and deidentifyResult
395 sdp = filt.get("sdpFilterResult")
396 if sdp:
397 if sdp.get("inspectResult", {}).get("matchState") == "MATCH_FOUND":
398 return True
399 # Only block on deidentifyResult if sanitization is not allowed
400 if sdp.get("deidentifyResult", {}).get("matchState") == "MATCH_FOUND":
401 if not allow_sanitization:
402 return True
403 # Fallback dict code removed; all cases handled above
404 return False
406 def _get_sanitized_content(self, armor_response: Mapping[str, Any]) -> str | None:
407 """
408 Get the sanitized content from a Model Armor response, if available.
409 Looks for sanitized text in deidentifyResult, and falls back to root-level fields if not found.
410 """
411 filters: Final = self._filter_result_items(armor_response)
413 # Prefer sanitized text from deidentifyResult if present
414 for filter_entry in filters:
415 sdp = filter_entry.get("sdpFilterResult")
416 if sdp:
417 deid = sdp.get("deidentifyResult", {})
418 sanitized = deid.get("data", {}).get("text", "")
419 # If Model Armor found something and returned a sanitized version, use it
420 if deid.get("matchState") == "MATCH_FOUND" and sanitized:
421 return sanitized
423 # If no deidentifyResult, optionally check for inspectResult (rare, but could have findings)
424 for filter_entry in filters:
425 sdp = filter_entry.get("sdpFilterResult")
426 if sdp:
427 inspect = sdp.get("inspectResult", {})
428 # If Model Armor flagged something but didn't sanitize, return None
429 if inspect.get("matchState") == "MATCH_FOUND":
430 return None
432 # Fallback: if Model Armor put sanitized text at the root, use it
433 return armor_response.get("sanitizedText") or armor_response.get("text")
435 @staticmethod
436 def _filter_result_items(armor_response: Mapping[str, Any]) -> Sequence[Any]:
437 """Every filter result in a scan response.
439 filterResults is a dict of named filters on most templates and a list on some, so both
440 shapes are flattened to the same list of filter entries.
441 """
442 filter_results: Final = armor_response.get("sanitizationResult", {}).get("filterResults", {})
443 if isinstance(filter_results, dict):
444 return list(filter_results.values())
445 if isinstance(filter_results, list):
446 return filter_results
447 return []
449 def _has_deidentify_match(self, armor_response: Mapping[str, object]) -> bool:
450 """Whether an SDP de-identify filter matched, i.e. Model Armor owes this response a redaction."""
451 for filter_entry in self._filter_result_items(armor_response):
452 sdp = filter_entry.get("sdpFilterResult")
453 if sdp and sdp.get("deidentifyResult", {}).get("matchState") == "MATCH_FOUND":
454 return True
455 return False
457 def _resolve_streaming_outcome(
458 self,
459 armor_response: Mapping[str, object],
460 assembled_response: object,
461 content: str,
462 ) -> tuple[bool, str | None]:
463 """Whether to block the buffered stream, and the rewrite to emit when it is not blocked.
465 A de-identify match only reaches here unblocked because masking is on, so the redaction it
466 stands for has to be both resolvable and emittable. Where it is neither, the buffered
467 original still carries what Model Armor matched on, so this fails closed instead of
468 releasing it.
469 """
470 if self._should_block_content(armor_response, allow_sanitization=self.mask_response_content):
471 return True, None
472 if not self.mask_response_content:
473 return False, None
475 sanitized_content: Final = self._get_sanitized_content(armor_response)
476 if not sanitized_content:
477 # No rewrite to apply. Harmless unless a match is outstanding, in which case applying
478 # nothing would hand back the very content that matched
479 return self._has_deidentify_match(armor_response), None
480 if sanitized_content == content:
481 return False, None
482 if not isinstance(assembled_response, ModelResponse):
483 verbose_proxy_logger.warning(
484 "Model Armor: sanitized content cannot be re-emitted on this streaming endpoint, "
485 "blocking the response instead"
486 )
487 return True, None
488 return False, sanitized_content
490 @staticmethod
491 def _append_armor_response(existing: object, armor_response: Mapping[str, object]) -> object:
492 """Accumulate scan responses so a later text scan does not drop an earlier file scan.
494 Returns the single response on its own (backward compatible) and a list once a request
495 carries more than one scan. A list (not a tuple) is required because the guardrail logging
496 pipeline (redact_nested_match_and_regex_keys and the StandardLoggingGuardrailInformation
497 dict | list[dict] contract) only recurses into dicts and lists when redacting and serializing.
498 """
499 if existing is None:
500 return armor_response
501 if isinstance(existing, list):
502 return [*existing, armor_response] # mutable-ok: logging pipeline requires list[dict], not tuple
503 return [existing, armor_response] # mutable-ok: logging pipeline requires list[dict], not tuple
505 def _process_response(
506 self,
507 response: dict | None,
508 request_data: dict,
509 start_time: float | None = None,
510 end_time: float | None = None,
511 duration: float | None = None,
512 event_type: GuardrailEventHooks | None = None,
513 original_inputs: dict | None = None,
514 ):
515 """
516 Override to store only the Model Armor API response, not the entire data dict.
517 This prevents circular references in logging.
518 """
519 metadata: Final = (
520 request_data.get(get_metadata_variable_name_from_kwargs(request_data)) or {}
521 if isinstance(request_data, dict)
522 else {}
523 )
524 guardrail_response: Final = metadata.get("_model_armor_response", {})
526 # Determine status – default to "success" but prefer the explicit value if present.
527 guardrail_status: Final[GuardrailStatus] = metadata.get("_model_armor_status", "success")
529 self.add_standard_logging_guardrail_information_to_request_data(
530 guardrail_json_response=guardrail_response,
531 request_data=request_data,
532 guardrail_status=guardrail_status,
533 duration=duration,
534 start_time=start_time,
535 end_time=end_time,
536 event_type=event_type,
537 )
538 return response
540 @staticmethod
541 def _unscannable_block_error(reason: str) -> HTTPException:
542 return HTTPException(
543 status_code=400,
544 detail={"error": f"Model Armor could not scan an attachment and blocked the request: {reason}"},
545 )
547 async def _scan_request_files(self, messages: Sequence[AllMessageValues], data: dict) -> None:
548 """Submit inline document/file attachments to Model Armor and block on any findings.
550 Each attachment is sent through the byte API and a MATCH_FOUND raises a 400 before the
551 request reaches the LLM. File scanning does not support masking (Model Armor returns
552 findings, not a sanitized document), so it only blocks. A file_id or remote URL reference
553 with no inline bytes and a document over the 4 MB byte limit are guardrail failures that
554 block unless the operator has opted into fail-open via fail_on_error=False.
556 skip_unscannable_attachments decouples reference-only attachments from fail_on_error: when
557 enabled, attachments Model Armor cannot scan (file_id, gs://, or http(s) references with no
558 inline bytes, and inline content whose base64 will not decode) pass through instead of
559 blocking, while fail_on_error still governs real Model Armor API errors.
560 """
561 from litellm.proxy.common_utils.callback_utils import (
562 add_guardrail_to_applied_guardrails_header,
563 )
565 plan: Final = plan_file_scans(messages)
566 attachments: Final = plan.attachments
567 skip_unscannable: Final = bool(self.optional_params.get("skip_unscannable_attachments", False))
568 if skip_unscannable and plan.unscannable_count > 0:
569 verbose_proxy_logger.warning(
570 "Model Armor: allowing %d unscannable attachment(s) through because "
571 "skip_unscannable_attachments is enabled",
572 plan.unscannable_count,
573 )
574 unscannable_references: Final = 0 if skip_unscannable else plan.unscannable_count
575 if not attachments and unscannable_references == 0:
576 return
578 add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
579 # Use the same metadata bucket the header helper writes to, so the logged Model Armor
580 # payload and status land where _process_response reads them on every route.
581 _, metadata = get_or_create_metadata_bucket(data)
582 fail_on_error: Final = bool(self.optional_params.get("fail_on_error", True))
584 if unscannable_references > 0:
585 reason = (
586 f"{unscannable_references} attachment(s) reference a document with no inline bytes "
587 "(file_id or remote URL) that Model Armor cannot scan"
588 )
589 verbose_proxy_logger.warning("Model Armor: %s", reason)
590 if fail_on_error:
591 metadata["_model_armor_status"] = "blocked"
592 raise self._unscannable_block_error(reason)
594 for attachment in attachments:
595 if len(attachment.file_bytes) > MODEL_ARMOR_MAX_FILE_SIZE_BYTES:
596 reason = (
597 f"attachment of {len(attachment.file_bytes)} bytes exceeds Model Armor's "
598 f"{MODEL_ARMOR_MAX_FILE_SIZE_BYTES} byte scan limit"
599 )
600 verbose_proxy_logger.warning("Model Armor: %s", reason)
601 if not fail_on_error:
602 continue
603 metadata["_model_armor_status"] = "blocked"
604 raise self._unscannable_block_error(reason)
606 try:
607 armor_response = await self.make_model_armor_request(
608 source="user_prompt",
609 request_data=data,
610 file_bytes=attachment.file_bytes,
611 file_type=attachment.byte_data_type,
612 )
613 except ModelArmorAPIError as e:
614 self._raise_if_fail_closed(e)
615 continue
616 except HTTPException:
617 raise
618 except Exception as e:
619 # Isolate transient errors per attachment so one failure does not leave the
620 # remaining attachments in the same request unscanned.
621 verbose_proxy_logger.error("Model Armor file scan error: %s", str(e), exc_info=True)
622 if fail_on_error:
623 raise
624 continue
626 # Model Armor returns findings for documents, not a sanitized file, so there is no
627 # masking fallback. Any finding must block, even when mask_request_content is enabled,
628 # otherwise a PII-only (SDP deidentify) document would pass through unscrubbed.
629 blocked = self._should_block_content(armor_response, allow_sanitization=False)
630 metadata["_model_armor_response"] = self._append_armor_response(
631 metadata.get("_model_armor_response"),
632 self._build_logging_response(armor_response),
633 )
634 if blocked or metadata.get("_model_armor_status") == "blocked":
635 metadata["_model_armor_status"] = "blocked"
636 else:
637 metadata["_model_armor_status"] = "success"
639 if blocked:
640 raise HTTPException(
641 status_code=400,
642 detail=self._build_block_error_detail("Content blocked by Model Armor", armor_response),
643 )
645 @log_guardrail_information
646 async def async_pre_call_hook(
647 self,
648 user_api_key_dict: UserAPIKeyAuth,
649 cache: DualCache,
650 data: dict,
651 call_type: CallTypesLiteral,
652 ) -> Exception | str | dict | None:
653 """Pre-call hook to sanitize user prompts."""
654 verbose_proxy_logger.debug("Inside Model Armor Pre-Call Hook")
656 from litellm.proxy.common_utils.callback_utils import (
657 add_guardrail_to_applied_guardrails_header,
658 )
660 event_type = GuardrailEventHooks.pre_call
661 if call_type == CallTypes.call_mcp_tool.value:
662 event_type = GuardrailEventHooks.pre_mcp_call
663 if self.should_run_guardrail(data=data, event_type=event_type) is not True:
664 return data
666 messages: Final = data.get("messages")
667 if not messages:
668 verbose_proxy_logger.warning("Model Armor: not running guardrail. No messages in data")
669 return data
671 # Extract content from messages using helper from common_utils
672 from litellm.litellm_core_utils.prompt_templates.common_utils import (
673 get_last_user_message,
674 )
676 await self._scan_request_files(messages=messages, data=data)
678 content: Final = get_last_user_message(messages)
679 if not content:
680 return data
682 # Make Model Armor request
683 try:
684 armor_response: Final = await self.make_model_armor_request(
685 content=content,
686 source="user_prompt",
687 request_data=data,
688 )
690 # Store the armor response for logging
691 # Attach Model Armor response + evaluation status directly to the per-request metadata to avoid
692 # race-conditions between concurrent requests which share the same guardrail instance.
693 # This ensures each request logs its own Model Armor response instead of a potentially stale value
694 # overwritten by another coroutine.
695 blocked: Final = self._should_block_content(armor_response, allow_sanitization=self.mask_request_content)
696 if isinstance(data, dict):
697 _, metadata = get_or_create_metadata_bucket(data) # ensures metadata exists and is unique per request
698 # Accumulate so a prior file scan on the same request is not overwritten by this text scan.
699 metadata["_model_armor_response"] = self._append_armor_response(
700 metadata.get("_model_armor_response"),
701 self._build_logging_response(armor_response),
702 )
703 # Pre-compute guardrail status for downstream logging. A blocked response will eventually raise
704 # an HTTPException, however in scenarios where the caller decides to ignore the exception (e.g.
705 # fail_on_error=False) we still want the correct status reflected.
706 if blocked or metadata.get("_model_armor_status") == "blocked":
707 metadata["_model_armor_status"] = "blocked"
708 else:
709 metadata["_model_armor_status"] = "success"
711 # Add guardrail to applied_guardrails BEFORE potential blocking
712 # This ensures guardrail is recorded even when it blocks the request
713 add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
715 # Check if content should be blocked
716 if blocked:
717 raise HTTPException(
718 status_code=400,
719 detail=self._build_block_error_detail("Content blocked by Model Armor", armor_response),
720 )
722 # If mask_request_content is enabled, update messages with sanitized content
723 if self.mask_request_content:
724 sanitized_content: Final = self._get_sanitized_content(armor_response)
725 if sanitized_content and sanitized_content != content:
726 # Use the helper to set the last user message with sanitized content
727 from litellm.litellm_core_utils.prompt_templates.common_utils import (
728 set_last_user_message,
729 )
731 data["messages"] = set_last_user_message(messages, sanitized_content)
733 except ModelArmorAPIError as e:
734 self._raise_if_fail_closed(e)
735 except HTTPException:
736 raise
737 except Exception as e:
738 verbose_proxy_logger.error("Model Armor pre-call error: %s", str(e), exc_info=True)
739 # Depending on configuration, either fail or continue
740 if self.optional_params.get("fail_on_error", True):
741 raise
743 return data
745 @log_guardrail_information
746 async def async_moderation_hook(
747 self,
748 data: dict,
749 user_api_key_dict: UserAPIKeyAuth,
750 call_type: CallTypesLiteral,
751 ) -> Exception | str | dict | None:
752 """During-call hook to sanitize user prompts in parallel with LLM call."""
753 verbose_proxy_logger.debug("Inside Model Armor Moderation Hook")
755 from litellm.proxy.common_utils.callback_utils import (
756 add_guardrail_to_applied_guardrails_header,
757 )
759 event_type = GuardrailEventHooks.during_call
760 if call_type == CallTypes.call_mcp_tool.value:
761 event_type = GuardrailEventHooks.during_mcp_call
762 if self.should_run_guardrail(data=data, event_type=event_type) is not True:
763 return data
765 messages: Final = data.get("messages")
766 if not messages:
767 verbose_proxy_logger.warning("Model Armor: not running guardrail. No messages in data")
768 return data
770 # Extract content from messages
771 from litellm.litellm_core_utils.prompt_templates.common_utils import (
772 get_last_user_message,
773 )
775 await self._scan_request_files(messages=messages, data=data)
777 content: Final = get_last_user_message(messages)
778 if not content:
779 return data
781 # Make Model Armor request
782 try:
783 armor_response: Final = await self.make_model_armor_request(
784 content=content,
785 source="user_prompt",
786 request_data=data,
787 )
789 blocked: Final = self._should_block_content(armor_response, allow_sanitization=self.mask_request_content)
790 # Store the armor response for logging
791 if isinstance(data, dict):
792 _, metadata = get_or_create_metadata_bucket(data)
793 # Accumulate so a prior file scan on the same request is not overwritten by this text scan.
794 metadata["_model_armor_response"] = self._append_armor_response(
795 metadata.get("_model_armor_response"),
796 self._build_logging_response(armor_response),
797 )
798 if blocked or metadata.get("_model_armor_status") == "blocked":
799 metadata["_model_armor_status"] = "blocked"
800 else:
801 metadata["_model_armor_status"] = "success"
803 # Add guardrail to applied_guardrails BEFORE potential blocking
804 # This ensures guardrail is recorded even when it blocks the request
805 add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
807 # Check if content should be blocked
808 if blocked:
809 raise HTTPException(
810 status_code=400,
811 detail=self._build_block_error_detail("Content blocked by Model Armor", armor_response),
812 )
814 # If mask_request_content is enabled, update messages with sanitized content
815 if self.mask_request_content:
816 sanitized_content: Final = self._get_sanitized_content(armor_response)
817 if sanitized_content and sanitized_content != content:
818 from litellm.litellm_core_utils.prompt_templates.common_utils import (
819 set_last_user_message,
820 )
822 data["messages"] = set_last_user_message(messages, sanitized_content)
824 except ModelArmorAPIError as e:
825 self._raise_if_fail_closed(e)
826 except HTTPException:
827 raise
828 except Exception as e:
829 verbose_proxy_logger.error("Model Armor moderation error: %s", str(e), exc_info=True)
830 if self.optional_params.get("fail_on_error", True):
831 raise
833 return data
835 @log_guardrail_information
836 async def async_post_call_success_hook(
837 self,
838 data: dict,
839 user_api_key_dict: UserAPIKeyAuth,
840 response,
841 ):
842 """Post-call hook to sanitize model responses."""
843 from litellm.proxy.common_utils.callback_utils import (
844 add_guardrail_response_to_standard_logging_object,
845 add_guardrail_to_applied_guardrails_header,
846 )
848 if self.should_run_guardrail(data=data, event_type=GuardrailEventHooks.post_call) is not True:
849 return
851 # Extract content from response
852 content: Final = self._extract_content_from_response(response)
853 if not content:
854 verbose_proxy_logger.debug("Model Armor: No text content to process in response, skipping guardrail")
855 return
857 # Make Model Armor request
858 try:
859 armor_response: Final = await self.make_model_armor_request(
860 content=content,
861 source="model_response",
862 request_data=data,
863 )
865 # Attach Model Armor response & status to this request's metadata to prevent race conditions
866 if isinstance(armor_response, dict):
867 model_armor_logged_object: Final = {
868 "model_armor_response": self._build_logging_response(armor_response),
869 "model_armor_status": (
870 "blocked"
871 if self._should_block_content(
872 armor_response,
873 allow_sanitization=self.mask_response_content,
874 )
875 else "success"
876 ),
877 }
878 standard_logging_guardrail_information: Final = StandardLoggingGuardrailInformation(
879 guardrail_name=self.guardrail_name,
880 guardrail_provider="model_armor",
881 guardrail_mode=GuardrailEventHooks.post_call,
882 guardrail_response=model_armor_logged_object,
883 guardrail_status="success",
884 start_time=data.get("start_time"),
885 )
886 add_guardrail_response_to_standard_logging_object(
887 litellm_logging_obj=data.get("litellm_logging_obj"),
888 guardrail_response=standard_logging_guardrail_information,
889 )
891 # Add guardrail to applied_guardrails BEFORE potential blocking
892 # This ensures guardrail is recorded even when it blocks the request
893 add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
895 # Check if content should be blocked
896 if self._should_block_content(armor_response, allow_sanitization=self.mask_response_content):
897 raise HTTPException(
898 status_code=400,
899 detail=self._build_block_error_detail("Response blocked by Model Armor", armor_response),
900 )
902 # If mask_response_content is enabled, update response with sanitized content
903 if self.mask_response_content:
904 sanitized_content: Final = self._get_sanitized_content(armor_response)
905 if sanitized_content and sanitized_content != content:
906 # Update response content
907 if isinstance(response, litellm.ModelResponse):
908 for choice in response.choices:
909 if isinstance(choice, Choices):
910 if choice.message.content:
911 choice.message.content = sanitized_content
913 except ModelArmorAPIError as e:
914 self._raise_if_fail_closed(e)
915 except HTTPException:
916 raise
917 except Exception as e:
918 verbose_proxy_logger.error("Model Armor post-call error: %s", str(e), exc_info=True)
919 if self.optional_params.get("fail_on_error", True):
920 raise
922 return response
924 @staticmethod
925 def _is_terminal_error_stream(all_chunks: Sequence[object]) -> bool:
926 """Whether the buffered stream is only the refusal an earlier guardrail in the chain emitted.
928 post_call guardrails are composed, so this hook can be handed the terminal error items a
929 preceding one produced. They carry no message to scan, and replacing them would hide the
930 refusal the client is owed.
931 """
932 if all(getattr(chunk, "type", None) == "error" for chunk in all_chunks):
933 return True
934 return is_sse_error_stream(all_chunks)
936 @staticmethod
937 def _classify_stream(all_chunks: Sequence[object]) -> _StreamSurface:
938 """Wire format the buffered chunks belong to."""
939 if is_raw_sse_stream(all_chunks):
940 return (
941 _StreamSurface.ANTHROPIC_MESSAGES if is_anthropic_sse_stream(all_chunks) else _StreamSurface.OPAQUE_SSE
942 )
943 if any(
944 isinstance(event_type := getattr(chunk, "type", None), str) and event_type.startswith("response.")
945 for chunk in all_chunks
946 ):
947 return _StreamSurface.RESPONSES
948 return _StreamSurface.CHAT_COMPLETIONS
950 @staticmethod
951 def _final_responses_api_response(all_chunks: Sequence[object]) -> ResponsesAPIResponse | None:
952 """Response body carried by a terminal ``/v1/responses`` event.
954 A stream cut short before it completes has to read as unassembled rather than as a clean
955 empty response: ``response.created`` also carries a body, but an empty one, and scanning
956 that would release every buffered delta unscanned.
957 """
958 return next(
959 (
960 body
961 for chunk in reversed(all_chunks)
962 if getattr(chunk, "type", None) in _RESPONSES_TERMINAL_EVENT_TYPES
963 and isinstance(body := getattr(chunk, "response", None), ResponsesAPIResponse)
964 ),
965 None,
966 )
968 @staticmethod
969 def _responses_api_response_text(response: ResponsesAPIResponse) -> str:
970 """Text to scan in a Responses API response, tool-call arguments included.
972 Tool calls are folded in because ``get_content_from_model_response`` folds them into what
973 the chat surface scans, and a Responses turn can carry its whole payload in them.
974 """
975 from litellm.llms.openai.responses.guardrail_translation.handler import (
976 OpenAIResponsesHandler,
977 )
979 texts: Final[list[str]] = [] # mutable-ok: the shared extractor below appends into caller-owned lists
980 tool_calls: Final[list[ChatCompletionToolCallChunk]] = [] # mutable-ok: the same extractor's tool-call sink
981 handler: Final = OpenAIResponsesHandler()
982 for output_idx, output_item in enumerate(response.output or ()):
983 handler._extract_output_text_and_images( # pyright: ignore[reportPrivateUsage] # the shared Responses output extractor; forking it would duplicate per-item parsing
984 output_item=output_item,
985 output_idx=output_idx,
986 texts_to_check=texts,
987 images_to_check=[], # mutable-ok: the extractor's images sink, unused here
988 task_mappings=[], # mutable-ok: the extractor's task-mapping sink, unused here
989 tool_calls_to_check=tool_calls,
990 )
991 return "".join((*texts, *(json.dumps(tool_call) for tool_call in tool_calls)))
993 def _extract_streaming_content(self, assembled_response: object) -> str:
994 """Text to scan from an assembled stream, for every endpoint shape this hook serves."""
995 if isinstance(assembled_response, ResponsesAPIResponse):
996 return self._responses_api_response_text(assembled_response)
997 return self._extract_content_from_response(assembled_response)
999 @staticmethod
1000 def _responses_delta_field(chunk: object) -> tuple[str, ...]:
1001 """Which field of the turn a delta event belongs to."""
1002 return tuple(str(getattr(chunk, attr, None)) for attr in _RESPONSES_DELTA_FIELD_ATTRS)
1004 @staticmethod
1005 def _responses_delta_field_texts(all_chunks: Sequence[object]) -> tuple[str, ...]:
1006 """Text each field of a ``/v1/responses`` turn has already spelled out in its delta events.
1008 One field's deltas are joined as they streamed, since a finding can be split across them,
1009 and separate fields stay apart, so a reasoning summary running into the visible answer
1010 cannot spell out a finding that neither of them carries.
1011 """
1012 deltas: Final = tuple(
1013 (ModelArmorGuardrail._responses_delta_field(chunk), delta)
1014 for chunk in all_chunks
1015 if getattr(chunk, "type", None) in _RESPONSES_DELTA_EVENT_TYPES
1016 and isinstance(delta := getattr(chunk, "delta", None), str)
1017 )
1018 return tuple(
1019 "".join(delta for field, delta in deltas if field == streamed_field)
1020 for streamed_field in dict.fromkeys(field for field, _ in deltas)
1021 )
1023 def _streaming_content_to_scan(
1024 self,
1025 assembled_response: object,
1026 all_chunks: Sequence[object],
1027 surface: _StreamSurface,
1028 ) -> str:
1029 """Text to scan for a buffered stream, which is everything the client is about to receive.
1031 A ``/v1/responses`` stream also spells out reasoning summaries and tool-call arguments in
1032 delta events that its terminal body never repeats, so every delta field the body does not
1033 already carry is scanned after it.
1034 """
1035 content: Final = self._extract_streaming_content(assembled_response)
1036 if surface is not _StreamSurface.RESPONSES:
1037 return content
1038 unscanned: Final = tuple(text for text in self._responses_delta_field_texts(all_chunks) if text not in content)
1039 return "\n".join(part for part in (content, *unscanned) if part)
1041 @staticmethod
1042 def _apply_sanitized_content(assembled_response: ModelResponse, sanitized_content: str) -> None:
1043 """Replace every non-empty choice message with the Model Armor sanitized text."""
1044 for choice in assembled_response.choices:
1045 if isinstance(choice, Choices) and choice.message.content:
1046 choice.message.content = sanitized_content
1048 @staticmethod
1049 def _assemble_chat_completion_stream(
1050 all_chunks: list[object], # mutable-ok: stream_chunk_builder only accepts a mutable list
1051 ) -> ModelResponse | TextCompletionResponse | None:
1052 """Assemble chat-completion chunks, returning ``None`` when they cannot be assembled."""
1053 from litellm.main import stream_chunk_builder
1055 try:
1056 return stream_chunk_builder(chunks=all_chunks)
1057 except Exception as exc:
1058 verbose_proxy_logger.warning("Model Armor: chat-completion stream assembly failed (%s)", exc)
1059 return None
1061 def _assemble_stream(
1062 self, all_chunks: Sequence[object], surface: _StreamSurface
1063 ) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | None:
1064 """Assemble the buffered stream into the scannable response its surface produces."""
1065 if surface is _StreamSurface.ANTHROPIC_MESSAGES:
1066 return assemble_anthropic_sse_stream(all_chunks, restore_identity=True)
1067 if surface is _StreamSurface.RESPONSES:
1068 return self._final_responses_api_response(all_chunks)
1069 if surface is _StreamSurface.OPAQUE_SSE:
1070 return None
1071 return self._assemble_chat_completion_stream(list(all_chunks))
1073 @staticmethod
1074 def _error_payload(exc: HTTPException) -> Mapping[str, object]:
1075 """Error object for a terminal stream item, carrying the status the frame would otherwise lose."""
1076 detail: Final = exc.detail if isinstance(exc.detail, Mapping) else {"message": str(exc.detail)}
1077 error_value: Final = detail.get("error", detail)
1078 return {
1079 **(dict(error_value) if isinstance(error_value, Mapping) else {"message": str(error_value)}),
1080 "code": str(exc.status_code),
1081 }
1083 @staticmethod
1084 def _build_responses_error_items(exc: HTTPException) -> Sequence[object] | None:
1085 """Responses API error events for a failure discovered after the stream started."""
1086 from litellm.llms.openai.responses.guardrail_translation.handler import (
1087 OpenAIResponsesHandler,
1088 )
1090 return OpenAIResponsesHandler().build_stream_error_items(exc, responses_so_far=None)
1092 def _stream_error_items(self, exc: HTTPException, *, surface: _StreamSurface) -> Sequence[object]:
1093 """Frame a guardrail failure as terminal stream items in this endpoint's wire format."""
1094 payload: Final = self._error_payload(exc)
1095 if surface is _StreamSurface.ANTHROPIC_MESSAGES:
1096 return anthropic_sse_error_frames(str(payload.get("message", "")))
1097 if surface is _StreamSurface.RESPONSES and (responses_items := self._build_responses_error_items(exc)):
1098 return responses_items
1099 # Also the fallback when a surface cannot frame its own error: create_response() reads the
1100 # status back out of this form, so the refusal keeps its code instead of arriving as a 200
1101 return (f"data: {json.dumps({'error': payload})}\n\n",)
1103 async def async_post_call_streaming_iterator_hook(
1104 self,
1105 user_api_key_dict: UserAPIKeyAuth,
1106 response: Any,
1107 request_data: dict,
1108 ) -> AsyncGenerator[ModelResponseStream, None]:
1109 """Process streaming response chunks."""
1111 from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
1112 from litellm.proxy.common_utils.callback_utils import (
1113 add_guardrail_to_applied_guardrails_header,
1114 )
1116 if self.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is not True:
1117 async for chunk in response:
1118 yield chunk
1119 return
1121 all_chunks: Final[Sequence[object]] = tuple([chunk async for chunk in response])
1123 if not all_chunks or self._is_terminal_error_stream(all_chunks):
1124 for chunk in all_chunks:
1125 yield chunk
1126 return
1128 surface: Final = self._classify_stream(all_chunks)
1130 # Build complete response
1131 assembled_response: Final = self._assemble_stream(all_chunks, surface)
1133 if assembled_response is None:
1134 if not self.optional_params.get("fail_on_error", True):
1135 verbose_proxy_logger.warning(
1136 "Model Armor: streamed response could not be assembled for scanning, "
1137 "forwarding it unscanned because fail_on_error is disabled"
1138 )
1139 for chunk in all_chunks:
1140 yield chunk
1141 return
1143 # Forwarding an unscannable stream would silently disable the guardrail, so fail closed
1144 add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name)
1145 for error_item in self._stream_error_items(
1146 HTTPException(
1147 status_code=500,
1148 detail=f"{self.guardrail_name}: streamed response could not be assembled for scanning, blocking it",
1149 ),
1150 surface=surface,
1151 ):
1152 yield error_item
1153 return
1155 # Extract content
1156 content: Final = self._streaming_content_to_scan(
1157 assembled_response=assembled_response, all_chunks=all_chunks, surface=surface
1158 )
1160 if not content:
1161 verbose_proxy_logger.debug("Model Armor: No text content in streaming response, skipping guardrail")
1162 for chunk in all_chunks:
1163 yield chunk
1164 return
1166 try:
1167 # Check with Model Armor
1168 armor_response: Final = await self.make_model_armor_request(
1169 content=content,
1170 source="model_response",
1171 request_data=request_data,
1172 )
1174 # Decide the outcome before recording it. Mirrors the non-streaming sibling: with
1175 # masking on, a de-identify match is a redaction to apply rather than a refusal, but
1176 # that only holds while the redaction can actually be delivered
1177 blocked, sanitized_content = self._resolve_streaming_outcome(
1178 armor_response=armor_response,
1179 assembled_response=assembled_response,
1180 content=content,
1181 )
1183 # Attach Model Armor response & status to this request's metadata to avoid race conditions
1184 if isinstance(request_data, dict):
1185 _, metadata = get_or_create_metadata_bucket(request_data)
1186 metadata["_model_armor_response"] = self._build_logging_response(armor_response)
1187 metadata["_model_armor_status"] = "blocked" if blocked else "success"
1189 # Add guardrail to applied_guardrails BEFORE potential blocking
1190 # This ensures guardrail is recorded even when it blocks the request
1191 add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name)
1193 if blocked:
1194 raise HTTPException(
1195 status_code=400,
1196 detail=self._build_block_error_detail(
1197 "Streaming response blocked by Model Armor",
1198 armor_response,
1199 ),
1200 )
1202 if sanitized_content is not None and isinstance(assembled_response, ModelResponse):
1203 self._apply_sanitized_content(assembled_response, sanitized_content)
1205 # Return sanitized stream
1206 if surface is _StreamSurface.ANTHROPIC_MESSAGES:
1207 for sse_chunk in anthropic_sse_chunks_from_response(assembled_response):
1208 yield sse_chunk
1209 return
1210 mock_response: Final = MockResponseIterator(model_response=assembled_response)
1211 async for chunk in mock_response:
1212 yield chunk
1213 return
1215 except ModelArmorAPIError as e:
1216 if self.optional_params.get("fail_on_error", True):
1217 for error_item in self._stream_error_items(
1218 HTTPException(status_code=500, detail=e.detail), surface=surface
1219 ):
1220 yield error_item
1221 return
1222 except HTTPException as e:
1223 # Yield the error as a terminal stream item so create_response() detects it and returns
1224 # a proper JSON error response with the correct status code. Raising from a generator
1225 # instead hits create_response's generic except and becomes a 500.
1226 for error_item in self._stream_error_items(e, surface=surface):
1227 yield error_item
1228 return
1229 except Exception as e:
1230 verbose_proxy_logger.error("Model Armor streaming error: %s", str(e), exc_info=True)
1231 if self.optional_params.get("fail_on_error", True):
1232 raise
1234 # Return original chunks if no sanitization needed
1235 for chunk in all_chunks:
1236 yield chunk
1238 @log_guardrail_information
1239 async def apply_guardrail(
1240 self,
1241 inputs: GenericGuardrailAPIInputs,
1242 request_data: dict,
1243 input_type: Literal["request", "response"],
1244 logging_obj: "LiteLLMLoggingObj | None" = None,
1245 ) -> GenericGuardrailAPIInputs:
1246 content: Final = "\n".join(text for text in inputs.get("texts") or () if text)
1247 if not content:
1248 return inputs
1250 source: Final[Literal["user_prompt", "model_response"]] = (
1251 "user_prompt" if input_type == "request" else "model_response"
1252 )
1253 start_time: Final = time.time()
1254 try:
1255 armor_response: Final = await self.make_model_armor_request(
1256 content=content, source=source, request_data=request_data
1257 )
1258 except (ModelArmorAPIError, httpx.HTTPError) as e:
1259 error_end_time: Final = time.time()
1260 self.add_standard_logging_guardrail_information_to_request_data(
1261 guardrail_json_response=str(e),
1262 request_data=request_data,
1263 guardrail_status="guardrail_failed_to_respond",
1264 guardrail_provider="model_armor",
1265 start_time=start_time,
1266 end_time=error_end_time,
1267 duration=error_end_time - start_time,
1268 )
1269 return inputs
1271 flagged: Final = self._should_block_content(armor_response, allow_sanitization=False)
1272 end_time: Final = time.time()
1273 self.add_standard_logging_guardrail_information_to_request_data(
1274 guardrail_json_response=self._build_logging_response(armor_response),
1275 request_data=request_data,
1276 guardrail_status="guardrail_flagged" if flagged else "success",
1277 guardrail_provider="model_armor",
1278 start_time=start_time,
1279 end_time=end_time,
1280 duration=end_time - start_time,
1281 )
1282 if flagged and not self._event_hook_is_event_type(GuardrailEventHooks.logging_only):
1283 raise HTTPException(
1284 status_code=400,
1285 detail=self._build_block_error_detail(
1286 "Response blocked by Model Armor" if input_type == "response" else "Content blocked by Model Armor",
1287 armor_response,
1288 ),
1289 )
1290 return inputs
1292 @staticmethod
1293 def get_config_model() -> type["GuardrailConfigModel"] | None:
1294 """
1295 Get the config model for the Model Armor guardrail.
1296 """
1297 from litellm.types.proxy.guardrails.guardrail_hooks.model_armor import (
1298 ModelArmorGuardrailConfigModel,
1299 )
1301 return ModelArmorGuardrailConfigModel