scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 123

feat: add Straiker guardrail integration (#33781)

Browse source 
* 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 <chandra@straiker.ai>
This commit is contained in:
yucheng-berri 2026-07-17 20:31:29 -07:00 • committed by GitHub
parent 93afde8605
commit f759c75466
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 1541 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

71
litellm/proxy/guardrails/guardrail_hooks/straiker/__init__.py Normal file
View file

@ -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,
}

541
litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py Normal file
View file

@ -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

1
litellm/types/guardrails.py
View file

@ -131,6 +131,7 @@ class SupportedGuardrailIntegrations(Enum):
SINGULR = "singulr"
HEADROOM = "headroom"
COMPRESR = "compresr"
STRAIKER = "straiker"
class Role(Enum):

169
litellm/types/proxy/guardrails/guardrail_hooks/straiker.py Normal file
View file

@ -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"

733
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py Normal file
View file

@ -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()
)

9
ui/litellm-dashboard/public/assets/logos/straiker.svg Normal file
View file

@ -0,0 +1,9 @@
<svg width="35" height="49" viewBox="0 0 35 49" fill="none" xmlns="http://www.w3.org/2000/svg">
<path d="M27.0752 3.56763L22.0752 19.5676L16.5869 17.8528L18.7236 11.012L8.49805 19.3782L33.0723 25.9319L34.1514 30.9358L12.1514 48.9358L7.58691 45.8528L12.5869 29.8528L18.0752 31.5676L15.9375 38.4075L26.1631 30.0413L1.58984 23.4885L0.510742 18.4846L22.5107 0.484619L27.0752 3.56763Z" fill="url(#paint0_linear_193_10422)"/>
<defs>
<linearGradient id="paint0_linear_193_10422" x1="26.0107" y1="9.48462" x2="7.51631" y2="40.1868" gradientUnits="userSpaceOnUse">
<stop stop-color="#FFC971"/>
<stop offset="1" stop-color="#F06D9A"/>
</linearGradient>
</defs>
</svg>
Side by side

After

Width:  |  Height:  |  Size: 658 B

6
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts
View file

@ -300,4 +300,10 @@ export const GUARDRAIL_PRESETS: Record<string, GuardrailPreset> = {
mode: "pre_call",
defaultOn: false,
},
straiker: {
provider: "Straiker",
guardrailNameSuggestion: "Straiker Guardrail",
mode: "pre_call",
defaultOn: false,
},
};

10
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts
View file

@ -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];

1
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx
View file

@ -166,6 +166,7 @@ export const guardrailLogoMap: Record<string, string> = {
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 } => {