Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py: 23%
134 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1"""Ovalix guardrail integration: pre- and post-call checks via the Tracker service.
3Use Ovalix Guardrails for your LLM calls. Supports pre_call (user input) and
4post_call (model output) checkpoints with optional correction/blocking.
5"""
7import datetime
8import hashlib
9import os
10from typing import TYPE_CHECKING, Any, Final, Literal
12import httpx
13from typing_extensions import ReadOnly, TypedDict, Unpack
15from litellm._logging import verbose_proxy_logger
16from litellm.exceptions import GuardrailRaisedException
17from litellm.integrations.custom_guardrail import (
18 CustomGuardrail,
19 log_guardrail_information,
20)
21from litellm.llms.custom_httpx.http_handler import (
22 get_async_httpx_client,
23 httpxSpecialProvider,
24)
25from litellm.types.guardrails import GuardrailEventHooks
26from litellm.types.utils import GenericGuardrailAPIInputs
28if TYPE_CHECKING: 28 ↛ 29line 28 didn't jump to line 29 because the condition on line 28 was never true
29 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
30 from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
33BLOCKED_BY_OVALIX_FALLBACK_MESSAGE: Final = "This message was blocked by Ovalix"
34BLOCKED_ACTION_TYPE: Final = "block"
37class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object):
38 """Base-class constructor options this guardrail forwards untouched to CustomGuardrail."""
40 supported_event_hooks: ReadOnly[list[GuardrailEventHooks]]
43class OvalixGuardrailMissingSecrets(Exception):
44 """Raised when required Ovalix config (API base, key, application/checkpoint IDs) is missing."""
47class OvalixGuardrailBlockedException(GuardrailRaisedException):
48 """
49 Raised when Ovalix blocks a message. Sets status_code=400 so the proxy
50 returns 400 and HTTP clients do not retry (they retry on 5xx).
51 """
53 status_code = 400
55 def __init__(
56 self,
57 guardrail_name: str | None = None,
58 message: str = "",
59 should_wrap_with_default_message: bool = True,
60 ):
61 super().__init__(
62 guardrail_name=guardrail_name,
63 message=message,
64 should_wrap_with_default_message=should_wrap_with_default_message,
65 blocked_content=True,
66 )
69class OvalixGuardrail(CustomGuardrail):
70 """
71 Ovalix guardrail: pre-prompt (pre_call) and post-prompt (post_call) checks
72 via the Tracker service, with application and checkpoint resolution from the
73 Monolith backend.
74 """
76 @classmethod
77 def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
78 return [
79 GuardrailEventHooks.pre_call,
80 GuardrailEventHooks.post_call,
81 ]
83 def __init__(
84 self,
85 tracker_api_base: str | None = None,
86 tracker_api_key: str | None = None,
87 application_id: str | None = None,
88 pre_checkpoint_id: str | None = None,
89 post_checkpoint_id: str | None = None,
90 **kwargs: Unpack[_CustomGuardrailOptions],
91 ):
92 self._tracker_api_base = tracker_api_base or os.environ.get("OVALIX_TRACKER_API_BASE")
93 self._tracker_api_key = tracker_api_key or os.environ.get("OVALIX_TRACKER_API_KEY")
94 self._application_id = application_id or os.environ.get("OVALIX_APPLICATION_ID")
95 self._pre_checkpoint_id = pre_checkpoint_id or os.environ.get("OVALIX_PRE_CHECKPOINT_ID")
96 self._post_checkpoint_id = post_checkpoint_id or os.environ.get("OVALIX_POST_CHECKPOINT_ID")
98 supported_event_hooks: Final = kwargs.get("supported_event_hooks", [])
100 self._validate_config(supported_event_hooks)
102 self._tracker_headers = httpx.Headers(
103 {
104 "Authorization": f"Bearer {self._tracker_api_key}",
105 "Content-Type": "application/json",
106 },
107 encoding="utf-8",
108 )
110 self._async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
112 forwarded: Final[_CustomGuardrailOptions] = {**kwargs, "supported_event_hooks": supported_event_hooks}
113 super().__init__(**forwarded)
114 verbose_proxy_logger.debug(
115 "Ovalix Guardrail initialized: tracker=%s, application_id=%s, pre_checkpoint_id=%s, post_checkpoint_id=%s",
116 self._tracker_api_base,
117 self._application_id,
118 self._pre_checkpoint_id,
119 self._post_checkpoint_id,
120 )
122 def _validate_config(self, supported_event_hooks: list[GuardrailEventHooks]) -> None:
123 """Ensure required secrets and checkpoint IDs are set; auto-add hooks when IDs are present."""
124 errors: Final[list[str]] = []
126 if not self._tracker_api_base:
127 errors.append("Tracker API base, set OVALIX_TRACKER_API_BASE or pass tracker_api_base")
128 if not self._tracker_api_key:
129 errors.append("Tracker API key, set OVALIX_TRACKER_API_KEY or pass tracker_api_key")
130 if not self._application_id:
131 errors.append("Application ID, set OVALIX_APPLICATION_ID or pass application_id")
132 if not self._pre_checkpoint_id and GuardrailEventHooks.pre_call in supported_event_hooks:
133 errors.append("Pre-checkpoint ID, set OVALIX_PRE_CHECKPOINT_ID or pass pre_checkpoint_id")
134 if not self._post_checkpoint_id and GuardrailEventHooks.post_call in supported_event_hooks:
135 errors.append("Post-checkpoint ID, set OVALIX_POST_CHECKPOINT_ID or pass post_checkpoint_id")
136 if not self._pre_checkpoint_id and not self._post_checkpoint_id:
137 errors.append(
138 "Pre-checkpoint ID or Post-checkpoint ID, set OVALIX_PRE_CHECKPOINT_ID or OVALIX_POST_CHECKPOINT_ID or pass pre_checkpoint_id or post_checkpoint_id"
139 )
141 if errors:
142 raise OvalixGuardrailMissingSecrets("Missing Ovalix guardrail configuration errors: " + ". ".join(errors))
144 # auto-add hooks when checkpoint IDs are present
145 if self._pre_checkpoint_id and GuardrailEventHooks.pre_call not in supported_event_hooks:
146 supported_event_hooks.append(GuardrailEventHooks.pre_call)
147 if self._post_checkpoint_id and GuardrailEventHooks.post_call not in supported_event_hooks:
148 supported_event_hooks.append(GuardrailEventHooks.post_call)
150 def _get_actor(self, data: dict) -> str:
151 """Return a stable actor identifier from request metadata (e.g. user email or id)."""
152 metadata: Final = data.get("metadata") or data.get("litellm_metadata") or {}
153 if metadata.get("user_api_key_user_email"):
154 return metadata["user_api_key_user_email"]
155 if metadata.get("user_api_key_user_id"):
156 return metadata["user_api_key_user_id"]
157 return "unknown"
159 def _get_tracker_actor_id(self, data: dict) -> str:
160 """Normalize the actor string into a short, stable id for Tracker API payloads."""
161 # NOTE: this hash is purely for normalization — it collapses an arbitrary actor
162 # string (email, user id, or "unknown") into a compact, fixed-length, consistent
163 # key. It is not a privacy/security measure and the actor value is not sensitive,
164 # so a plain SHA-256 (truncated) is sufficient; no salting/KDF is needed here.
165 actor_id: Final = self._get_actor(data).encode()
166 normalized_actor_id: Final = hashlib.sha256(actor_id).hexdigest()[:8]
167 return normalized_actor_id
169 def _get_session_id(self, data: dict) -> str:
170 """Return a unique identifier for the chat/session (actor + date + application_id)."""
171 actor_hash: Final = self._get_tracker_actor_id(data)
172 today: Final = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d")
173 return f"{actor_hash}_{today}_{self._application_id}"
175 async def _call_checkpoint(
176 self,
177 content: str,
178 checkpoint_id: str,
179 actor: str,
180 session_id: str,
181 ) -> dict[str, Any]:
182 """Call the Ovalix Tracker checkpoint API and return the JSON response."""
183 application_id: Final = self._application_id
184 if not application_id or not checkpoint_id:
185 raise ValueError("Ovalix: application_id or checkpoint_id not resolved")
187 url: Final = f"{self._tracker_api_base}/tracking/custom_application/checkpoint"
188 headers: Final = dict(self._tracker_headers)
189 payload: Final = {
190 "application_id": application_id,
191 "checkpoint_id": checkpoint_id,
192 "actor": actor,
193 "session_id": session_id,
194 "data_type": "TEXT",
195 "data": {"content": content},
196 }
197 response: Final = await self._async_handler.post(url, headers=headers, json=payload)
198 response.raise_for_status()
199 return response.json()
201 @log_guardrail_information
202 async def apply_guardrail(
203 self,
204 inputs: GenericGuardrailAPIInputs,
205 request_data: dict,
206 input_type: Literal["request", "response"],
207 logging_obj: "LiteLLMLoggingObj | None" = None,
208 ) -> GenericGuardrailAPIInputs:
209 """
210 Apply Ovalix guardrail to the given inputs (request or response text).
212 Used by the unified guardrail flow and the /apply_guardrail API.
213 For "request", uses the pre-checkpoint; for "response", uses the post-checkpoint.
215 Args:
216 inputs: Guardrail API inputs (e.g. texts to check).
217 request_data: Full request payload (messages, metadata, response).
218 input_type: "request" (pre_call) or "response" (post_call).
219 logging_obj: Optional logging context.
221 Returns:
222 Updated inputs (e.g. with replaced/corrected texts, or unchanged).
223 """
224 if not self._pre_checkpoint_id and not self._post_checkpoint_id:
225 return inputs
227 tracker_actor_id: Final = self._get_tracker_actor_id(request_data)
228 session_id: Final = self._get_session_id(request_data)
229 texts: Final = inputs.get("texts") or []
230 if not texts or not isinstance(texts, list):
231 return inputs
233 if input_type == "response":
234 if not self._post_checkpoint_id:
235 return inputs
236 corrected_llm_responses: Final = await self._generate_post_guardrail_llm_texts(
237 texts, tracker_actor_id, session_id, self._post_checkpoint_id
238 )
239 return {**inputs, "texts": corrected_llm_responses}
241 if self._pre_checkpoint_id:
242 post_guardrail_texts: Final = await self._generate_post_guardrail_llm_texts(
243 texts, tracker_actor_id, session_id, self._pre_checkpoint_id
244 )
245 return {**inputs, "texts": post_guardrail_texts}
246 return inputs
248 async def _generate_post_guardrail_llm_texts(
249 self, texts: list[str], actor: str, session_id: str, checkpoint_id: str
250 ) -> list[str]:
251 """Generate post-guardrail LLM responses for the given LLM responses."""
252 post_guardrail_texts: Final[list[str]] = []
254 is_first_response = True
255 for llm_response in reversed(texts):
256 try:
257 resp = await self._call_checkpoint(llm_response, checkpoint_id, actor, session_id)
258 except Exception as e:
259 verbose_proxy_logger.exception("Ovalix apply_guardrail checkpoint call failed: %s", e)
260 raise GuardrailRaisedException(
261 guardrail_name=self.guardrail_name,
262 message=f"Ovalix guardrail error: {e}",
263 should_wrap_with_default_message=False,
264 ) from e
266 action_type = (resp.get("action_type") or "").lower()
267 blocking_message = self._get_trackers_corrected_message(resp) or BLOCKED_BY_OVALIX_FALLBACK_MESSAGE
268 if action_type == BLOCKED_ACTION_TYPE and is_first_response:
269 self._block_current_message(blocking_message)
270 elif action_type == BLOCKED_ACTION_TYPE:
271 post_guardrail_texts.insert(0, blocking_message)
272 else:
273 corrected_text = self._get_trackers_corrected_message(resp) or llm_response
274 post_guardrail_texts.insert(0, corrected_text)
275 is_first_response = False
276 return post_guardrail_texts
278 def _block_current_message(self, blocking_message: str) -> None:
279 """Raise OvalixGuardrailBlockedException with the given message (no default wrapper)."""
280 raise OvalixGuardrailBlockedException(
281 guardrail_name=self.guardrail_name,
282 message=blocking_message,
283 should_wrap_with_default_message=False,
284 )
286 def _get_trackers_corrected_message(self, resp: dict) -> str | None:
287 """Extract corrected/blocking message content from Tracker checkpoint response."""
288 modified: Final = resp.get("modified_data")
289 if isinstance(modified, dict) and "content" in modified:
290 return modified["content"]
291 return None
293 @staticmethod
294 def get_config_model() -> type["GuardrailConfigModel"] | None:
295 from litellm.types.proxy.guardrails.guardrail_hooks.ovalix import (
296 OvalixGuardrailConfigModel,
297 )
299 return OvalixGuardrailConfigModel