Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py: 12%
294 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1import copy
2import os
3from collections.abc import Mapping, Sequence
4from datetime import datetime
5from string import Formatter
6from types import MappingProxyType
7from typing import Final, Literal
9from fastapi import HTTPException
11import litellm
12from litellm._logging import verbose_proxy_logger
13from litellm.integrations.custom_guardrail import (
14 DEFAULT_ADVISORY_MESSAGE,
15 CustomGuardrail,
16)
17from litellm.llms.base_llm.guardrail_translation.utils import (
18 effective_skip_system_message_for_guardrail,
19 effective_skip_tool_message_for_guardrail,
20 filter_messages_by_skip_flags,
21 merge_guardrailed_scoped_messages,
22)
23from litellm.llms.custom_httpx.http_handler import (
24 get_async_httpx_client,
25 httpxSpecialProvider,
26)
27from litellm.proxy._types import UserAPIKeyAuth
28from litellm.proxy.guardrails._content_utils import (
29 apply_redacted_messages_back,
30 build_inspection_messages,
31 has_non_string_content,
32)
33from litellm.secret_managers.main import get_secret_str
34from litellm.types.guardrails import GuardrailEventHooks, LitellmParams
35from litellm.types.llms.openai import AllMessageValues
36from litellm.types.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import (
37 LakeraAIBreakdownItem,
38 LakeraAIRequest,
39 LakeraAIResponse,
40)
41from litellm.types.utils import CallTypesLiteral, GuardrailStatus, ModelResponse
43_DETECTOR_CATEGORY_PHRASES: Final[Mapping[str, str]] = MappingProxyType(
44 {
45 "prompt_injection": "a potential prompt injection attempt",
46 "prompt_attack": "a potential prompt injection attempt",
47 "pii": "personally identifiable information",
48 "moderated_content": "policy-violating content",
49 }
50)
53def humanize_lakera_block_reasons(breakdown: Sequence[LakeraAIBreakdownItem] | None) -> str:
54 """
55 Turn a Lakera v2 ``breakdown`` list into a plain-language reason string
56 suitable for an advisory message shown to the LLM (e.g. "a potential
57 prompt injection attempt, personally identifiable information").
59 Falls back to a generic phrase when breakdown is empty or every detected
60 detector_type is unrecognized.
61 """
62 if not breakdown:
63 return "a content safety concern"
65 categories: Final = (
66 (item.get("detector_type") or "").split("/")[0] for item in breakdown if item.get("detected", False)
67 )
68 phrases: Final = tuple(
69 dict.fromkeys(
70 _DETECTOR_CATEGORY_PHRASES.get(category) or category.replace("_", " ")
71 for category in categories
72 if category
73 )
74 )
75 return ", ".join(phrases) if phrases else "a content safety concern"
78def _template_uses_reason_placeholder(template: str) -> bool:
79 """True if ``template`` has a real ``{reason}`` format field, not just the
80 literal substring -- an escaped ``{{reason}}`` contains the substring but
81 formats to a literal "{reason}", never substituting the actual value."""
82 return any(field_name == "reason" for _, field_name, _, _ in Formatter().parse(template))
85def _pre_masking_scope_indices(
86 guardrail: "LakeraAIGuardrail",
87 messages: Sequence[object],
88) -> tuple[int, ...]:
89 """Indices into ``messages`` that mask-in-place can safely target: has
90 non-empty string content, and survives the same skip_system_message_in_guardrail
91 / skip_tool_message_in_guardrail scoping ``filter_messages_by_skip_flags``
92 applies. Content is guaranteed to already be a plain string here -- masking
93 is only attempted when ``has_non_string_content(data)`` is False.
95 Preserved in original order, so it lines up positionally with the
96 ``messages_for_lakera`` list _build_lakera_inspection_messages/skip-filtering
97 produces from the same input: both apply the identical "has text" and
98 "not skipped by role" predicates over the same original sequence. Role
99 comparison is lowercased to match filter_messages_by_skip_flags's own
100 normalization (via its _message_role helper) -- an uppercase-cased
101 "System"/"TOOL" role must be excluded by both or the two lists disagree
102 on length and the caller's strict positional zip raises."""
103 skip_system: Final = effective_skip_system_message_for_guardrail(guardrail)
104 skip_tool: Final = effective_skip_tool_message_for_guardrail(guardrail)
105 return tuple(
106 idx
107 for idx, message in enumerate(messages)
108 if isinstance(message, dict)
109 and isinstance(message.get("content"), str)
110 and message["content"]
111 and not (skip_system and str(message.get("role") or "").lower() == "system")
112 and not (skip_tool and str(message.get("role") or "").lower() == "tool")
113 )
116def _apply_redacted_messages_back_preserving_fields(
117 guardrail: "LakeraAIGuardrail",
118 data: dict[str, object], # mutable-ok: writes the redacted result back into the caller's request dict in place
119 redacted_messages: Sequence[AllMessageValues],
120) -> None:
121 """Write masked content back to ``data["messages"]`` without losing fields
122 the synthetic role/content-only ``redacted_messages`` never carried (e.g. a
123 tool message's tool_call_id, an assistant message's tool_calls, name,
124 cache_control). Falls back to the shared, wholesale-replacing
125 apply_redacted_messages_back when ``data["messages"]`` isn't a list (a pure
126 Responses-API ``input`` string, with no chat messages to merge into)."""
127 original_messages: Final = data.get("messages")
128 if not isinstance(original_messages, list):
129 redacted_list: Final = list(redacted_messages) # mutable-ok: apply_redacted_messages_back requires a list
130 apply_redacted_messages_back(data, redacted_list)
131 return
132 scope_indices: Final = _pre_masking_scope_indices(guardrail, original_messages)
133 guardrailed_scoped: Final = tuple(
134 { # mutable-ok: fresh dict per iteration, not stored beyond this comprehension
135 **original_messages[original_idx],
136 "content": redacted["content"],
137 }
138 for original_idx, redacted in zip(scope_indices, redacted_messages, strict=True)
139 )
140 data["messages"] = merge_guardrailed_scoped_messages(
141 full_messages=original_messages,
142 scoped_indices=scope_indices,
143 guardrailed_scoped=guardrailed_scoped, # pyright: ignore[reportArgumentType] # plain dicts satisfy AllMessageValues's TypedDict shape at runtime
144 )
147def _has_combined_messages_and_input(data: Mapping[str, object]) -> bool:
148 """True if ``data`` carries both ``messages`` and ``input``.
149 build_inspection_messages flattens both into one synthetic list, so
150 mask-in-place would write input-derived content into data["messages"]
151 (and vice versa) even when a message dropped for having no text
152 coincidentally keeps the raw message count unchanged."""
153 return isinstance(data.get("messages"), list) and data.get("input") is not None
156def _has_responses_instructions(guardrail: "LakeraAIGuardrail", data: Mapping[str, object]) -> bool:
157 """True if ``data`` carries a Responses-API ``instructions`` field that
158 Lakera actually inspected. _build_lakera_inspection_messages includes
159 ``instructions`` as a synthetic system message so Lakera can inspect it,
160 but apply_redacted_messages_back has no path to rewrite
161 ``data["instructions"]`` -- masking here would either leave unredacted
162 content in the real instructions field the model reads, or write a
163 redacted duplicate into data["messages"] instead, which the Responses
164 API never consumes.
166 When skip_system_message_in_guardrail excludes that synthetic system
167 message before it ever reaches Lakera, none of this applies: Lakera never
168 saw ``instructions``, so it can't have flagged anything there, and
169 forcing a hard block anyway would defeat the whole point of the skip
170 flag for a response that only carries PII in the (maskable) non-system
171 content."""
172 instructions: Final = data.get("instructions")
173 return (
174 isinstance(instructions, str)
175 and bool(instructions)
176 and not effective_skip_system_message_for_guardrail(guardrail)
177 )
180def _breakdown_has_pii_violation(lakera_response: LakeraAIResponse | None) -> bool:
181 """True if any PII-category detector fired, regardless of whether other,
182 non-PII detectors (prompt injection, moderated content) also fired.
183 Unlike ``_is_only_pii_violation``, this doesn't require PII to be the
184 *only* thing detected -- it's used to decide whether masking/blocking is
185 even relevant at all before advisory mode's own logic runs."""
186 if not lakera_response:
187 return False
188 breakdown: Final = lakera_response.get("breakdown") or ()
189 return any(
190 item.get("detected", False) and (item.get("detector_type") or "").startswith("pii/") for item in breakdown
191 )
194def _unmaskable_reason(
195 guardrail: "LakeraAIGuardrail",
196 data: dict[str, object],
197 lakera_response: LakeraAIResponse | None,
198) -> str | None:
199 """Why a PII-only violation on ``data`` can't be masked in place, or None when it can."""
200 if has_non_string_content(data):
201 return "multimodal content, masking would drop the image/audio parts"
202 if _has_combined_messages_and_input(data):
203 return "messages and input are both present, so the write-back is positionally ambiguous"
204 if "messages" in data and not isinstance(data.get("messages"), list):
205 return "a messages key that isn't a list, so there's nothing to merge the redacted content into"
206 if not _has_responses_instructions(guardrail, data):
207 return "no write-back path for the redacted content"
208 if not (lakera_response or {}).get("payload"):
209 return "Lakera reported no locations to redact, so payload=true is likely off"
210 return None
213def _build_lakera_inspection_messages(data: Mapping[str, object]) -> Sequence[Mapping[str, str]]:
214 """Like build_inspection_messages, but also covers the Responses-API
215 ``instructions`` field, placed first since litellm later converts it
216 into the model's leading system message and a prompt-injection detector
217 should see the same conversation order the model actually receives.
219 Kept local to Lakera rather than folded into the shared
220 _content_utils.build_inspection_messages helper: doing that once made
221 ``instructions`` visible to every guardrail sharing that helper (AIM,
222 presidio, bedrock, ...), but only Lakera has a masking-safety-guard
223 (_has_responses_instructions) accounting for apply_redacted_messages_back
224 having no write-back path for data["instructions"] -- other guardrails
225 would have silently mishandled a PII/redaction hit found there."""
226 instructions: Final = data.get("instructions")
227 leading: Final[Sequence[Mapping[str, str]]] = (
228 [{"role": "system", "content": instructions}] # mutable-ok: fresh list/dict, not stored
229 if isinstance(instructions, str) and instructions
230 else [] # mutable-ok: fresh empty list, not stored
231 )
232 return [ # mutable-ok: fresh list, not stored
233 *leading,
234 *build_inspection_messages(dict(data)), # mutable-ok: fresh shallow copy for the dict[str, Any] param
235 ]
238class LakeraAIGuardrail(CustomGuardrail):
239 @classmethod
240 def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
241 return [
242 GuardrailEventHooks.pre_call,
243 GuardrailEventHooks.during_call,
244 GuardrailEventHooks.post_call,
245 ]
247 def __init__(
248 self,
249 api_key: str | None = None,
250 api_base: str | None = None,
251 project_id: str | None = None,
252 payload: bool | None = True,
253 breakdown: bool | None = True,
254 metadata: dict | None = None,
255 dev_info: bool | None = True,
256 on_flagged: Literal["block", "monitor", "inject_system_message"] | None = "block",
257 skip_system_message_in_guardrail: bool | None = None,
258 skip_tool_message_in_guardrail: bool | None = None,
259 advisory_system_message: str | None = None,
260 **kwargs,
261 ):
262 """
263 Initialize the LakeraAIGuardrail class.
265 This guardrail only supports the chat completions endpoint (/v1/chat/completions).
266 It is not supported for the Responses API, /v1/messages, MCP, A2A, or other endpoints.
268 This calls: https://api.lakera.ai/v2/guard
270 Args:
271 api_key: Optional[str] = None,
272 api_base: Optional[str] = None,
273 project_id: Optional[str] = None,
274 payload: Optional[bool] = True,
275 breakdown: Optional[bool] = True,
276 metadata: Optional[Dict] = None,
277 dev_info: Optional[bool] = True,
278 on_flagged: Optional[str] = "block", Action to take when content is flagged:
279 "block", "monitor", or "inject_system_message"
280 skip_system_message_in_guardrail: Optional[bool] = None,
281 skip_tool_message_in_guardrail: Optional[bool] = None,
282 advisory_system_message: Optional[str] = None, custom advisory message template
283 (must contain a {reason} placeholder) used when on_flagged="inject_system_message".
284 Defaults to a generic message when unset.
285 """
286 self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
287 self.lakera_api_key = api_key or os.environ.get("LAKERA_API_KEY") or ""
288 self.project_id = project_id
289 self.api_base = api_base or get_secret_str("LAKERA_API_BASE") or "https://api.lakera.ai"
290 self.payload: bool | None = payload
291 self.breakdown: bool | None = breakdown
292 self.metadata: dict | None = metadata
293 self.dev_info: bool | None = dev_info
294 self.skip_system_message_in_guardrail = skip_system_message_in_guardrail
295 self.skip_tool_message_in_guardrail = skip_tool_message_in_guardrail
296 self.on_flagged = on_flagged or "block"
297 self.advisory_system_message = advisory_system_message
298 kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks()))
299 super().__init__(**kwargs)
300 self._validate_advisory_config(
301 on_flagged=self.on_flagged,
302 advisory_system_message=self.advisory_system_message,
303 payload=self.payload,
304 breakdown=self.breakdown,
305 )
307 def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None:
308 """
309 The base implementation blindly ``setattr``s every field on ``litellm_params``
310 (including ``on_flagged``/``advisory_system_message``/``payload``/``breakdown``)
311 onto this live instance with no revalidation, so an in-place config update (via
312 the DB/UI, without a restart) could otherwise reintroduce the exact invalid
313 on_flagged combinations __init__ rejects. Validate the prospective post-update
314 state *before* mutating, so a rejected update leaves the live instance untouched
315 instead of raising after it's already been corrupted.
317 The base setattr also writes ``litellm_params.mode`` onto a new ``self.mode``
318 attribute rather than the ``self.event_hook`` dispatch actually reads
319 (LitellmParams has no field literally named ``event_hook``), so without the
320 explicit sync below a hot reload that changes mode would pass validation but
321 keep dispatching on the stale event_hook.
322 """
323 new_event_hook: Final = litellm_params.mode or self.event_hook
324 prospective_payload: Final = litellm_params.payload
325 prospective_breakdown: Final = litellm_params.breakdown
326 self._validate_advisory_config(
327 on_flagged=litellm_params.on_flagged or self.on_flagged,
328 advisory_system_message=litellm_params.advisory_system_message,
329 payload=self.payload if prospective_payload is None else prospective_payload,
330 breakdown=self.breakdown if prospective_breakdown is None else prospective_breakdown,
331 )
332 super().update_in_memory_litellm_params(litellm_params=litellm_params)
333 self.event_hook = new_event_hook
335 def _validate_advisory_config(
336 self,
337 on_flagged: str,
338 advisory_system_message: str | None,
339 payload: bool | None,
340 breakdown: bool | None,
341 ) -> None:
342 if on_flagged == "inject_system_message" and advisory_system_message is not None:
343 if not _template_uses_reason_placeholder(advisory_system_message):
344 raise ValueError(
345 "Invalid advisory_system_message template: must include a real {reason} "
346 "placeholder (not an escaped {{reason}}) so the LLM sees why the request was flagged."
347 )
348 try:
349 advisory_system_message.format(reason="placeholder")
350 except (KeyError, IndexError, ValueError) as e:
351 raise ValueError(
352 f"Invalid advisory_system_message template: {e}. The template must be a valid "
353 "str.format() string using only the {reason} placeholder."
354 ) from e
355 if on_flagged == "inject_system_message" and not (payload and breakdown):
356 raise ValueError(
357 "on_flagged='inject_system_message' requires payload=True and breakdown=True: advisory "
358 "mode masks any detected PII before appending the advisory note, and that masking can "
359 "only happen when Lakera's response carries both the violation breakdown and the "
360 "payload location data. Without them, PII would be forwarded to the model unredacted."
361 )
363 def _build_advisory_message(self, lakera_response: LakeraAIResponse | None) -> str:
364 """Format the advisory message shown to the LLM when on_flagged='inject_system_message'."""
365 reason: Final = humanize_lakera_block_reasons(lakera_response.get("breakdown") if lakera_response else None)
366 template: Final = self.advisory_system_message or DEFAULT_ADVISORY_MESSAGE
367 return template.format(reason=reason)
369 def _filter_skipped_messages(
370 self, messages: Sequence[AllMessageValues]
371 ) -> tuple[tuple[AllMessageValues, ...], bool]:
372 return filter_messages_by_skip_flags(self, messages)
374 async def call_v2_guard(
375 self,
376 messages: Sequence[AllMessageValues],
377 request_data: dict,
378 event_type: GuardrailEventHooks,
379 ) -> tuple[LakeraAIResponse, dict]:
380 """
381 Call the Lakera AI v2 guard API.
382 """
383 status: GuardrailStatus = "success"
384 exception_str: str = ""
385 start_time: Final[datetime] = datetime.now()
386 lakera_response: LakeraAIResponse | None = None
387 request: dict = {}
388 masked_entity_count: Final[dict] = {}
389 try:
390 request = dict(
391 LakeraAIRequest(
392 messages=messages,
393 project_id=self.project_id,
394 payload=self.payload,
395 breakdown=self.breakdown,
396 metadata=self.metadata,
397 dev_info=self.dev_info,
398 )
399 )
400 verbose_proxy_logger.debug("Lakera AI v2 guard request: %s", request)
401 response: Final = await self.async_handler.post(
402 url=f"{self.api_base}/v2/guard",
403 headers={"Authorization": f"Bearer {self.lakera_api_key}"},
404 json=request,
405 )
406 verbose_proxy_logger.debug("Lakera AI v2 guard response: %s", response.json())
407 lakera_response = LakeraAIResponse(**response.json())
408 return lakera_response, masked_entity_count
409 except Exception as e:
410 status = "guardrail_failed_to_respond"
411 exception_str = str(e)
412 raise e
413 finally:
414 ####################################################
415 # Create Guardrail Trace for logging on Langfuse, Datadog, etc.
416 ####################################################
417 guardrail_json_response: Exception | str | dict | list[dict] = {}
418 if status == "success":
419 copy_lakera_response_dict: Final = dict(copy.deepcopy(lakera_response)) if lakera_response else {}
420 # payload contains PII, we don't want to log it
421 copy_lakera_response_dict.pop("payload")
422 guardrail_json_response = copy_lakera_response_dict
423 else:
424 guardrail_json_response = exception_str
425 self.add_standard_logging_guardrail_information_to_request_data(
426 guardrail_json_response=guardrail_json_response,
427 guardrail_status=status,
428 request_data=request_data,
429 start_time=start_time.timestamp(),
430 end_time=datetime.now().timestamp(),
431 duration=(datetime.now() - start_time).total_seconds(),
432 masked_entity_count=masked_entity_count,
433 event_type=event_type,
434 )
436 def _mask_pii_in_messages(
437 self,
438 messages: Sequence[AllMessageValues],
439 lakera_response: LakeraAIResponse | None,
440 masked_entity_count: dict,
441 ) -> Sequence[AllMessageValues]:
442 """
443 Return a copy of messages with any detected PII replaced by
444 “[MASKED <TYPE>]” tokens.
445 """
446 payload: Final = lakera_response.get("payload") if lakera_response else None
447 if not payload:
448 return messages
450 messages = copy.deepcopy(messages)
451 # For each message, find its detections on the fly
452 for idx, msg in enumerate(messages):
453 content = msg.get("content", "")
454 if not content:
455 continue
457 # For v1, we only support masking content strings
458 if not isinstance(content, str):
459 continue
461 # Filter only detections for this message
462 detected_modifications = [d for d in payload if d.get("message_id") == idx]
463 if not detected_modifications:
464 continue
466 # Apply masks from end to start so earlier indices remain valid after each replacement
467 detected_modifications = sorted(
468 detected_modifications,
469 key=lambda d: (d.get("start", 0), d.get("end", 0)),
470 reverse=True,
471 )
473 for modification in detected_modifications:
474 start, end = modification.get("start", 0), modification.get("end", 0)
476 # Extract the type (e.g. 'credit_card' → 'CREDIT_CARD')
477 detector_type = modification.get("detector_type", "")
478 if not detector_type:
479 continue
481 typ = detector_type.split("/")[-1].upper() or "PII"
482 mask = f"[MASKED {typ}]"
483 if start is not None and end is not None:
484 content = self.mask_content_in_string(
485 content_string=content,
486 mask_string=mask,
487 start_index=start,
488 end_index=end,
489 )
490 masked_entity_count[typ] = masked_entity_count.get(typ, 0) + 1
492 msg["content"] = content
493 return messages
495 def _mask_unwritable_instructions_pii_in_place(
496 self,
497 data: dict[str, object], # mutable-ok: writes the redacted result back into the caller's request dict in place
498 inspected_messages: Sequence[AllMessageValues],
499 lakera_response: LakeraAIResponse | None,
500 masked_entity_count: dict[str, int],
501 ) -> bool:
502 """Mask a body whose only obstacle to mask-in-place is the Responses-API
503 ``instructions`` field, writing the redacted instructions straight into
504 ``data["instructions"]``: apply_redacted_messages_back has no path for
505 that field and would fold the instructions text into ``data["input"]``.
506 Returns False without masking anything when _unmaskable_reason names an
507 obstacle this can't get around."""
508 if _unmaskable_reason(self, data, lakera_response) is not None:
509 return False
510 redacted: Final = self._mask_pii_in_messages(
511 messages=inspected_messages,
512 lakera_response=lakera_response,
513 masked_entity_count=masked_entity_count,
514 )
515 # _build_lakera_inspection_messages puts instructions first and
516 # _filter_skipped_messages kept it, so index 0 is the instructions.
517 data["instructions"] = redacted[0]["content"]
518 _apply_redacted_messages_back_preserving_fields(self, data, redacted[1:])
519 return True
521 async def async_pre_call_hook(
522 self,
523 user_api_key_dict: UserAPIKeyAuth,
524 cache: litellm.DualCache,
525 data: dict,
526 call_type: CallTypesLiteral,
527 ) -> Exception | str | dict | None:
528 from litellm.proxy.common_utils.callback_utils import (
529 add_guardrail_to_applied_guardrails_header,
530 )
532 verbose_proxy_logger.debug("Lakera AI: pre_call_hook")
534 event_type: Final[GuardrailEventHooks] = GuardrailEventHooks.pre_call
535 if self.should_run_guardrail(data=data, event_type=event_type) is not True:
536 verbose_proxy_logger.debug("Lakera AI: not running guardrail. Guardrail is disabled.")
537 return data
539 # Covers multimodal list content + Responses-API input/instructions.
540 inspection_messages: Final = _build_lakera_inspection_messages(data)
541 if not inspection_messages:
542 verbose_proxy_logger.warning("Lakera AI: not running guardrail. No inspectable text in data")
543 return data
545 new_messages, _ = self._filter_skipped_messages(
546 inspection_messages # pyright: ignore[reportArgumentType] # build_inspection_messages returns plain dicts, not typed message unions
547 )
548 if not new_messages:
549 verbose_proxy_logger.warning(
550 "Lakera AI: not running guardrail. All inspectable text was excluded by "
551 "skip_system_message_in_guardrail/skip_tool_message_in_guardrail"
552 )
553 return data
555 # Mask-in-place can only preserve non-text parts (images, audio) when
556 # the original content is a plain string, and can only merge a
557 # redacted result back into data["messages"] by position when
558 # messages and input aren't both present at once (build_inspection_messages
559 # flattens both into one list, so a position could mean either).
560 # Degrade to block-on-detect in either case. Skip-flag-excluded and
561 # no-text messages, and messages carrying fields beyond role/content
562 # (tool_call_id, name, tool_calls, cache_control), are otherwise
563 # handled safely by _apply_redacted_messages_back_preserving_fields's
564 # scope-index merge, which never touches a message outside the scope
565 # it actually redacted instead of reconstructing the list from scratch.
566 is_multimodal_input: Final = (
567 has_non_string_content(data)
568 or _has_combined_messages_and_input(data)
569 or _has_responses_instructions(self, data)
570 )
572 #########################################################
573 ########## 1. Make the Lakera AI v2 guard API request ##########
574 #########################################################
575 lakera_guardrail_response, masked_entity_count = await self.call_v2_guard(
576 messages=new_messages,
577 request_data=data,
578 event_type=GuardrailEventHooks.pre_call,
579 )
581 #########################################################
582 ########## 2. Handle flagged content ##########
583 #########################################################
584 if lakera_guardrail_response.get("flagged") is True:
585 is_pii_only_violation: Final = self._is_only_pii_violation(lakera_guardrail_response)
586 # PII-only violations get masked in place regardless of on_flagged: there's
587 # no reason to expose raw PII to satisfy an advisory note, and masking is
588 # strictly safer than either blocking or appending an advisory message next
589 # to unredacted PII.
590 if is_pii_only_violation and not is_multimodal_input:
591 redacted_messages: Final = self._mask_pii_in_messages(
592 messages=new_messages,
593 lakera_response=lakera_guardrail_response,
594 masked_entity_count=masked_entity_count,
595 )
596 _apply_redacted_messages_back_preserving_fields(self, data, redacted_messages)
597 verbose_proxy_logger.debug("Lakera AI: Masked PII in messages instead of blocking request")
598 elif self.on_flagged == "inject_system_message":
599 if _breakdown_has_pii_violation(lakera_guardrail_response) and is_multimodal_input:
600 # There's PII in the mix and nothing here can be safely masked,
601 # so an advisory note next to this raw, unredacted PII would be
602 # no safer than a note next to nothing. Degrade to blocking
603 # instead, same as this on_flagged setting already does when
604 # the advisory itself has no field it can be delivered into.
605 raise self._get_http_exception_for_blocked_guardrail(lakera_guardrail_response)
606 masked_pii_before_advisory: Final = _breakdown_has_pii_violation(lakera_guardrail_response)
607 if masked_pii_before_advisory:
608 # A mixed violation (PII plus something else, e.g. prompt
609 # injection): mask whatever Lakera returned location data for
610 # before advising about what remains, so the advisory is never
611 # shown next to raw PII that could have been redacted.
612 mixed_redacted_messages: Final = self._mask_pii_in_messages(
613 messages=new_messages,
614 lakera_response=lakera_guardrail_response,
615 masked_entity_count=masked_entity_count,
616 )
617 _apply_redacted_messages_back_preserving_fields(self, data, mixed_redacted_messages)
618 advisory_delivered: Final = self.inject_advisory_message(
619 data, self._build_advisory_message(lakera_guardrail_response)
620 )
621 if advisory_delivered:
622 verbose_proxy_logger.warning(
623 "Lakera Guardrail: Advisory mode - violation detected, %sappended advisory system message",
624 "masked PII and " if masked_pii_before_advisory else "",
625 )
626 else:
627 # Structured Responses-API input (a list, not a plain string)
628 # has no field this can safely append into -- degrade to
629 # blocking rather than silently letting the flagged request
630 # through with no advisory ever reaching the model.
631 raise self._get_http_exception_for_blocked_guardrail(lakera_guardrail_response)
632 elif self.on_flagged == "monitor":
633 # Monitor means "don't block", not "don't redact": until the mask
634 # branch above started skipping shapes it can't write back to, a
635 # PII-only violation was masked whatever on_flagged said.
636 masked_in_place: Final = is_pii_only_violation and self._mask_unwritable_instructions_pii_in_place(
637 data=data,
638 inspected_messages=new_messages,
639 lakera_response=lakera_guardrail_response,
640 masked_entity_count=masked_entity_count,
641 )
642 if masked_in_place:
643 verbose_proxy_logger.warning(
644 "Lakera Guardrail: Monitoring mode - PII detected, masked in place and allowing request"
645 )
646 elif is_pii_only_violation:
647 verbose_proxy_logger.error(
648 "Lakera Guardrail: Monitoring mode - PII detected but NOT masked, forwarding unredacted "
649 "content to the model (reason: %s)",
650 _unmaskable_reason(self, data, lakera_guardrail_response),
651 )
652 else:
653 verbose_proxy_logger.warning(
654 "Lakera Guardrail: Monitoring mode - violation detected but allowing request"
655 )
656 elif self.on_flagged == "block":
657 # Either non-PII violations, or PII on multimodal input
658 # (which cannot be masked in place without dropping
659 # image/audio parts) — raise the standard block error.
660 raise self._get_http_exception_for_blocked_guardrail(lakera_guardrail_response)
662 #########################################################
663 ########## 3. Add the guardrail to the applied guardrails header ##########
664 #########################################################
665 add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
667 return data
669 async def async_moderation_hook(
670 self,
671 data: dict,
672 user_api_key_dict: UserAPIKeyAuth,
673 call_type: CallTypesLiteral,
674 ):
675 from litellm.proxy.common_utils.callback_utils import (
676 add_guardrail_to_applied_guardrails_header,
677 )
679 event_type: Final[GuardrailEventHooks] = GuardrailEventHooks.during_call
680 if self.should_run_guardrail(data=data, event_type=event_type) is not True:
681 return
683 # Covers multimodal list content + Responses-API input/instructions.
684 inspection_messages: Final = _build_lakera_inspection_messages(data)
685 if not inspection_messages:
686 verbose_proxy_logger.warning("Lakera AI: not running guardrail. No inspectable text in data")
687 return
689 new_messages, _ = self._filter_skipped_messages(
690 inspection_messages # pyright: ignore[reportArgumentType] # build_inspection_messages returns plain dicts, not typed message unions
691 )
692 if not new_messages:
693 verbose_proxy_logger.warning(
694 "Lakera AI: not running guardrail. All inspectable text was excluded by "
695 "skip_system_message_in_guardrail/skip_tool_message_in_guardrail"
696 )
697 return
699 #########################################################
700 ########## 1. Make the Lakera AI v2 guard API request ##########
701 #########################################################
702 lakera_guardrail_response, _ = await self.call_v2_guard(
703 messages=new_messages,
704 request_data=data,
705 event_type=GuardrailEventHooks.during_call,
706 )
708 #########################################################
709 ########## 2. Handle flagged content ##########
710 #########################################################
711 if lakera_guardrail_response.get("flagged") is True:
712 # during_call runs concurrently with the LLM dispatch (see
713 # ProxyLogging.during_call_hook / common_request_processing.py), with
714 # no pre-call barrier: in the common path, the provider call already
715 # binds its messages kwarg before this coroutine gets a chance to run,
716 # let alone before the masking helper's own network round trip
717 # completes. Unlike async_pre_call_hook, mask-in-place here can never
718 # reliably reach the outgoing request, so PII is never masked in this
719 # hook -- only blocked (which still works, since raising here blocks
720 # the response from reaching the caller regardless of dispatch timing)
721 # or, for non-PII violations, logged and allowed same as monitor mode.
722 if self.on_flagged == "inject_system_message":
723 if _breakdown_has_pii_violation(lakera_guardrail_response):
724 raise self._get_http_exception_for_blocked_guardrail(lakera_guardrail_response)
725 verbose_proxy_logger.warning(
726 "Lakera Guardrail: Advisory mode has no effect during during_call; "
727 "violation detected but allowing request"
728 )
729 elif self.on_flagged == "monitor":
730 verbose_proxy_logger.warning(
731 "Lakera Guardrail: Monitoring mode - violation detected but allowing request"
732 )
733 elif self.on_flagged == "block":
734 raise self._get_http_exception_for_blocked_guardrail(lakera_guardrail_response)
736 #########################################################
737 ########## 3. Add the guardrail to the applied guardrails header ##########
738 #########################################################
739 add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
741 return data
743 async def async_post_call_success_hook(
744 self,
745 data: dict,
746 user_api_key_dict: UserAPIKeyAuth,
747 response,
748 ):
749 """
750 Post-call hook for Lakera guardrail.
751 """
752 from litellm.proxy.common_utils.callback_utils import (
753 add_guardrail_to_applied_guardrails_header,
754 )
756 event_type: Final[GuardrailEventHooks] = GuardrailEventHooks.post_call
757 if self.should_run_guardrail(data=data, event_type=event_type) is not True:
758 return response
760 messages_or_none: Final[list[AllMessageValues] | None] = data.get("messages")
761 original_messages, _ = self._filter_skipped_messages(messages_or_none or [])
763 # Extract assistant messages from the response, keeping only role/content.
764 # Track choice indices so we write masked content back to the correct choice
765 # when some choices have null content (e.g. tool-call-only).
766 response_messages: Final[list[AllMessageValues]] = []
767 choice_indices: Final[list[int]] = []
768 response_dict: Final = response.model_dump() if hasattr(response, "model_dump") else {}
769 for i, choice in enumerate(response_dict.get("choices", [])):
770 msg = choice.get("message")
771 if not msg:
772 continue
773 role = msg.get("role")
774 content = msg.get("content")
775 if role and content:
776 response_messages.append({"role": role, "content": content})
777 choice_indices.append(i)
779 # Use a copy of original_messages so _mask_pii_in_messages does not mutate data["messages"]
780 post_call_messages: Final = list(copy.deepcopy(original_messages)) + response_messages # mutable-ok: needs list
782 # Call Lakera guardrail
783 lakera_guardrail_response, _ = await self.call_v2_guard(
784 messages=post_call_messages,
785 request_data=data,
786 event_type=GuardrailEventHooks.post_call,
787 )
789 # Handle flagged content
790 if lakera_guardrail_response.get("flagged") is True:
791 # If only PII violations exist, mask the PII in the response and allow
792 if self._is_only_pii_violation(lakera_guardrail_response):
793 masked_entity_count: Final[dict[str, int]] = {}
794 masked_messages: Final = self._mask_pii_in_messages(
795 messages=post_call_messages,
796 lakera_response=lakera_guardrail_response,
797 masked_entity_count=masked_entity_count,
798 )
799 assistant_messages: Final = masked_messages[len(original_messages) :]
800 for idx, msg in enumerate(assistant_messages):
801 if idx < len(choice_indices):
802 choice_idx = choice_indices[idx]
803 response_dict["choices"][choice_idx]["message"]["content"] = msg.get("content", "")
804 add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
805 return ModelResponse(**response_dict)
807 # inject_system_message has nothing left to inject into once a response
808 # already exists, so it is treated the same as monitor: log and allow.
809 if self.on_flagged in ("monitor", "inject_system_message"):
810 verbose_proxy_logger.warning(
811 "Lakera Guardrail: Post-call violation detected (on_flagged=%s) - allowing response",
812 self.on_flagged,
813 )
814 elif self.on_flagged == "block":
815 raise self._get_http_exception_for_blocked_guardrail(lakera_guardrail_response)
817 # Record applied guardrail
818 add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
820 return response
822 def _is_only_pii_violation(self, lakera_response: LakeraAIResponse | None) -> bool:
823 """
824 Returns True if there are only PII violations in the response.
825 """
826 if not lakera_response:
827 return False
829 # Check breakdown field for detected violations
830 breakdown: Final = lakera_response.get("breakdown", []) or []
831 if not breakdown:
832 return False
834 has_violations = False
835 for item in breakdown:
836 if item.get("detected", False):
837 has_violations = True
838 detector_type = item.get("detector_type", "") or ""
839 if not detector_type.startswith("pii/"):
840 return False
842 # Return True only if there are violations and they are all PII
843 return has_violations
845 def _get_http_exception_for_blocked_guardrail(self, lakera_response: LakeraAIResponse | None) -> HTTPException:
846 """
847 Get the HTTP exception for a blocked guardrail, similar to Bedrock's implementation.
848 """
849 return HTTPException(
850 status_code=400,
851 detail={
852 "error": "Violated guardrail policy",
853 "lakera_guardrail_response": lakera_response,
854 },
855 )