mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-14 04:09:09 +00:00
feat(guardrails): add Wingback guardrail integration (#46136)
* feat(guardrails): add Wingback guardrail integration
Mirrored from #38685 at d73b66544d (net diff against its merge base, applied on main)
* fix(guardrails): dispatch Wingback through the unified guardrail
The proxy only runs a unified guardrail when apply_guardrail is defined on
the guardrail's own class, so WingbackGuardrail now binds the parent's
decorated method in its class body. Chat, messages and responses pre, post
and during call checks reach Wingback again, and each check logs once.
Also types the initializer and config model for the lint gates and uses
monkeypatch.setenv in the unit test
* chore(guardrails): drop the Wingback dispatch comment
The unit and integration hook tests already fail when the apply_guardrail binding is removed, so the comment is redundant under the repo comment rule
---------
Co-authored-by: bsk-wb <bharath@wingback.ai>
Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com>
This commit is contained in:
parent
9621a60d27
commit
c3ca8a5140
14 changed files with 623 additions and 0 deletions
|
|
@ -17248,6 +17248,18 @@
|
|||
"description": "Custom message when a guardrail blocks an action. Supports placeholders like {tool_name}, {rule_id}, and {default_message}.",
|
||||
"title": "Violation Message Template"
|
||||
},
|
||||
"wingback_app_id": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"description": "Wingback external gateway integration name for request attribution (sent as additional_provider_specific_params.wingback_app_id).",
|
||||
"title": "Wingback App Id"
|
||||
},
|
||||
"xecguard_model": {
|
||||
"anyOf": [
|
||||
{
|
||||
|
|
|
|||
|
|
@ -0,0 +1,43 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm.types.guardrails import SupportedGuardrailIntegrations
|
||||
|
||||
from .wingback import WingbackGuardrail
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.guardrails import Guardrail, LitellmParams
|
||||
|
||||
|
||||
def initialize_guardrail(
|
||||
litellm_params: LitellmParams,
|
||||
guardrail: Guardrail,
|
||||
) -> WingbackGuardrail:
|
||||
import litellm
|
||||
|
||||
_wingback_callback: Final = WingbackGuardrail(
|
||||
api_base=litellm_params.api_base,
|
||||
api_key=litellm_params.api_key,
|
||||
wingback_app_id=litellm_params.wingback_app_id,
|
||||
additional_provider_specific_params=litellm_params.additional_provider_specific_params,
|
||||
unreachable_fallback=litellm_params.unreachable_fallback or "fail_closed",
|
||||
fail_on_error=litellm_params.fail_on_error,
|
||||
extra_headers=litellm_params.extra_headers,
|
||||
guardrail_name=guardrail["guardrail_name"],
|
||||
event_hook=litellm_params.mode,
|
||||
default_on=litellm_params.default_on,
|
||||
timeout=litellm_params.timeout,
|
||||
)
|
||||
|
||||
litellm.logging_callback_manager.add_litellm_callback(_wingback_callback)
|
||||
return _wingback_callback
|
||||
|
||||
|
||||
guardrail_initializer_registry: Final = {
|
||||
SupportedGuardrailIntegrations.WINGBACK.value: initialize_guardrail,
|
||||
}
|
||||
|
||||
guardrail_class_registry: Final = {
|
||||
SupportedGuardrailIntegrations.WINGBACK.value: WingbackGuardrail,
|
||||
}
|
||||
|
|
@ -0,0 +1,68 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Final, Literal
|
||||
|
||||
from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api.generic_guardrail_api import (
|
||||
GenericGuardrailAPI,
|
||||
)
|
||||
from litellm.types.guardrails import Mode
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.wingback import (
|
||||
WingbackGuardrailConfigModel,
|
||||
)
|
||||
|
||||
GUARDRAIL_NAME: Final = "wingback"
|
||||
DEFAULT_WINGBACK_API_BASE: Final = "https://api.wingback.ai/connectors"
|
||||
|
||||
|
||||
class WingbackGuardrail(GenericGuardrailAPI):
|
||||
"""Wingback runtime security via the LiteLLM Generic Guardrail API contract."""
|
||||
|
||||
apply_guardrail = GenericGuardrailAPI.apply_guardrail # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # parent annotates request_data as a bare dict
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_base: str | None = None,
|
||||
api_key: str | None = None,
|
||||
wingback_app_id: str | None = None,
|
||||
additional_provider_specific_params: Mapping[str, object] | None = None,
|
||||
unreachable_fallback: Literal["fail_closed", "fail_open"] = "fail_closed",
|
||||
fail_on_error: bool | None = True,
|
||||
extra_headers: Sequence[str] | None = None,
|
||||
guardrail_name: str | None = None,
|
||||
event_hook: str | Sequence[str] | Mode | None = None,
|
||||
default_on: bool | None = False,
|
||||
timeout: float | None = None,
|
||||
) -> None:
|
||||
resolved_api_base: Final = api_base or os.environ.get("WINGBACK_API_BASE") or DEFAULT_WINGBACK_API_BASE
|
||||
resolved_api_key: Final = api_key or os.environ.get("WINGBACK_INTEGRATION_API_KEY")
|
||||
|
||||
existing_params: Final = additional_provider_specific_params
|
||||
additional_params: Final = (
|
||||
{
|
||||
**existing_params,
|
||||
"wingback_app_id": wingback_app_id,
|
||||
}
|
||||
if wingback_app_id and existing_params is not None and "wingback_app_id" not in existing_params
|
||||
else (
|
||||
{"wingback_app_id": wingback_app_id} if wingback_app_id and existing_params is None else existing_params
|
||||
)
|
||||
)
|
||||
|
||||
super().__init__( # pyright: ignore[reportUnknownMemberType] # parent takes untyped extra_headers and **kwargs
|
||||
api_base=resolved_api_base,
|
||||
api_key=resolved_api_key,
|
||||
additional_provider_specific_params=additional_params,
|
||||
unreachable_fallback=unreachable_fallback,
|
||||
fail_on_error=fail_on_error,
|
||||
extra_headers=list(extra_headers) if extra_headers is not None else None,
|
||||
guardrail_name=guardrail_name or GUARDRAIL_NAME,
|
||||
event_hook=event_hook,
|
||||
default_on=default_on,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def get_config_model() -> type[WingbackGuardrailConfigModel] | None:
|
||||
return WingbackGuardrailConfigModel
|
||||
|
|
@ -73,6 +73,9 @@ from litellm.types.proxy.guardrails.guardrail_hooks.typesafe import (
|
|||
from litellm.types.proxy.guardrails.guardrail_hooks.vigil_guard import (
|
||||
VigilGuardGuardrailConfigModel,
|
||||
)
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.wingback import (
|
||||
WingbackGuardrailConfigModel,
|
||||
)
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.xecguard import (
|
||||
XecGuardConfigModel,
|
||||
)
|
||||
|
|
@ -152,6 +155,7 @@ class SupportedGuardrailIntegrations(Enum):
|
|||
AGENT_365 = "agent_365"
|
||||
LLM_SHIELD_PROXY = "llm_shield_proxy"
|
||||
CONDUCT = "conduct"
|
||||
WINGBACK = "wingback"
|
||||
DECISION_MODEL = "decision_model"
|
||||
|
||||
|
||||
|
|
@ -1347,6 +1351,7 @@ class LitellmParams( # pyright: ignore[reportIncompatibleVariableOverride] # o
|
|||
VigilGuardGuardrailConfigModel,
|
||||
SingulrGuardrailConfigModel,
|
||||
Agent365GuardrailConfigModel,
|
||||
WingbackGuardrailConfigModel,
|
||||
):
|
||||
guardrail: str = Field(description="The type of guardrail integration to use")
|
||||
decision_model: str | None = Field(
|
||||
|
|
|
|||
48
litellm/types/proxy/guardrails/guardrail_hooks/wingback.py
Normal file
48
litellm/types/proxy/guardrails/guardrail_hooks/wingback.py
Normal file
|
|
@ -0,0 +1,48 @@
|
|||
from typing import Literal
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
from .base import GuardrailConfigModel
|
||||
|
||||
|
||||
class WingbackGuardrailConfigModel(GuardrailConfigModel): # pyright: ignore[reportMissingTypeArgument] # parameterizing adds an optional_params field to the UI form
|
||||
"""Configuration parameters for the Wingback guardrail."""
|
||||
|
||||
api_key: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Wingback external gateway integration API key (wbk_eg_*). "
|
||||
"If not provided, the WINGBACK_INTEGRATION_API_KEY environment variable is checked."
|
||||
),
|
||||
)
|
||||
api_base: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Wingback connectors service base URL. LiteLLM appends "
|
||||
"/beta/litellm_basic_guardrail_api. Defaults to https://api.wingback.ai/connectors "
|
||||
"or WINGBACK_API_BASE when unset."
|
||||
),
|
||||
)
|
||||
wingback_app_id: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Wingback external gateway integration name for request attribution "
|
||||
"(sent as additional_provider_specific_params.wingback_app_id)."
|
||||
),
|
||||
)
|
||||
unreachable_fallback: Literal["fail_closed", "fail_open"] = Field(
|
||||
default="fail_closed",
|
||||
description=(
|
||||
"Behavior when the Wingback connectors service is unreachable. "
|
||||
"fail_closed blocks traffic when Wingback is unavailable (default, matches generic_guardrail_api). "
|
||||
"fail_open allows traffic for monitor-only rollouts."
|
||||
),
|
||||
)
|
||||
fail_on_error: bool | None = Field(
|
||||
default=True,
|
||||
description=("If False, allow requests when the guardrail returns an error other than an explicit block."),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def ui_friendly_name() -> str:
|
||||
return "Wingback"
|
||||
|
|
@ -149,6 +149,7 @@ PROVIDERS: Final = (
|
|||
pytest.param("pillar", "pillar", {}, "pre_call", False, id="pillar"),
|
||||
pytest.param("cisco", "cisco_ai_defense", {}, "pre_call", False, id="cisco-ai-defense"),
|
||||
pytest.param("vigil", "vigil_guard", {}, "pre_call", False, id="vigil-guard"),
|
||||
pytest.param("wingback", "wingback", {}, "pre_call", False, id="wingback"),
|
||||
pytest.param("singulr", "singulr", {}, "pre_call", False, id="singulr"),
|
||||
pytest.param("headroom", "headroom", {}, "pre_call", True, id="headroom"),
|
||||
pytest.param("onyx", "onyx", {}, "post_call", False, id="onyx"),
|
||||
|
|
|
|||
157
tests/integration/observability/test_wingback_guardrail.py
Normal file
157
tests/integration/observability/test_wingback_guardrail.py
Normal file
|
|
@ -0,0 +1,157 @@
|
|||
import json
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, eventually, gateway_from_environment, object_value
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.process import owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
|
||||
API_KEY: Final = "wbk_eg_synthetic"
|
||||
APP_ID: Final = "synthetic-app"
|
||||
DENIED: Final = "synthetic wingback denied marker"
|
||||
DENIAL: Final = "synthetic wingback denial"
|
||||
PRE_CALL: Final = "wingback-pre"
|
||||
POST_CALL: Final = "wingback-post"
|
||||
|
||||
|
||||
def _verdict(request: Request) -> Reply:
|
||||
texts: Final = json.loads(request.body).get("texts") or []
|
||||
denied: Final = any(DENIED in text for text in texts)
|
||||
result: Final = {"action": "BLOCKED", "blocked_reason": DENIAL} if denied else {"action": "NONE"}
|
||||
return Reply(body=json.dumps(result).encode())
|
||||
|
||||
|
||||
def _echo_provider(request: Request) -> Reply:
|
||||
prompt: Final = json.loads(request.body)["messages"][-1]["content"]
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": "chatcmpl-" + uuid.uuid4().hex,
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "gpt-4o-mini",
|
||||
"choices": [
|
||||
{"index": 0, "message": {"role": "assistant", "content": f"echo {prompt}"}, "finish_reason": "stop"}
|
||||
],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
def _wingback(name: str, mode: str, policy: str) -> dict[str, object]:
|
||||
return {
|
||||
"guardrail_name": name,
|
||||
"litellm_params": {
|
||||
"guardrail": "wingback",
|
||||
"mode": mode,
|
||||
"default_on": False,
|
||||
"api_base": policy,
|
||||
"api_key": API_KEY,
|
||||
"wingback_app_id": APP_ID,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Rig:
|
||||
proxy: Gateway
|
||||
policy: Wire
|
||||
provider: Wire
|
||||
model: str
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
root: Final = tmp_path_factory.mktemp("wingback")
|
||||
with (
|
||||
gateway_from_environment() as gateway,
|
||||
wire_server(_verdict) as policy,
|
||||
wire_server(_echo_provider) as provider,
|
||||
):
|
||||
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
config["litellm_settings"]["cache"] = False
|
||||
config["guardrails"] = [
|
||||
_wingback(PRE_CALL, "pre_call", policy.url),
|
||||
_wingback(POST_CALL, "post_call", policy.url),
|
||||
]
|
||||
path: Final = root / "wingback.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
with (
|
||||
owned_proxy_process(gateway, root, {}, config=path, workers=2) as owned,
|
||||
owned.gateway.scenario() as scenario,
|
||||
):
|
||||
model: Final = scenario.model(
|
||||
model="openai/gpt-4o-mini", api_base=provider.url + "/v1", api_key="synthetic-openai-key"
|
||||
)
|
||||
yield Rig(owned.gateway, policy, provider, model)
|
||||
|
||||
|
||||
def _chat(rig: Rig, guardrail: str, text: str) -> httpx.Response:
|
||||
rig.policy.drain()
|
||||
rig.provider.drain()
|
||||
return rig.proxy.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": rig.model, "messages": [{"role": "user", "content": text}], "guardrails": [guardrail]},
|
||||
)
|
||||
|
||||
|
||||
def _assert_wingback_request(request: Request, input_type: str, text: str) -> None:
|
||||
assert request.target == "/beta/litellm_basic_guardrail_api", request.target
|
||||
assert request.headers.get("x-api-key") == API_KEY
|
||||
body: Final = json.loads(request.body)
|
||||
assert body["input_type"] == input_type, body
|
||||
assert any(text in sent for sent in body["texts"]), body
|
||||
assert body["additional_provider_specific_params"] == {"wingback_app_id": APP_ID}, body
|
||||
|
||||
|
||||
def test_pre_call_denial_never_reaches_the_provider(rig: Rig) -> None:
|
||||
response: Final = _chat(rig, PRE_CALL, f"{DENIED} {uuid.uuid4().hex}")
|
||||
|
||||
assert response.status_code == 400 and DENIAL in response.text, response.text
|
||||
assert rig.provider.drain() == ()
|
||||
checks: Final = rig.policy.drain()
|
||||
assert len(checks) == 1, checks
|
||||
_assert_wingback_request(checks[0], "request", DENIED)
|
||||
|
||||
|
||||
def test_post_call_denial_blocks_the_provider_response(rig: Rig) -> None:
|
||||
response: Final = _chat(rig, POST_CALL, f"{DENIED} {uuid.uuid4().hex}")
|
||||
|
||||
assert response.status_code == 400 and DENIAL in response.text, response.text
|
||||
assert len(rig.provider.drain()) == 1
|
||||
checks: Final = rig.policy.drain()
|
||||
assert len(checks) == 1, checks
|
||||
_assert_wingback_request(checks[0], "response", f"echo {DENIED}")
|
||||
|
||||
|
||||
def test_allowed_call_is_checked_once_and_logged_once(rig: Rig) -> None:
|
||||
prompt: Final = "synthetic wingback allowed prompt " + uuid.uuid4().hex
|
||||
response: Final = _chat(rig, PRE_CALL, prompt)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json()["choices"][0]["message"]["content"] == f"echo {prompt}"
|
||||
assert len(rig.provider.drain()) == 1
|
||||
assert response.headers["x-litellm-applied-guardrails"] == PRE_CALL, dict(response.headers)
|
||||
checks: Final = rig.policy.drain()
|
||||
assert len(checks) == 1, checks
|
||||
_assert_wingback_request(checks[0], "request", prompt)
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (response.json()["id"],)),
|
||||
lambda values: len(values) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
entries: Final = object_value(rows[0]["metadata"])["guardrail_information"]
|
||||
assert isinstance(entries, list) and len(entries) == 1, entries
|
||||
entry: Final = object_value(entries[0])
|
||||
assert entry["guardrail_name"] == PRE_CALL, entry
|
||||
assert entry["guardrail_mode"] == "pre_call", entry
|
||||
assert entry["guardrail_status"] == "success", entry
|
||||
265
tests/unit/proxy/guardrails/guardrail_hooks/test_wingback.py
Normal file
265
tests/unit/proxy/guardrails/guardrail_hooks/test_wingback.py
Normal file
|
|
@ -0,0 +1,265 @@
|
|||
import json
|
||||
import os
|
||||
from collections.abc import Iterator, Mapping
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.exceptions import GuardrailRaisedException
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.guardrails.guardrail_hooks.wingback import (
|
||||
guardrail_class_registry,
|
||||
guardrail_initializer_registry,
|
||||
initialize_guardrail,
|
||||
)
|
||||
from litellm.proxy.guardrails.guardrail_hooks.wingback.wingback import (
|
||||
DEFAULT_WINGBACK_API_BASE,
|
||||
WingbackGuardrail,
|
||||
)
|
||||
from litellm.types.guardrails import SupportedGuardrailIntegrations
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.wingback import (
|
||||
WingbackGuardrailConfigModel,
|
||||
)
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.types.utils import Choices, GenericGuardrailAPIInputs, Message, ModelResponse
|
||||
|
||||
|
||||
def test_wingback_guard_registry():
|
||||
assert SupportedGuardrailIntegrations.WINGBACK.value in guardrail_initializer_registry
|
||||
assert SupportedGuardrailIntegrations.WINGBACK.value in guardrail_class_registry
|
||||
assert guardrail_class_registry["wingback"].get_config_model().ui_friendly_name() == "Wingback"
|
||||
|
||||
|
||||
def test_wingback_config_model_defaults():
|
||||
assert WingbackGuardrailConfigModel.ui_friendly_name() == "Wingback"
|
||||
assert WingbackGuardrailConfigModel.model_fields["unreachable_fallback"].default == "fail_closed"
|
||||
|
||||
|
||||
def test_initialize_guardrail_passes_extra_headers():
|
||||
litellm_params = MagicMock()
|
||||
litellm_params.api_base = "http://localhost:8101"
|
||||
litellm_params.api_key = "wbk_eg_test"
|
||||
litellm_params.wingback_app_id = None
|
||||
litellm_params.additional_provider_specific_params = None
|
||||
litellm_params.unreachable_fallback = "fail_closed"
|
||||
litellm_params.fail_on_error = True
|
||||
litellm_params.mode = "pre_call"
|
||||
litellm_params.default_on = True
|
||||
litellm_params.extra_headers = ["x-request-id"]
|
||||
litellm_params.timeout = 12.5
|
||||
|
||||
with patch("litellm.logging_callback_manager.add_litellm_callback"):
|
||||
instance = initialize_guardrail(
|
||||
litellm_params,
|
||||
{"guardrail_name": "wingback-runtime-security"},
|
||||
)
|
||||
|
||||
assert instance.extra_headers == ["x-request-id"]
|
||||
assert instance.timeout == 12.5
|
||||
|
||||
|
||||
class TestWingbackGuardrail:
|
||||
def setup_method(self):
|
||||
for key in ["WINGBACK_INTEGRATION_API_KEY", "WINGBACK_API_BASE"]:
|
||||
if key in os.environ:
|
||||
del os.environ[key]
|
||||
|
||||
def teardown_method(self):
|
||||
for key in ["WINGBACK_INTEGRATION_API_KEY", "WINGBACK_API_BASE"]:
|
||||
if key in os.environ:
|
||||
del os.environ[key]
|
||||
|
||||
def test_initialization_with_defaults(self, monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setenv("WINGBACK_INTEGRATION_API_KEY", "wbk_eg_test_key")
|
||||
|
||||
guardrail = WingbackGuardrail(
|
||||
guardrail_name="wingback-runtime-security",
|
||||
event_hook="pre_call",
|
||||
default_on=True,
|
||||
)
|
||||
|
||||
assert guardrail.api_base == f"{DEFAULT_WINGBACK_API_BASE}/beta/litellm_basic_guardrail_api"
|
||||
assert guardrail.headers["x-api-key"] == "wbk_eg_test_key"
|
||||
assert guardrail.unreachable_fallback == "fail_closed"
|
||||
|
||||
def test_initialization_with_custom_api_base_and_app_id(self):
|
||||
guardrail = WingbackGuardrail(
|
||||
api_base="http://localhost:8101",
|
||||
api_key="wbk_eg_custom",
|
||||
wingback_app_id="local-litellm",
|
||||
guardrail_name="wingback-runtime-security",
|
||||
event_hook="post_call",
|
||||
default_on=False,
|
||||
unreachable_fallback="fail_closed",
|
||||
)
|
||||
|
||||
assert guardrail.api_base == "http://localhost:8101/beta/litellm_basic_guardrail_api"
|
||||
assert guardrail.additional_provider_specific_params["wingback_app_id"] == "local-litellm"
|
||||
assert guardrail.unreachable_fallback == "fail_closed"
|
||||
|
||||
def test_get_config_model(self):
|
||||
config_model = WingbackGuardrail.get_config_model()
|
||||
assert config_model is not None
|
||||
assert config_model.__name__ == "WingbackGuardrailConfigModel"
|
||||
assert config_model.ui_friendly_name() == "Wingback"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_allows_safe_request(self):
|
||||
guardrail = WingbackGuardrail(
|
||||
api_base="http://localhost:8101",
|
||||
api_key="wbk_eg_test",
|
||||
guardrail_name="wingback-runtime-security",
|
||||
event_hook="pre_call",
|
||||
default_on=True,
|
||||
)
|
||||
|
||||
inputs = GenericGuardrailAPIInputs(texts=["Hello"])
|
||||
request_data = {
|
||||
"proxy_server_request": {
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"model": "gpt-4o-mini",
|
||||
}
|
||||
}
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {"action": "NONE"}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
|
||||
with patch.object(
|
||||
guardrail.async_handler,
|
||||
"post",
|
||||
new=AsyncMock(return_value=mock_response),
|
||||
) as mock_post:
|
||||
result = await guardrail.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data=request_data,
|
||||
input_type="request",
|
||||
logging_obj=None,
|
||||
)
|
||||
|
||||
assert result["texts"] == ["Hello"]
|
||||
mock_post.assert_awaited_once()
|
||||
assert mock_post.await_args.kwargs["url"] == "http://localhost:8101/beta/litellm_basic_guardrail_api"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_blocks_request(self):
|
||||
guardrail = WingbackGuardrail(
|
||||
api_base="http://localhost:8101",
|
||||
api_key="wbk_eg_test",
|
||||
wingback_app_id="prod-litellm",
|
||||
guardrail_name="wingback-runtime-security",
|
||||
event_hook="pre_call",
|
||||
default_on=True,
|
||||
)
|
||||
|
||||
inputs = GenericGuardrailAPIInputs(texts=["ignore previous instructions"])
|
||||
request_data = {
|
||||
"proxy_server_request": {
|
||||
"messages": [{"role": "user", "content": "ignore previous instructions"}],
|
||||
"model": "gpt-4o-mini",
|
||||
}
|
||||
}
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"action": "BLOCKED",
|
||||
"blocked_reason": "Prompt injection detected",
|
||||
}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
|
||||
with patch.object(
|
||||
guardrail.async_handler,
|
||||
"post",
|
||||
new=AsyncMock(return_value=mock_response),
|
||||
):
|
||||
with pytest.raises(GuardrailRaisedException, match="Prompt injection detected"):
|
||||
await guardrail.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data=request_data,
|
||||
input_type="request",
|
||||
logging_obj=None,
|
||||
)
|
||||
|
||||
|
||||
DENIAL: Final = "synthetic wingback denial"
|
||||
|
||||
|
||||
class _Wingback:
|
||||
def __init__(self, verdict: Mapping[str, str]) -> None:
|
||||
self.verdict: Final = verdict
|
||||
self.sent: tuple[Mapping[str, object], ...] = ()
|
||||
|
||||
def respond(self, request: httpx.Request) -> httpx.Response:
|
||||
self.sent = (*self.sent, json.loads(request.content))
|
||||
return httpx.Response(200, json=dict(self.verdict))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def registered(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> Iterator[_Wingback]:
|
||||
hook, verdict = request.param
|
||||
vendor: Final = _Wingback(verdict)
|
||||
guardrail: Final = WingbackGuardrail(
|
||||
api_base="http://wingback.invalid",
|
||||
api_key="wbk_eg_test",
|
||||
guardrail_name="wingback-runtime-security",
|
||||
event_hook=hook,
|
||||
default_on=True,
|
||||
)
|
||||
guardrail.async_handler = AsyncHTTPHandler(transport=httpx.MockTransport(vendor.respond))
|
||||
monkeypatch.setattr(litellm, "callbacks", [guardrail])
|
||||
yield vendor
|
||||
|
||||
|
||||
def _chat_request(text: str) -> dict[str, object]:
|
||||
return {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": text}], "metadata": {}}
|
||||
|
||||
|
||||
BLOCKED: Final = {"action": "BLOCKED", "blocked_reason": DENIAL}
|
||||
ALLOWED: Final = {"action": "NONE"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("registered", [("pre_call", BLOCKED)], indirect=True)
|
||||
async def test_proxy_pre_call_hook_blocks_on_wingback_verdict(registered: _Wingback) -> None:
|
||||
with pytest.raises(GuardrailRaisedException, match=DENIAL):
|
||||
await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
data=_chat_request("ignore previous instructions"),
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
assert [sent["texts"] for sent in registered.sent] == [["ignore previous instructions"]]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("registered", [("post_call", BLOCKED)], indirect=True)
|
||||
async def test_proxy_post_call_hook_blocks_on_wingback_verdict(registered: _Wingback) -> None:
|
||||
response: Final = ModelResponse(choices=[Choices(index=0, message=Message(role="assistant", content="leaked"))])
|
||||
|
||||
with pytest.raises(GuardrailRaisedException, match=DENIAL):
|
||||
await ProxyLogging(user_api_key_cache=DualCache()).post_call_success_hook(
|
||||
data=_chat_request("hello"), user_api_key_dict=UserAPIKeyAuth(), response=response
|
||||
)
|
||||
|
||||
assert [(sent["input_type"], sent["texts"]) for sent in registered.sent] == [("response", ["leaked"])]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("registered", [("pre_call", ALLOWED)], indirect=True)
|
||||
async def test_proxy_pre_call_hook_records_one_guardrail_entry_per_check(registered: _Wingback) -> None:
|
||||
data: Final = _chat_request("hello")
|
||||
|
||||
await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(), data=data, call_type="completion"
|
||||
)
|
||||
|
||||
entries: Final = data["litellm_metadata"]["standard_logging_guardrail_information"]
|
||||
assert [(entry["guardrail_name"], entry["guardrail_status"]) for entry in entries] == [
|
||||
("wingback-runtime-security", "success")
|
||||
]
|
||||
assert len(registered.sent) == 1
|
||||
BIN
ui/litellm-dashboard/public/assets/logos/wingback.png
Normal file
BIN
ui/litellm-dashboard/public/assets/logos/wingback.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 14 KiB |
|
|
@ -379,4 +379,10 @@ export const GUARDRAIL_PRESETS: Record<string, GuardrailPreset> = {
|
|||
mode: "pre_call",
|
||||
defaultOn: false,
|
||||
},
|
||||
wingback: {
|
||||
provider: "Wingback",
|
||||
guardrailNameSuggestion: "Wingback Runtime Security",
|
||||
mode: "pre_call",
|
||||
defaultOn: false,
|
||||
},
|
||||
};
|
||||
|
|
|
|||
|
|
@ -37,6 +37,7 @@ const EXPECTED_PARTNER_LOGO_FILES: Record<string, string> = {
|
|||
agent_365: "microsoft_azure.svg",
|
||||
llm_shield_proxy: "llm_shield_proxy.svg",
|
||||
conduct: "conduct.png",
|
||||
wingback: "wingback.png",
|
||||
};
|
||||
|
||||
describe("guardrail_garden_data logos", () => {
|
||||
|
|
|
|||
|
|
@ -557,6 +557,16 @@ export const PARTNER_GUARDRAIL_CARDS: GuardrailCardInfo[] = [
|
|||
tags: ["Security", "Prompt Injection", "PII", "Policy"],
|
||||
providerKey: "Conduct",
|
||||
},
|
||||
{
|
||||
id: "wingback",
|
||||
name: "Wingback",
|
||||
description:
|
||||
"Wingback runtime security for LiteLLM — PII detection, prompt injection, toxic content, and policy enforcement via the Wingback connectors service.",
|
||||
category: "partner",
|
||||
logo: guardrailLogoMap["Wingback"],
|
||||
tags: ["Security", "PII", "Prompt Injection", "Runtime"],
|
||||
providerKey: "Wingback",
|
||||
},
|
||||
];
|
||||
|
||||
export const ALL_CARDS = [...LITELLM_CONTENT_FILTER_CARDS, ...DECISION_MODEL_CARDS, ...PARTNER_GUARDRAIL_CARDS];
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ import promptguardLogo from "../../../../../public/assets/logos/promptguard.svg"
|
|||
import qohashLogo from "../../../../../public/assets/logos/qohash.jpg";
|
||||
import repelloAiLogo from "../../../../../public/assets/logos/repelloai.png";
|
||||
import straikerLogo from "../../../../../public/assets/logos/straiker.svg";
|
||||
import wingbackLogo from "../../../../../public/assets/logos/wingback.png";
|
||||
import xecguardLogo from "../../../../../public/assets/logos/xecguard.svg";
|
||||
import zscalerLogo from "../../../../../public/assets/logos/zscaler.svg";
|
||||
|
||||
|
|
@ -323,6 +324,7 @@ export const guardrailLogoMap = {
|
|||
"Microsoft Agent 365": microsoftAzureLogo.src,
|
||||
"LLM Shield Proxy": llmShieldProxyLogo.src,
|
||||
"Conduct Guard": conductLogo.src,
|
||||
Wingback: wingbackLogo.src,
|
||||
"Decision Model": litellmLogo.src,
|
||||
} satisfies Record<string, string>;
|
||||
|
||||
|
|
|
|||
5
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
5
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -36624,6 +36624,11 @@ export interface components {
|
|||
* @description Custom message when a guardrail blocks an action. Supports placeholders like {tool_name}, {rule_id}, and {default_message}.
|
||||
*/
|
||||
violation_message_template?: string | null;
|
||||
/**
|
||||
* Wingback App Id
|
||||
* @description Wingback external gateway integration name for request attribution (sent as additional_provider_specific_params.wingback_app_id).
|
||||
*/
|
||||
wingback_app_id?: string | null;
|
||||
/**
|
||||
* Xecguard Model
|
||||
* @description XecGuard scanning model identifier. Defaults to 'xecguard_v2'.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue