from __future__ import annotations import asyncio import json import random from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn from urllib.parse import urlsplit import httpx from pydantic import BaseModel, TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm._version import version as litellm_version from litellm.exceptions import ( BadRequestError, GuardrailRaisedException, ModifyResponseException, Timeout, ) from litellm.integrations.custom_guardrail import ( CustomGuardrail, get_session_id_from_request_data, log_guardrail_information, ) from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.prompt_templates.factory import resolve_structured_messages from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) from litellm.types.guardrails import GuardrailEventHooks, Mode from litellm.types.proxy.guardrails.guardrail_hooks.straiker import ( STRAIKER_WEBHOOK_SCHEMA_VERSION, StraikerGuardrailConfigModel, StraikerWebhookApplication, StraikerWebhookContent, StraikerWebhookContext, StraikerWebhookEvent, StraikerWebhookIdentity, StraikerWebhookRequest, StraikerWebhookResponse, StraikerWebhookStream, StraikerWebhookUsage, ) from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel GUARDRAIL_NAME: Final = "straiker" DEFAULT_BLOCK_MESSAGE: Final = "Content violates policy" DEFAULT_API_BASE: Final = "https://api.prod.straiker.ai" DEFAULT_MAX_PAYLOAD_BYTES: Final = 524288 WEBHOOK_PATH: Final = "/api/v1/detect/webhook" RETRY_STATUS: Final = frozenset({408, 429, 500, 502, 503, 504}) UNREACHABLE_STATUS: Final = frozenset({502, 503, 504}) _APPLICATION_METADATA_KEYS: Final = frozenset({"agent_id", "app_name"}) _OPAQUE_METADATA_SCALAR_TYPES: Final = (str, int, float, bool) _JSON_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) @dataclass(frozen=True, slots=True) class _WebhookFailure: message: str is_unreachable: bool def _as_dict(value: object) -> dict: return value if isinstance(value, dict) else {} def _merged_metadata(request_data: dict) -> dict: return { **_as_dict(request_data.get("metadata")), **_as_dict(request_data.get("litellm_metadata")), } def _as_optional_str(value: object) -> str | None: return value if isinstance(value, str) and value else None def _build_webhook_metadata(request_data: dict, default_metadata: dict[str, str]) -> dict[str, object] | None: out: Final[dict[str, object]] = {} for key, value in _as_dict(request_data.get("metadata")).items(): if key in _APPLICATION_METADATA_KEYS or key.startswith("user_api"): continue if key == "session_id": continue if isinstance(value, _OPAQUE_METADATA_SCALAR_TYPES): out[key] = value out.update(default_metadata) return out or None def _extract_identity(request_data: dict) -> StraikerWebhookIdentity: meta: Final = _merged_metadata(request_data) return StraikerWebhookIdentity( litellm_key=_as_optional_str(meta.get("user_api_key_alias")) or _as_optional_str(meta.get("user_api_key_hash")) or _as_optional_str(meta.get("user_api_key_token")), litellm_team=_as_optional_str(meta.get("user_api_key_team_alias")) or _as_optional_str(meta.get("user_api_key_team_id")), litellm_user_id=_as_optional_str(meta.get("user_api_key_user_id")), litellm_user_email=_as_optional_str(meta.get("user_api_key_user_email")), litellm_org_id=_as_optional_str(meta.get("user_api_key_org_id")), end_user_id=_as_optional_str(meta.get("user_api_key_end_user_id")), ) def _resolve_provider(request_data: dict, model: str | None) -> str | None: litellm_params: Final = _as_dict(request_data.get("litellm_params")) custom_llm_provider: Final = request_data.get("custom_llm_provider") or litellm_params.get("custom_llm_provider") if custom_llm_provider: return custom_llm_provider if not model: return None try: _, provider, _, _ = get_llm_provider( model=model, api_base=request_data.get("api_base") or litellm_params.get("api_base"), api_key=request_data.get("api_key") or litellm_params.get("api_key"), ) except BadRequestError: return None return provider or None def _resolve_destination(request_data: dict) -> str | None: litellm_params: Final = _as_dict(request_data.get("litellm_params")) api_base: Final = request_data.get("api_base") or litellm_params.get("api_base") if not isinstance(api_base, str): return None try: return urlsplit(api_base).hostname except ValueError: return None def _route_has_translation(request_data: dict) -> bool: from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route from litellm.llms import load_guardrail_translation_mappings route: Final = _as_dict(request_data.get("litellm_metadata")).get("user_api_key_request_route") if not isinstance(route, str) or not route: return False mappings: Final = load_guardrail_translation_mappings() return any(call_type in mappings for call_type in get_call_types_for_route(route) or ()) def _request_structured_messages(request_data: dict) -> list[dict[str, Any]] | None: messages: Final = request_data.get("messages") if messages: return messages if isinstance(messages, list) else None if not _route_has_translation(request_data): return None return resolve_structured_messages(messages=None, request_kwargs=request_data) def _hook_name(value: object) -> str: return value.value if isinstance(value, GuardrailEventHooks) else str(value) def _configured_modes(event_hook: object) -> list[str] | None: if isinstance(event_hook, list): names = [_hook_name(v) for v in event_hook] elif isinstance(event_hook, (str, GuardrailEventHooks)): names = [_hook_name(event_hook)] elif isinstance(event_hook, Mode): default: Final = event_hook.default if isinstance(event_hook.default, list) else [event_hook.default] tags: Final = [v for value in event_hook.tags.values() for v in (value if isinstance(value, list) else [value])] names = [_hook_name(v) for v in (*default, *tags) if v is not None] else: return None return list(dict.fromkeys(names)) or None def _resolve_call_surface(logging_obj: LiteLLMLoggingObj | None, request_data: dict) -> str: call_type: Final = ( (getattr(logging_obj, "call_type", None) if logging_obj is not None else None) or request_data.get("call_type") or request_data.get("litellm_call_type") ) return call_type if isinstance(call_type, str) and call_type else "unknown" def _jsonable_dict(value: object) -> dict[str, object] | None: if isinstance(value, BaseModel): return _JSON_DICT_ADAPTER.validate_python(value.model_dump(mode="json", exclude_none=True)) if isinstance(value, dict): return _JSON_DICT_ADAPTER.validate_python(value) return None def _opaque_dict_list(value: object) -> list[dict[str, object]] | None: if not isinstance(value, list): return None items: Final = tuple(plain for item in value if (plain := _jsonable_dict(item)) is not None) return list(items) if items else None def _choice_terminal_reason(choice: object) -> str | None: if isinstance(choice, dict): return _as_optional_str(choice.get("finish_reason")) or _as_optional_str(choice.get("stop_reason")) return _as_optional_str(getattr(choice, "finish_reason", None)) or _as_optional_str( getattr(choice, "stop_reason", None) ) def _response_finish_reason(response: Any) -> str | None: if response is None: return None if isinstance(response, dict): top = _as_optional_str(response.get("finish_reason")) or _as_optional_str(response.get("stop_reason")) if top: return top choices = response.get("choices") if not isinstance(choices, list): return None for choice in choices: reason = _choice_terminal_reason(choice) if reason: return reason return None top = _as_optional_str(getattr(response, "finish_reason", None)) or _as_optional_str( getattr(response, "stop_reason", None) ) if top: return top choices = getattr(response, "choices", None) if not isinstance(choices, list): return None for choice in choices: reason = _choice_terminal_reason(choice) if reason: return reason return None def _as_optional_int(value: object) -> int | None: return value if isinstance(value, int) and not isinstance(value, bool) else None def _usage_token_count(usage: object, openai_key: str, anthropic_key: str) -> int | None: get: Final = usage.get if isinstance(usage, dict) else lambda key: getattr(usage, key, None) openai_count: Final = _as_optional_int(get(openai_key)) return openai_count if openai_count is not None else _as_optional_int(get(anthropic_key)) def _build_usage(response: object) -> StraikerWebhookUsage | None: usage: Final = response.get("usage") if isinstance(response, dict) else getattr(response, "usage", None) if usage is None: return None input_tokens: Final = _usage_token_count(usage, "prompt_tokens", "input_tokens") output_tokens: Final = _usage_token_count(usage, "completion_tokens", "output_tokens") if input_tokens is None and output_tokens is None: return None return StraikerWebhookUsage(input_tokens=input_tokens, output_tokens=output_tokens) def _is_streamed_request(request_data: dict) -> bool: if request_data.get("stream") is True: return True body: Final = _as_dict(_as_dict(request_data.get("proxy_server_request")).get("body")) return body.get("stream") is True class StraikerGuardrail(CustomGuardrail): @staticmethod def get_config_model() -> type[GuardrailConfigModel]: return StraikerGuardrailConfigModel @classmethod def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: return [ GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call, ] def __init__( self, api_key: str, api_base: str = DEFAULT_API_BASE, source: str = "LiteLLM Gateway", timeout: float = 5.0, max_retries: int = 2, initial_backoff: float = 0.1, max_backoff: float = 2.0, unreachable_fallback: Literal["fail_open", "fail_closed"] = "fail_closed", fail_on_error: bool = True, max_payload_bytes: int = DEFAULT_MAX_PAYLOAD_BYTES, custom_headers: dict[str, str] | None = None, metadata: dict[str, str] | None = None, verbose: bool = False, async_handler: httpx.AsyncClient | None = None, **kwargs: object, ) -> None: if not api_key: raise ValueError("api_key must be non-empty") if unreachable_fallback not in ("fail_open", "fail_closed"): raise ValueError(f"unreachable_fallback must be 'fail_open' or 'fail_closed'; got {unreachable_fallback!r}") self.api_key = api_key self.api_base = api_base.rstrip("/") self.source = source self.timeout = float(timeout) self.max_retries = max(0, int(max_retries)) self.initial_backoff = max(0.0, float(initial_backoff)) self.max_backoff = max(self.initial_backoff, float(max_backoff)) self.unreachable_fallback = unreachable_fallback self.fail_on_error = fail_on_error self.max_payload_bytes = int(max_payload_bytes) self.custom_headers = dict(custom_headers) if custom_headers else {} self.default_metadata = dict(metadata) if metadata else {} self.verbose = bool(verbose) self.streaming_end_of_stream_only = True self.streaming_buffer_until_moderated = True self.async_handler = async_handler or get_async_httpx_client( llm_provider=httpxSpecialProvider.GuardrailCallback, ) kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) super().__init__(**kwargs) self.configured_modes = _configured_modes(self.event_hook) def _webhook_url(self) -> str: return f"{self.api_base}{WEBHOOK_PATH}" def _headers(self) -> dict[str, str]: reserved: Final = {"authorization", "content-type", "x-straiker-webhook-format"} extra: Final = {k: v for k, v in self.custom_headers.items() if k.lower() not in reserved} return { "Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json", "X-Straiker-Webhook-Format": "litellm", **extra, } def _build_application(self, request_data: dict) -> StraikerWebhookApplication: meta: Final = _merged_metadata(request_data) agent_id: Final = _as_optional_str(meta.get("agent_id")) return StraikerWebhookApplication( source=agent_id or self.source, name=_as_optional_str(meta.get("app_name")), ) def _build_context( self, request_data: dict, model: str | None, logging_obj: LiteLLMLoggingObj | None, ) -> StraikerWebhookContext: return StraikerWebhookContext( call_surface=_resolve_call_surface(logging_obj, request_data), mode=self.configured_modes, model=model, model_provider=_resolve_provider(request_data, model), destination=_resolve_destination(request_data), session_id=get_session_id_from_request_data(request_data), litellm_call_id=getattr(logging_obj, "litellm_call_id", None) if logging_obj else None, litellm_trace_id=getattr(logging_obj, "litellm_trace_id", None) if logging_obj else None, litellm_version=litellm_version, ) def _build_envelope( self, *, inputs: GenericGuardrailAPIInputs, request_data: dict, input_type: Literal["request", "response"], logging_obj: LiteLLMLoggingObj | None, ) -> StraikerWebhookRequest: model: Final = inputs.get("model") or request_data.get("model") call_id: Final = getattr(logging_obj, "litellm_call_id", None) if logging_obj else None event_id: Final = f"{call_id or 'litellm'}:{input_type}" content: Final = StraikerWebhookContent( texts=list(inputs.get("texts") or []), images=list(inputs.get("images") or []), structured_messages=_opaque_dict_list(inputs.get("structured_messages")), tools=_opaque_dict_list(inputs.get("tools")), tool_calls=_opaque_dict_list(inputs.get("tool_calls")), ) if input_type == "request": event = StraikerWebhookEvent(type="pre_call", id=event_id) return StraikerWebhookRequest( event=event, request=content, context=self._build_context(request_data, model, logging_obj), identity=_extract_identity(request_data), application=self._build_application(request_data), metadata=_build_webhook_metadata(request_data, self.default_metadata), ) response_obj: Final = request_data.get("response") content.finish_reason = _response_finish_reason(response_obj) request_content: Final = StraikerWebhookContent( structured_messages=_opaque_dict_list(_request_structured_messages(request_data)), ) phase: Final[Literal["none", "assembled"]] = "assembled" if _is_streamed_request(request_data) else "none" event = StraikerWebhookEvent(type="post_call", id=event_id, stream=StraikerWebhookStream(phase=phase)) return StraikerWebhookRequest( event=event, request=request_content, response=content, context=self._build_context(request_data, model, logging_obj), identity=_extract_identity(request_data), application=self._build_application(request_data), usage=_build_usage(response_obj), metadata=_build_webhook_metadata(request_data, self.default_metadata), ) async def _post_webhook(self, payload: dict) -> tuple[StraikerWebhookResponse | None, _WebhookFailure | None]: try: body = json.dumps(payload).encode("utf-8") except (TypeError, ValueError, OverflowError) as error: return None, _WebhookFailure(f"request serialization failed: {error}", is_unreachable=False) body_bytes: Final = len(body) if body_bytes > self.max_payload_bytes: return None, _WebhookFailure( f"payload {body_bytes}B exceeds max_payload_bytes {self.max_payload_bytes}", is_unreachable=False, ) url: Final = self._webhook_url() headers: Final = self._headers() attempts: Final = self.max_retries + 1 last_failure: _WebhookFailure | None = None if self.verbose: verbose_proxy_logger.info( json.dumps( { "event": "straiker.webhook_request", "url": url, "bytes": body_bytes, "payload": payload, }, default=str, ) ) for attempt in range(attempts): try: resp = await self.async_handler.post(url, content=body, headers=headers, timeout=self.timeout) if resp.status_code == 200: try: body = resp.json() parsed = StraikerWebhookResponse.model_validate(body) except (ValidationError, json.JSONDecodeError) as ve: return None, _WebhookFailure(f"invalid response schema: {ve}", is_unreachable=False) if self.verbose: verbose_proxy_logger.info( json.dumps( { "event": "straiker.webhook_response", "status_code": resp.status_code, "body": body, }, default=str, ) ) return parsed, None last_failure = _WebhookFailure( f"HTTP {resp.status_code}: {resp.text[:200]}", is_unreachable=resp.status_code in UNREACHABLE_STATUS, ) if resp.status_code not in RETRY_STATUS: return None, last_failure except (httpx.RequestError, asyncio.TimeoutError, Timeout) as e: last_failure = _WebhookFailure(f"{type(e).__name__}: {e}", is_unreachable=True) except (json.JSONDecodeError, TypeError, ValueError) as e: return None, _WebhookFailure(f"{type(e).__name__}: {e}", is_unreachable=False) if attempt < attempts - 1: backoff = min(self.initial_backoff * (2**attempt), self.max_backoff) await asyncio.sleep(random.uniform(0, backoff)) return None, last_failure or _WebhookFailure("unknown error", is_unreachable=True) def _record( self, *, request_data: dict, logging_obj: LiteLLMLoggingObj | None, parsed: StraikerWebhookResponse, ) -> None: if not self.verbose: return response_obj: Final = request_data.get("response") hidden: Final = getattr(response_obj, "_hidden_params", None) if isinstance(hidden, dict): straiker_hidden: Final = hidden.setdefault("straiker", {}) if isinstance(straiker_hidden, dict): straiker_hidden.update({"action": parsed.action, "turn_id": parsed.turn_id}) def _fail( self, *, inputs: GenericGuardrailAPIInputs, request_data: dict, input_type: Literal["request", "response"], error: str, is_unreachable: bool, ) -> GenericGuardrailAPIInputs: fail_open: Final = (is_unreachable and self.unreachable_fallback == "fail_open") or not self.fail_on_error verbose_proxy_logger.error( json.dumps( { "event": "straiker.error", "input_type": input_type, "error": error, "fail_open": fail_open, }, default=str, ) ) if fail_open: return inputs self._block( request_data=request_data, input_type=input_type, message=f"Straiker detection unavailable: {error}", ) def _block( self, *, request_data: dict, input_type: Literal["request", "response"], message: str, blocked_content: bool = False, ) -> NoReturn: if input_type == "request": raise GuardrailRaisedException( guardrail_name=self.guardrail_name or GUARDRAIL_NAME, message=message, should_wrap_with_default_message=False, blocked_content=blocked_content, ) raise ModifyResponseException( message=message, model=request_data.get("model", "unknown") or "unknown", request_data=request_data, guardrail_name=self.guardrail_name or GUARDRAIL_NAME, original_response=request_data.get("response"), ) @staticmethod def _intervened_inputs( inputs: GenericGuardrailAPIInputs, parsed: StraikerWebhookResponse, ) -> GenericGuardrailAPIInputs: return_inputs: Final[GenericGuardrailAPIInputs] = {} return_inputs.update(inputs) if parsed.texts is not None: return_inputs["texts"] = parsed.texts return return_inputs @log_guardrail_information async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, request_data: dict, input_type: Literal["request", "response"], logging_obj: LiteLLMLoggingObj | None = None, ) -> GenericGuardrailAPIInputs: try: envelope: Final = self._build_envelope( inputs=inputs, request_data=request_data, input_type=input_type, logging_obj=logging_obj, ) payload: Final = envelope.model_dump(mode="json", exclude_none=True) except (ValidationError, TypeError, ValueError) as error: return self._fail( inputs=inputs, request_data=request_data, input_type=input_type, error=str(error), is_unreachable=False, ) parsed, failure = await self._post_webhook(payload) if failure is not None: return self._fail( inputs=inputs, request_data=request_data, input_type=input_type, error=failure.message, is_unreachable=failure.is_unreachable, ) if parsed is None: return self._fail( inputs=inputs, request_data=request_data, input_type=input_type, error="empty response from Straiker", is_unreachable=False, ) self._record(request_data=request_data, logging_obj=logging_obj, parsed=parsed) if parsed.schema_version is not None and parsed.schema_version != STRAIKER_WEBHOOK_SCHEMA_VERSION: verbose_proxy_logger.warning( json.dumps( { "event": "straiker.schema_drift", "expected": STRAIKER_WEBHOOK_SCHEMA_VERSION, "received": parsed.schema_version, } ) ) if parsed.action == "BLOCKED": self._block( request_data=request_data, input_type=input_type, message=parsed.blocked_reason or DEFAULT_BLOCK_MESSAGE, blocked_content=True, ) if parsed.action == "GUARDRAIL_INTERVENED": is_streamed_response: Final = input_type == "response" and _is_streamed_request(request_data) if parsed.texts is None or is_streamed_response: self._block( request_data=request_data, input_type=input_type, message=parsed.blocked_reason or DEFAULT_BLOCK_MESSAGE, blocked_content=True, ) return self._intervened_inputs(inputs, parsed) return inputs