mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-17 04:34:25 +00:00
Address PR review: fail_closed default and typing fixes
- Default unreachable_fallback to fail_closed (matches generic_guardrail_api) - Tighten Wingback initializer types; avoid mutating param dicts - Use PEP 604 union for logging_obj annotation - Add config model default coverage test
This commit is contained in:
parent
0068403868
commit
8763574221
4 changed files with 31 additions and 22 deletions
|
|
@ -1,4 +1,4 @@
|
|||
from typing import TYPE_CHECKING, Any, Final
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm.types.guardrails import SupportedGuardrailIntegrations
|
||||
|
||||
|
|
@ -8,27 +8,24 @@ if TYPE_CHECKING:
|
|||
from litellm.types.guardrails import Guardrail, LitellmParams
|
||||
|
||||
|
||||
def _get_config_value(litellm_params: Any, attribute_name: str) -> Any | None:
|
||||
return getattr(litellm_params, attribute_name, None)
|
||||
|
||||
|
||||
def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"):
|
||||
def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> WingbackGuardrail:
|
||||
import litellm
|
||||
|
||||
additional_params: Final = dict(
|
||||
getattr(litellm_params, "additional_provider_specific_params", None) or {}
|
||||
base_params = litellm_params.additional_provider_specific_params or {}
|
||||
wingback_app_id = litellm_params.wingback_app_id
|
||||
additional_params: Final = (
|
||||
{**base_params, "wingback_app_id": wingback_app_id}
|
||||
if wingback_app_id and "wingback_app_id" not in base_params
|
||||
else base_params
|
||||
)
|
||||
wingback_app_id = _get_config_value(litellm_params, "wingback_app_id")
|
||||
if wingback_app_id:
|
||||
additional_params.setdefault("wingback_app_id", wingback_app_id)
|
||||
|
||||
_wingback_callback: Final = WingbackGuardrail(
|
||||
api_base=litellm_params.api_base,
|
||||
api_key=litellm_params.api_key,
|
||||
wingback_app_id=wingback_app_id,
|
||||
additional_provider_specific_params=additional_params,
|
||||
unreachable_fallback=_get_config_value(litellm_params, "unreachable_fallback") or "fail_open",
|
||||
fail_on_error=_get_config_value(litellm_params, "fail_on_error"),
|
||||
unreachable_fallback=litellm_params.unreachable_fallback or "fail_closed",
|
||||
fail_on_error=litellm_params.fail_on_error,
|
||||
guardrail_name=guardrail.get("guardrail_name", ""),
|
||||
event_hook=litellm_params.mode,
|
||||
default_on=litellm_params.default_on,
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@
|
|||
# +-------------------------------------------------------------+
|
||||
|
||||
import os
|
||||
from typing import TYPE_CHECKING, Final, Literal, Optional
|
||||
from typing import TYPE_CHECKING, Final, Literal
|
||||
|
||||
from litellm.integrations.custom_guardrail import log_guardrail_information
|
||||
from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api.generic_guardrail_api import (
|
||||
|
|
@ -41,12 +41,15 @@ class WingbackGuardrail(GenericGuardrailAPI):
|
|||
)
|
||||
resolved_api_key: Final = api_key or os.environ.get("WINGBACK_INTEGRATION_API_KEY")
|
||||
|
||||
additional_params: Final = dict(kwargs.pop("additional_provider_specific_params", None) or {})
|
||||
if wingback_app_id and "wingback_app_id" not in additional_params:
|
||||
additional_params["wingback_app_id"] = wingback_app_id
|
||||
existing_params = kwargs.pop("additional_provider_specific_params", None) or {}
|
||||
additional_params: Final = (
|
||||
{**existing_params, "wingback_app_id": wingback_app_id}
|
||||
if wingback_app_id and "wingback_app_id" not in existing_params
|
||||
else existing_params
|
||||
)
|
||||
|
||||
kwargs.setdefault("guardrail_name", GUARDRAIL_NAME)
|
||||
kwargs.setdefault("unreachable_fallback", "fail_open")
|
||||
kwargs.setdefault("unreachable_fallback", "fail_closed")
|
||||
|
||||
super().__init__(
|
||||
api_base=resolved_api_base,
|
||||
|
|
@ -61,7 +64,7 @@ class WingbackGuardrail(GenericGuardrailAPI):
|
|||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional["LiteLLMLoggingObj"] = None,
|
||||
logging_obj: "LiteLLMLoggingObj | None" = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
"""
|
||||
Apply Wingback to the given inputs.
|
||||
|
|
|
|||
|
|
@ -31,10 +31,11 @@ class WingbackGuardrailConfigModel(GuardrailConfigModel):
|
|||
),
|
||||
)
|
||||
unreachable_fallback: Literal["fail_closed", "fail_open"] = Field(
|
||||
default="fail_open",
|
||||
default="fail_closed",
|
||||
description=(
|
||||
"Behavior when the Wingback connectors service is unreachable. "
|
||||
"fail_open allows traffic (monitor mode); fail_closed blocks traffic (enforce mode)."
|
||||
"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(
|
||||
|
|
|
|||
|
|
@ -15,6 +15,9 @@ from litellm.proxy.guardrails.guardrail_hooks.wingback.wingback import (
|
|||
WingbackGuardrail,
|
||||
)
|
||||
from litellm.types.guardrails import SupportedGuardrailIntegrations
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.wingback import (
|
||||
WingbackGuardrailConfigModel,
|
||||
)
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
|
||||
|
|
@ -24,6 +27,11 @@ def test_wingback_guard_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"
|
||||
|
||||
|
||||
class TestWingbackGuardrail:
|
||||
def setup_method(self):
|
||||
for key in ["WINGBACK_INTEGRATION_API_KEY", "WINGBACK_API_BASE"]:
|
||||
|
|
@ -46,7 +54,7 @@ class TestWingbackGuardrail:
|
|||
|
||||
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_open"
|
||||
assert guardrail.unreachable_fallback == "fail_closed"
|
||||
|
||||
def test_initialization_with_custom_api_base_and_app_id(self):
|
||||
guardrail = WingbackGuardrail(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue