From f759c75466f0475362a5016ad4cd03c1ecbbd515 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Fri, 17 Jul 2026 20:31:29 -0700 Subject: [PATCH] feat: add Straiker guardrail integration (#33781) * feat: add Straiker guardrail integration Implements LLM security guardrails via Straiker with prompt and response inspection, multi-mode execution (pre_call, post_call), and configurable blocking or redaction of flagged content across providers, streaming, images, and tool calls. * fix(guardrails): harden straiker source attribution and error-path consistency Use the operator-configured source for Straiker application attribution instead of a caller-supplied agent_id metadata value, so a caller cannot spoof which application a detection is attributed to. Make _fail reuse _block so a post_call error raises ModifyResponseException like a deliberate post_call block rather than GuardrailRaisedException, and type the blocking helper as NoReturn so the type checker enforces that execution never falls through the BLOCKED branch. Serialize the webhook payload once and send it as raw content to avoid re-serializing on the size check and on every retry. * fix(guardrails): read straiker config and metadata from all supported shapes Handle a dict optional_params in _get_config_value so nested guardrail settings loaded from YAML or the DB (timeout, unreachable_fallback, and the rest) are applied instead of silently falling back to defaults; previously only attribute-style access was supported. Build the webhook metadata bag from the merged metadata so client tags stored under litellm_metadata on routes like /v1/messages reach Straiker the same way identity and application fields already do, and widen the internal-key skip prefix to user_api so proxy-injected budget values are not forwarded. * fix(guardrails): fail safe on straiker interventions without redactions Block instead of passing content through when Straiker returns GUARDRAIL_INTERVENED without replacement texts, so a positive intervention verdict can never silently forward the original flagged content. Fix the streamed-request detection to read the request body from proxy_server_request.body, where the proxy stores it, instead of a top-level body key that is never populated; the previous fallback was dead, so a streamed response whose stream flag was not lifted to the top level would have been redacted rather than blocked while buffering replayed the original chunks. * revert(guardrails): restore straiker caller agent_id application attribution Restore the original behavior where a request-scoped agent_id in metadata sets the Straiker application source, falling back to the configured source. This is the integration's intended per-application attribution; litellm already resolves a key-owned agent_id ahead of any caller-supplied value, so a configured key cannot be spoofed. * revert(guardrails): restore straiker webhook metadata scoping Restore the original behavior where the Straiker webhook metadata bag is built from request-scoped metadata only. Forwarding litellm_metadata was a scope change to what the integration sends to Straiker; keep the author's intended scoping. * fix(guardrails): keep proxy key material out of straiker webhook metadata Widen the internal-key skip prefix from user_api_key_ to user_api so the proxy-injected user_api_key hash and user_api_end_user_max_budget are not copied into the Straiker webhook metadata bag. The narrower prefix missed the bare user_api_key name, leaking the hashed key to the vendor. Keeps the request-scoped metadata source unchanged. --------- Co-authored-by: cs-mehta --- .../guardrail_hooks/straiker/__init__.py | 71 ++ .../guardrail_hooks/straiker/straiker.py | 541 +++++++++++++ litellm/types/guardrails.py | 1 + .../guardrails/guardrail_hooks/straiker.py | 169 ++++ .../guardrail_hooks/test_straiker.py | 733 ++++++++++++++++++ .../public/assets/logos/straiker.svg | 9 + .../_components/guardrail_garden_configs.ts | 6 + .../_components/guardrail_garden_data.ts | 10 + .../_components/guardrail_info_helpers.tsx | 1 + 9 files changed, 1541 insertions(+) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/straiker/__init__.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py create mode 100644 litellm/types/proxy/guardrails/guardrail_hooks/straiker.py create mode 100644 tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py create mode 100644 ui/litellm-dashboard/public/assets/logos/straiker.svg diff --git a/litellm/proxy/guardrails/guardrail_hooks/straiker/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/straiker/__init__.py new file mode 100644 index 00000000000..ba4c712764e --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/straiker/__init__.py @@ -0,0 +1,71 @@ +from typing import TYPE_CHECKING + +import litellm +from litellm.types.guardrails import SupportedGuardrailIntegrations + +from .straiker import StraikerGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + +_OPTIONAL_INIT_FIELDS = ( + "timeout", + "max_retries", + "initial_backoff", + "max_backoff", + "unreachable_fallback", + "fail_on_error", + "max_payload_bytes", + "custom_headers", + "metadata", + "verbose", +) + + +def _get_config_value(litellm_params: "LitellmParams", optional_params: object, attribute_name: str) -> object: + if optional_params is not None: + if isinstance(optional_params, dict): + value = optional_params.get(attribute_name) + else: + value = getattr(optional_params, attribute_name, None) + if value is not None: + return value + return getattr(litellm_params, attribute_name, None) + + +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): + optional_params = getattr(litellm_params, "optional_params", None) + api_key = litellm_params.api_key + if not api_key: + raise ValueError("api_key is required for straiker") + + api_base = litellm_params.api_base or "https://api.prod.straiker.ai" + default_app = getattr(litellm_params, "default_app", None) or getattr(litellm_params, "source", None) + source = default_app if isinstance(default_app, str) and default_app else "LiteLLM Gateway" + kwargs: dict[str, object] = { + field: value + for field in _OPTIONAL_INIT_FIELDS + for value in [_get_config_value(litellm_params, optional_params, field)] + if value is not None + } + _callback = StraikerGuardrail( + api_key=api_key, + api_base=api_base if isinstance(api_base, str) else "https://api.prod.straiker.ai", + source=source, + guardrail_name=guardrail.get("guardrail_name", "straiker"), + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + **kwargs, + ) + + litellm.logging_callback_manager.add_litellm_callback(_callback) + return _callback + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.STRAIKER.value: initialize_guardrail, +} + +guardrail_class_registry = { + SupportedGuardrailIntegrations.STRAIKER.value: StraikerGuardrail, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py b/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py new file mode 100644 index 00000000000..5c9f93fc2cd --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py @@ -0,0 +1,541 @@ +from __future__ import annotations + +import asyncio +import json +import random +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Literal, NoReturn +from urllib.parse import urlsplit + +import httpx +from pydantic import 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.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + httpxSpecialProvider, +) +from litellm.types.guardrails import GuardrailEventHooks +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, Usage + +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 = "straiker" +DEFAULT_BLOCK_MESSAGE = "Content violates policy" +DEFAULT_API_BASE = "https://api.prod.straiker.ai" +DEFAULT_MAX_PAYLOAD_BYTES = 524288 +WEBHOOK_PATH = "/api/v1/detect/webhook" +RETRY_STATUS = frozenset({408, 429, 500, 502, 503, 504}) +UNREACHABLE_STATUS = frozenset({502, 503, 504}) +_APPLICATION_METADATA_KEYS = frozenset({"agent_id", "app_name"}) +_OPAQUE_METADATA_SCALAR_TYPES = (str, int, float, bool) + + +@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: 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 = _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 = _as_dict(request_data.get("litellm_params")) + custom_llm_provider = 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 = _as_dict(request_data.get("litellm_params")) + api_base = 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 _resolve_call_surface(logging_obj: LiteLLMLoggingObj | None, request_data: dict) -> str: + call_type = ( + (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 _response_finish_reason(response: Any) -> str | None: + choices = getattr(response, "choices", None) + if not isinstance(choices, list): + return None + for choice in choices: + reason = getattr(choice, "finish_reason", None) + if isinstance(reason, str) and reason: + return reason + return None + + +def _build_usage(response: object) -> StraikerWebhookUsage | None: + usage = getattr(response, "usage", None) + if not isinstance(usage, Usage): + return None + input_tokens = usage.prompt_tokens + output_tokens = usage.completion_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 = _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) + + def _webhook_url(self) -> str: + return f"{self.api_base}{WEBHOOK_PATH}" + + def _headers(self) -> dict[str, str]: + reserved = {"authorization", "content-type", "x-straiker-webhook-format"} + extra = {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 = _merged_metadata(request_data) + agent_id = _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), + 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 = inputs.get("model") or request_data.get("model") + call_id = getattr(logging_obj, "litellm_call_id", None) if logging_obj else None + event_id = f"{call_id or 'litellm'}:{input_type}" + + content = StraikerWebhookContent( + texts=list(inputs.get("texts") or []), + images=list(inputs.get("images") or []), + structured_messages=inputs.get("structured_messages"), + tools=inputs.get("tools"), + tool_calls=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 = request_data.get("response") + content.finish_reason = _response_finish_reason(response_obj) + original_messages = request_data.get("messages") + request_content = StraikerWebhookContent( + structured_messages=original_messages if isinstance(original_messages, list) else None, + ) + phase: 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 = 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 = self._webhook_url() + headers = self._headers() + attempts = 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 = request_data.get("response") + hidden = getattr(response_obj, "_hidden_params", None) + if isinstance(hidden, dict): + straiker_hidden = 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 = (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, + ) -> NoReturn: + if input_type == "request": + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name or GUARDRAIL_NAME, + message=message, + should_wrap_with_default_message=False, + ) + 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: 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 = self._build_envelope( + inputs=inputs, + request_data=request_data, + input_type=input_type, + logging_obj=logging_obj, + ) + payload = 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, + ) + if parsed.action == "GUARDRAIL_INTERVENED": + is_streamed_response = 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, + ) + return self._intervened_inputs(inputs, parsed) + return inputs diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 86e69467dbf..5b611971154 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -131,6 +131,7 @@ class SupportedGuardrailIntegrations(Enum): SINGULR = "singulr" HEADROOM = "headroom" COMPRESR = "compresr" + STRAIKER = "straiker" class Role(Enum): diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/straiker.py b/litellm/types/proxy/guardrails/guardrail_hooks/straiker.py new file mode 100644 index 00000000000..b4375237917 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/straiker.py @@ -0,0 +1,169 @@ +from __future__ import annotations + +from typing import Literal + +from pydantic import BaseModel, ConfigDict, Field + +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk +from litellm.types.utils import ChatCompletionMessageToolCall + +from .base import GuardrailConfigModel + +StraikerWebhookEventType = Literal["pre_call", "post_call"] +StraikerWebhookStreamPhase = Literal["none", "assembled"] +StraikerWebhookAction = Literal["NONE", "BLOCKED", "GUARDRAIL_INTERVENED"] + +STRAIKER_WEBHOOK_SCHEMA_VERSION = "1" + + +class StraikerWebhookStream(BaseModel): + phase: StraikerWebhookStreamPhase = "none" + index: int | None = None + + +class StraikerWebhookEvent(BaseModel): + type: StraikerWebhookEventType + id: str + stream: StraikerWebhookStream = Field(default_factory=StraikerWebhookStream) + + +class StraikerWebhookContent(BaseModel): + model_config = ConfigDict(extra="ignore") + + texts: list[str] = Field(default_factory=list) + images: list[str] = Field(default_factory=list) + structured_messages: list[AllMessageValues] | None = None + tools: list[dict[str, object]] | None = None + tool_calls: list[ChatCompletionToolCallChunk] | list[ChatCompletionMessageToolCall] | None = None + finish_reason: str | None = None + + +class StraikerWebhookUsage(BaseModel): + input_tokens: int | None = None + output_tokens: int | None = None + + +class StraikerWebhookContext(BaseModel): + call_surface: str + model: str | None = None + model_provider: str | None = None + destination: str | None = None + session_id: str | None = None + litellm_call_id: str | None = None + litellm_trace_id: str | None = None + litellm_version: str | None = None + + +class StraikerWebhookIdentity(BaseModel): + litellm_key: str | None = None + litellm_team: str | None = None + litellm_user_id: str | None = None + litellm_user_email: str | None = None + litellm_org_id: str | None = None + end_user_id: str | None = None + + +class StraikerWebhookApplication(BaseModel): + source: str + name: str | None = None + + +class StraikerWebhookRequest(BaseModel): + schema_version: str = STRAIKER_WEBHOOK_SCHEMA_VERSION + event: StraikerWebhookEvent + request: StraikerWebhookContent + response: StraikerWebhookContent | None = None + context: StraikerWebhookContext + identity: StraikerWebhookIdentity + application: StraikerWebhookApplication + usage: StraikerWebhookUsage | None = None + metadata: dict[str, object] | None = None + + +class StraikerWebhookResponse(BaseModel): + model_config = ConfigDict(extra="allow") + + action: StraikerWebhookAction = "NONE" + blocked_reason: str | None = None + texts: list[str] | None = None + schema_version: str | None = None + turn_id: str | None = Field(default=None, alias="turnId") + + +class StraikerGuardrailConfigModelOptionalParams(BaseModel): + timeout: float | None = Field( + default=5.0, + gt=0.0, + description="Per-attempt HTTP timeout in seconds.", + ) + max_retries: int | None = Field( + default=2, + ge=0, + description="Retries on transient HTTP (408/429/5xx) and network errors.", + ) + initial_backoff: float | None = Field( + default=0.1, + ge=0.0, + description="Initial retry backoff in seconds.", + ) + max_backoff: float | None = Field( + default=2.0, + ge=0.0, + description="Maximum retry backoff in seconds.", + ) + unreachable_fallback: Literal["fail_open", "fail_closed"] | None = Field( + default="fail_closed", + description="Behavior when Straiker is unreachable after retries.", + ) + fail_on_error: bool | None = Field( + default=True, + description=( + "Behavior on any guardrail error, not just unreachability. True (default) blocks " + "the request on error; False logs and allows the request to proceed." + ), + ) + max_payload_bytes: int | None = Field( + default=524288, + gt=0, + description="Maximum serialized webhook payload size sent to Straiker.", + ) + custom_headers: dict[str, str] | None = Field( + default=None, + description="Additional HTTP headers sent to Straiker, excluding Authorization and the webhook-format header.", + ) + metadata: dict[str, str] | None = Field( + default=None, + description=( + "Default metadata key/values added to the webhook metadata bag on every request. " + "On key conflict with request-derived metadata, these configured values win." + ), + ) + verbose: bool | None = Field( + default=False, + description="Log webhook request/response payloads and record action/turn_id in response hidden params.", + ) + + +class StraikerGuardrailConfigModel(GuardrailConfigModel[StraikerGuardrailConfigModelOptionalParams]): + api_key: str = Field( + min_length=1, + description="Straiker DefendAI environment API key (Bearer token). Env: STRAIKER_API_KEY.", + json_schema_extra={"secret": True}, + ) + + api_base: str | None = Field( + default="https://api.prod.straiker.ai", + description="Straiker API base URL. Use the regional variant for non-US tenants.", + ) + + default_app: str | None = Field( + default="LiteLLM Gateway", + description=( + "Default application registered in the Straiker Defend Console. " + "Overridden per-request by metadata.agent_id when present." + ), + ) + + @staticmethod + def ui_friendly_name() -> str: + return "Straiker" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py new file mode 100644 index 00000000000..ca57118ee9d --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py @@ -0,0 +1,733 @@ +import json +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest + +from litellm.exceptions import GuardrailRaisedException, ModifyResponseException +from litellm.proxy.guardrails.guardrail_hooks.straiker import initialize_guardrail +from litellm.proxy.guardrails.guardrail_hooks.straiker.straiker import ( + StraikerGuardrail, +) +from litellm.proxy.guardrails.guardrail_registry import ( + guardrail_class_registry, + guardrail_initializer_registry, +) +from litellm.types.proxy.guardrails.guardrail_hooks.straiker import ( + StraikerGuardrailConfigModel, + StraikerGuardrailConfigModelOptionalParams, +) +from litellm.types.utils import Choices, Message, ModelResponse, Usage + + +def _mock_response(action: str, turn_id: str = "turn-1", schema_version: str = "1", **extra) -> MagicMock: + resp = MagicMock(spec=httpx.Response) + resp.status_code = 200 + resp.json.return_value = { + "schema_version": schema_version, + "action": action, + "turn_id": turn_id, + **extra, + } + resp.text = "" + return resp + + +def _make_guardrail(**overrides) -> StraikerGuardrail: + defaults = dict( + api_key="test-key", + api_base="https://test.straiker.ai", + max_retries=0, + guardrail_name="straiker", + event_hook="pre_call", + async_handler=MagicMock(spec=httpx.AsyncClient), + ) + defaults.update(overrides) + g = StraikerGuardrail(**defaults) + g.async_handler.post = AsyncMock() + return g + + +def _logging_obj() -> MagicMock: + obj = MagicMock() + obj.litellm_call_id = "call-123" + obj.litellm_trace_id = "trace-456" + obj.call_type = "acompletion" + return obj + + +def _posted_payload(g: StraikerGuardrail) -> dict: + return json.loads(g.async_handler.post.call_args.kwargs["content"]) + + +def test_registry_membership(): + assert "straiker" in guardrail_initializer_registry + assert guardrail_class_registry["straiker"] is StraikerGuardrail + + +def test_config_model_wiring(): + assert StraikerGuardrailConfigModel.ui_friendly_name() == "Straiker" + assert StraikerGuardrail.get_config_model() is StraikerGuardrailConfigModel + fields = StraikerGuardrailConfigModel.model_fields + assert "api_key" in fields + assert "api_base" in fields + assert "default_app" in fields + assert "source" not in fields + assert "optional_params" in fields + assert "timeout" not in fields + assert "verbose" not in fields + + +def test_init_rejects_empty_api_key(): + with pytest.raises(ValueError): + StraikerGuardrail(api_key="") + + +def test_init_rejects_invalid_fallback(): + with pytest.raises(ValueError): + StraikerGuardrail(api_key="k", unreachable_fallback="nope") + + +def test_supported_hooks_limited_to_pre_and_post(): + from litellm.types.guardrails import GuardrailEventHooks + + assert StraikerGuardrail.get_supported_event_hooks() == [ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + ] + + +def test_during_call_mode_rejected_at_init(): + with pytest.raises(ValueError): + StraikerGuardrail(api_key="k", event_hook="during_call") + + +def test_streaming_attrs_hardcoded_to_buffered(): + g = _make_guardrail() + assert g.streaming_buffer_until_moderated is True + assert g.streaming_end_of_stream_only is True + + +def test_streaming_flags_not_configurable(): + fields = StraikerGuardrailConfigModelOptionalParams.model_fields + assert "streaming_buffer_until_moderated" not in fields + assert "streaming_end_of_stream_only" not in fields + assert "streaming_sampling_rate" not in fields + + +def test_initializer_builds_working_callback(): + from litellm.types.guardrails import LitellmParams + + params = LitellmParams(guardrail="straiker", mode="pre_call", api_key="abc", api_base="https://x.straiker.ai") + callback = initialize_guardrail(params, {"guardrail_name": "straiker"}) + assert isinstance(callback, StraikerGuardrail) + assert callback.api_base == "https://x.straiker.ai" + + +def test_initializer_maps_default_app_to_source(): + from litellm.types.guardrails import LitellmParams + + params = LitellmParams( + guardrail="straiker", + mode="pre_call", + api_key="abc", + default_app="My App", + ) + callback = initialize_guardrail(params, {"guardrail_name": "straiker"}) + assert callback.source == "My App" + + +def test_initializer_reads_optional_params_flattened_like_ui(): + from litellm.types.guardrails import LitellmParams + + params = LitellmParams( + guardrail="straiker", + mode="pre_call", + api_key="abc", + api_base="https://x.straiker.ai", + timeout=9.5, + verbose=True, + unreachable_fallback="fail_open", + ) + callback = initialize_guardrail(params, {"guardrail_name": "straiker"}) + assert isinstance(callback, StraikerGuardrail) + assert callback.timeout == 9.5 + assert callback.verbose is True + assert callback.unreachable_fallback == "fail_open" + assert callback.api_base == "https://x.straiker.ai" + + +def test_initializer_reads_nested_optional_params(): + from types import SimpleNamespace + + from litellm.types.guardrails import LitellmParams + + params = LitellmParams.model_construct( + guardrail="straiker", + mode="pre_call", + api_key="abc", + api_base="https://x.straiker.ai", + optional_params=SimpleNamespace( + timeout=7.25, + verbose=True, + unreachable_fallback="fail_open", + ), + ) + callback = initialize_guardrail(params, {"guardrail_name": "straiker"}) + assert isinstance(callback, StraikerGuardrail) + assert callback.timeout == 7.25 + assert callback.verbose is True + assert callback.unreachable_fallback == "fail_open" + + +def test_initializer_reads_dict_optional_params(): + from litellm.types.guardrails import LitellmParams + + params = LitellmParams.model_construct( + guardrail="straiker", + mode="pre_call", + api_key="abc", + api_base="https://x.straiker.ai", + optional_params={"timeout": 7.25, "verbose": True, "unreachable_fallback": "fail_open"}, + ) + callback = initialize_guardrail(params, {"guardrail_name": "straiker"}) + assert isinstance(callback, StraikerGuardrail) + assert callback.timeout == 7.25 + assert callback.verbose is True + assert callback.unreachable_fallback == "fail_open" + + +@pytest.mark.asyncio +async def test_request_envelope_transport_and_shape(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + inputs = {"texts": ["hello world"], "model": "gpt-4o-mini"} + request_data = { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello world"}], + "metadata": {"user_api_key_alias": "team-key", "agent_id": "chatbot-app", "app_name": "Chatbot"}, + } + + out = await g.apply_guardrail(inputs=inputs, request_data=request_data, input_type="request", logging_obj=_logging_obj()) + + assert out is inputs + url = g.async_handler.post.call_args.args[0] + assert url == "https://test.straiker.ai/api/v1/detect/webhook" + headers = g.async_handler.post.call_args.kwargs["headers"] + assert headers["X-Straiker-Webhook-Format"] == "litellm" + assert headers["Authorization"] == "Bearer test-key" + + payload = _posted_payload(g) + assert payload["schema_version"] == "1" + assert payload["event"]["type"] == "pre_call" + assert payload["event"]["id"] == "call-123:request" + assert payload["request"]["texts"] == ["hello world"] + assert payload["context"]["litellm_call_id"] == "call-123" + assert payload["identity"]["litellm_key"] == "team-key" + assert payload["application"] == {"source": "chatbot-app", "name": "Chatbot"} + assert "session_id" not in payload["application"] + assert "user_name" not in payload["application"] + assert "user_role" not in payload["application"] + assert "response" not in payload + assert "metadata" not in payload + + +@pytest.mark.asyncio +async def test_webhook_metadata_session_id_and_opaque_passthrough(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={ + "model": "m", + "litellm_session_id": "sess-from-litellm", + "metadata": { + "agent_id": "chatbot-app", + "app_name": "Chatbot", + "user_api_key_alias": "team-key", + "custom_tag": "experiment-7", + "client_ip": "10.0.0.1", + }, + }, + input_type="request", + logging_obj=_logging_obj(), + ) + payload = _posted_payload(g) + assert payload["application"] == {"source": "chatbot-app", "name": "Chatbot"} + assert payload["identity"]["litellm_key"] == "team-key" + assert payload["context"]["session_id"] == "sess-from-litellm" + assert "session_id" not in payload["metadata"] + assert payload["metadata"] == { + "custom_tag": "experiment-7", + "client_ip": "10.0.0.1", + } + + +@pytest.mark.asyncio +async def test_webhook_metadata_never_forwards_proxy_internal_keys(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={ + "model": "m", + "metadata": { + "custom_tag": "experiment-7", + "user_api_key": "sk-hashed-secret", + "user_api_end_user_max_budget": 12.5, + }, + }, + input_type="request", + logging_obj=_logging_obj(), + ) + assert _posted_payload(g)["metadata"] == {"custom_tag": "experiment-7"} + + +@pytest.mark.asyncio +async def test_default_metadata_injected_and_config_wins_on_clash(): + g = _make_guardrail(metadata={"tenant": "acme", "custom_tag": "config-value"}) + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={ + "model": "m", + "metadata": {"custom_tag": "request-value", "client_ip": "10.0.0.1"}, + }, + input_type="request", + logging_obj=_logging_obj(), + ) + assert _posted_payload(g)["metadata"] == { + "client_ip": "10.0.0.1", + "custom_tag": "config-value", + "tenant": "acme", + } + + +@pytest.mark.asyncio +async def test_default_metadata_present_without_request_metadata(): + g = _make_guardrail(metadata={"tenant": "acme"}) + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={"model": "m"}, + input_type="request", + logging_obj=_logging_obj(), + ) + assert _posted_payload(g)["metadata"] == {"tenant": "acme"} + + +@pytest.mark.asyncio +async def test_context_session_id_from_request_metadata(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={"model": "m", "metadata": {"session_id": "sess-meta"}}, + input_type="request", + logging_obj=_logging_obj(), + ) + payload = _posted_payload(g) + assert payload["context"]["session_id"] == "sess-meta" + assert "metadata" not in payload + + +@pytest.mark.asyncio +async def test_identity_key_and_team_coalesce_alias_over_id(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={ + "model": "m", + "metadata": { + "user_api_key_alias": "prod-key", + "user_api_key_hash": "hash-abc", + "user_api_key_team_alias": "growth", + "user_api_key_team_id": "team-9", + }, + }, + input_type="request", + logging_obj=_logging_obj(), + ) + identity = _posted_payload(g)["identity"] + assert identity["litellm_key"] == "prod-key" + assert identity["litellm_team"] == "growth" + assert "key" not in identity + assert "team" not in identity + + +@pytest.mark.asyncio +async def test_identity_key_and_team_fall_back_to_hash_and_id(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={ + "model": "m", + "metadata": { + "user_api_key_hash": "hash-abc", + "user_api_key_team_id": "team-9", + }, + }, + input_type="request", + logging_obj=_logging_obj(), + ) + identity = _posted_payload(g)["identity"] + assert identity["litellm_key"] == "hash-abc" + assert identity["litellm_team"] == "team-9" + + +@pytest.mark.asyncio +async def test_identity_end_user_from_resolved_metadata(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={ + "model": "m", + "metadata": { + "user_api_key_end_user_id": "eu-meta", + "user_api_key_user_id": "default_user_id", + }, + "user": "eu-body", + }, + input_type="request", + logging_obj=_logging_obj(), + ) + identity = _posted_payload(g)["identity"] + assert identity["end_user_id"] == "eu-meta" + assert identity["litellm_user_id"] == "default_user_id" + assert _posted_payload(g)["application"] == {"source": g.source} + + +@pytest.mark.asyncio +async def test_identity_end_user_absent_without_resolved_metadata(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={"model": "m", "user": "eu-body", "metadata": {"user_api_key_user_id": "default_user_id"}}, + input_type="request", + logging_obj=_logging_obj(), + ) + assert "end_user_id" not in _posted_payload(g)["identity"] + + +@pytest.mark.asyncio +async def test_application_source_from_agent_id(): + g = _make_guardrail(source="litellm") + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={"model": "m", "metadata": {"agent_id": "analytics-app", "app_name": "Analytics"}}, + input_type="request", + logging_obj=_logging_obj(), + ) + assert _posted_payload(g)["application"] == {"source": "analytics-app", "name": "Analytics"} + +@pytest.mark.asyncio +async def test_request_block_raises_guardrail_exception_with_reason(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("BLOCKED", blocked_reason="prompt injection") + with pytest.raises(GuardrailRaisedException) as exc: + await g.apply_guardrail( + inputs={"texts": ["attack"]}, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + assert "prompt injection" in str(exc.value) + + +@pytest.mark.asyncio +async def test_guardrail_intervened_writes_back_modified_text_only(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("GUARDRAIL_INTERVENED", texts=["[redacted]"]) + inputs = {"texts": ["my ssn is 123"], "images": ["img-a"]} + out = await g.apply_guardrail( + inputs=inputs, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + assert out["texts"] == ["[redacted]"] + assert out["images"] == ["img-a"] + + +@pytest.mark.asyncio +async def test_streamed_response_intervention_converts_to_block(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("GUARDRAIL_INTERVENED", texts=["[redacted]"]) + response = ModelResponse( + choices=[Choices(finish_reason="stop", index=0, message=Message(content="secret", role="assistant"))], + model="gpt-4o-mini", + ) + request_data = { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "p"}], + "stream": True, + "response": response, + } + with pytest.raises(ModifyResponseException): + await g.apply_guardrail( + inputs={"texts": ["secret"], "model": "gpt-4o-mini"}, + request_data=request_data, + input_type="response", + logging_obj=_logging_obj(), + ) + + +@pytest.mark.asyncio +async def test_non_streamed_response_intervention_redacts(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("GUARDRAIL_INTERVENED", texts=["[redacted]"]) + response = ModelResponse( + choices=[Choices(finish_reason="stop", index=0, message=Message(content="secret", role="assistant"))], + model="gpt-4o-mini", + ) + request_data = { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "p"}], + "response": response, + } + out = await g.apply_guardrail( + inputs={"texts": ["secret"], "model": "gpt-4o-mini"}, + request_data=request_data, + input_type="response", + logging_obj=_logging_obj(), + ) + assert out["texts"] == ["[redacted]"] + + +@pytest.mark.asyncio +async def test_guardrail_intervened_without_texts_blocks(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("GUARDRAIL_INTERVENED") + with pytest.raises(GuardrailRaisedException): + await g.apply_guardrail( + inputs={"texts": ["my ssn is 123"]}, + request_data={"model": "m"}, + input_type="request", + logging_obj=_logging_obj(), + ) + + +@pytest.mark.asyncio +async def test_streamed_via_proxy_server_request_body_converts_to_block(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("GUARDRAIL_INTERVENED", texts=["[redacted]"]) + response = ModelResponse( + choices=[Choices(finish_reason="stop", index=0, message=Message(content="secret", role="assistant"))], + model="gpt-4o-mini", + ) + request_data = { + "model": "gpt-4o-mini", + "proxy_server_request": {"body": {"stream": True}}, + "response": response, + } + with pytest.raises(ModifyResponseException): + await g.apply_guardrail( + inputs={"texts": ["secret"], "model": "gpt-4o-mini"}, + request_data=request_data, + input_type="response", + logging_obj=_logging_obj(), + ) + + +@pytest.mark.asyncio +async def test_response_envelope_and_block_replaces_response(): + g = _make_guardrail(verbose=True) + g.async_handler.post.return_value = _mock_response("BLOCKED") + response = ModelResponse( + choices=[Choices(finish_reason="stop", index=0, message=Message(content="secret", role="assistant"))], + model="gpt-4o-mini", + ) + request_data = { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "original prompt"}], + "stream": True, + "response": response, + } + with pytest.raises(ModifyResponseException) as exc: + await g.apply_guardrail( + inputs={"texts": ["secret"], "model": "gpt-4o-mini"}, + request_data=request_data, + input_type="response", + logging_obj=_logging_obj(), + ) + + assert exc.value.original_response is response + payload = _posted_payload(g) + assert payload["event"]["type"] == "post_call" + assert payload["event"]["stream"]["phase"] == "assembled" + assert payload["response"]["texts"] == ["secret"] + assert payload["response"]["finish_reason"] == "stop" + assert payload["request"]["structured_messages"] == [{"role": "user", "content": "original prompt"}] + + +@pytest.mark.asyncio +async def test_post_call_fail_closed_raises_modify_response_exception(): + g = _make_guardrail(unreachable_fallback="fail_closed") + g.async_handler.post.side_effect = httpx.ConnectError("boom") + response = ModelResponse( + choices=[Choices(finish_reason="stop", index=0, message=Message(content="secret", role="assistant"))], + model="gpt-4o-mini", + ) + request_data = {"model": "gpt-4o-mini", "response": response} + with pytest.raises(ModifyResponseException) as exc: + await g.apply_guardrail( + inputs={"texts": ["secret"], "model": "gpt-4o-mini"}, + request_data=request_data, + input_type="response", + logging_obj=_logging_obj(), + ) + assert exc.value.original_response is response + + +@pytest.mark.asyncio +async def test_usage_tokens_on_post_call(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + response = ModelResponse( + choices=[Choices(finish_reason="stop", index=0, message=Message(content="hi", role="assistant"))], + model="gpt-4o-mini", + usage=Usage(prompt_tokens=11, completion_tokens=7, total_tokens=18), + ) + await g.apply_guardrail( + inputs={"texts": ["hi"], "model": "gpt-4o-mini"}, + request_data={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hey"}], "response": response}, + input_type="response", + logging_obj=_logging_obj(), + ) + usage = _posted_payload(g)["usage"] + assert usage == {"input_tokens": 11, "output_tokens": 7} + + +@pytest.mark.asyncio +async def test_usage_absent_on_pre_call(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data={"model": "gpt-4o-mini"}, + input_type="request", + logging_obj=_logging_obj(), + ) + assert "usage" not in _posted_payload(g) + + +@pytest.mark.asyncio +async def test_allow_returns_inputs_unchanged(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + inputs = {"texts": ["fine"]} + out = await g.apply_guardrail( + inputs=inputs, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + assert out is inputs + + +@pytest.mark.asyncio +async def test_unreachable_fail_closed_blocks(): + g = _make_guardrail(unreachable_fallback="fail_closed") + g.async_handler.post.side_effect = httpx.ConnectError("boom") + with pytest.raises(GuardrailRaisedException): + await g.apply_guardrail( + inputs={"texts": ["x"]}, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + + +@pytest.mark.asyncio +async def test_unreachable_fail_open_passes_through(): + g = _make_guardrail(unreachable_fallback="fail_open") + g.async_handler.post.side_effect = httpx.ConnectError("boom") + inputs = {"texts": ["x"]} + out = await g.apply_guardrail( + inputs=inputs, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + assert out is inputs + + +@pytest.mark.asyncio +async def test_fail_on_error_false_allows_on_bad_status(): + g = _make_guardrail(unreachable_fallback="fail_closed", fail_on_error=False) + bad = MagicMock(spec=httpx.Response) + bad.status_code = 400 + bad.text = "bad request" + g.async_handler.post.return_value = bad + inputs = {"texts": ["x"]} + out = await g.apply_guardrail( + inputs=inputs, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + assert out is inputs + + +@pytest.mark.asyncio +async def test_non_retryable_status_fail_closed_blocks(): + g = _make_guardrail(unreachable_fallback="fail_closed", fail_on_error=True) + bad = MagicMock(spec=httpx.Response) + bad.status_code = 401 + bad.text = "unauthorized" + g.async_handler.post.return_value = bad + with pytest.raises(GuardrailRaisedException): + await g.apply_guardrail( + inputs={"texts": ["x"]}, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + + +@pytest.mark.asyncio +async def test_payload_size_guard_fails_closed(): + g = _make_guardrail(max_payload_bytes=10) + inputs = {"texts": ["x" * 5000]} + with pytest.raises(GuardrailRaisedException): + await g.apply_guardrail( + inputs=inputs, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + g.async_handler.post.assert_not_called() + + +@pytest.mark.asyncio +async def test_payload_size_guard_blocks_even_with_fail_open(): + g = _make_guardrail(max_payload_bytes=10, unreachable_fallback="fail_open") + inputs = {"texts": ["x" * 5000]} + with pytest.raises(GuardrailRaisedException): + await g.apply_guardrail( + inputs=inputs, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + g.async_handler.post.assert_not_called() + + +@pytest.mark.asyncio +async def test_invalid_response_schema_blocks_even_with_fail_open(): + g = _make_guardrail(unreachable_fallback="fail_open") + bad = MagicMock(spec=httpx.Response) + bad.status_code = 200 + bad.json.return_value = {"action": "NOT_A_VALID_ACTION"} + bad.text = "" + g.async_handler.post.return_value = bad + with pytest.raises(GuardrailRaisedException): + await g.apply_guardrail( + inputs={"texts": ["x"]}, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + + +@pytest.mark.asyncio +async def test_unreachable_http_status_fail_open_passes(): + g = _make_guardrail(unreachable_fallback="fail_open") + resp = MagicMock(spec=httpx.Response) + resp.status_code = 503 + resp.text = "service unavailable" + g.async_handler.post.return_value = resp + inputs = {"texts": ["x"]} + out = await g.apply_guardrail( + inputs=inputs, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + assert out is inputs + + +@pytest.mark.asyncio +async def test_unreachable_http_status_fail_closed_blocks(): + g = _make_guardrail(unreachable_fallback="fail_closed") + resp = MagicMock(spec=httpx.Response) + resp.status_code = 503 + resp.text = "service unavailable" + g.async_handler.post.return_value = resp + with pytest.raises(GuardrailRaisedException): + await g.apply_guardrail( + inputs={"texts": ["x"]}, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) diff --git a/ui/litellm-dashboard/public/assets/logos/straiker.svg b/ui/litellm-dashboard/public/assets/logos/straiker.svg new file mode 100644 index 00000000000..bdfe0405736 --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/straiker.svg @@ -0,0 +1,9 @@ + + + + + + + + + diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts index 2ad5819b5f0..a40587cb3ae 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts @@ -300,4 +300,10 @@ export const GUARDRAIL_PRESETS: Record = { mode: "pre_call", defaultOn: false, }, + straiker: { + provider: "Straiker", + guardrailNameSuggestion: "Straiker Guardrail", + mode: "pre_call", + defaultOn: false, + }, }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts index f81277f13c3..ba11d3d400d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts @@ -442,6 +442,16 @@ export const PARTNER_GUARDRAIL_CARDS: GuardrailCardInfo[] = [ tags: ["Security", "Policy", "Prompt Injection"], providerKey: "Repelloai", }, + { + id: "straiker", + name: "Straiker", + description: + "Defend AI Agentic Guardrails: Indirect/Direct Prompt Injection, Tool Misuse, Malicious MCP and Skills", + category: "partner", + logo: `${ASSET_PREFIX}straiker.svg`, + tags: ["Agentic", "Prompt Injection", "Tool Misuse", "MCP", "Skills"], + providerKey: "Straiker", + }, ]; export const ALL_CARDS = [...LITELLM_CONTENT_FILTER_CARDS, ...PARTNER_GUARDRAIL_CARDS]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx index e8c4810ca69..a2873797096 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx @@ -166,6 +166,7 @@ export const guardrailLogoMap: Record = { Akto: `${asset_logos_folder}akto.svg`, "Qostodian Nexus": `${asset_logos_folder}qohash.jpg`, "RepelloAI Argus": `${asset_logos_folder}repelloai.png`, + Straiker: `${asset_logos_folder}straiker.svg`, }; export const getGuardrailLogoAndName = (guardrailValue: string): { logo: string; displayName: string } => {