mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
name resolve failed policy flag
This commit is contained in:
Shalom Jamil
parent
f93552cb1a
commit
a33c6247e6
4 changed files with 322 additions and 32 deletions
|
|
@ -21,6 +21,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
|
|||
post_checkpoint_id = getattr(litellm_params, "post_checkpoint_id", None)
|
||||
file_checkpoint_id = getattr(litellm_params, "file_checkpoint_id", None)
|
||||
enable_routing_cache = getattr(litellm_params, "enable_routing_cache", None)
|
||||
fail_if_no_application = getattr(litellm_params, "fail_if_no_application", None)
|
||||
|
||||
_ovalix_callback = OvalixGuardrail(
|
||||
guardrail_name=guardrail.get("guardrail_name", ""),
|
||||
|
|
@ -31,6 +32,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
|
|||
post_checkpoint_id=post_checkpoint_id,
|
||||
file_checkpoint_id=file_checkpoint_id,
|
||||
enable_routing_cache=enable_routing_cache,
|
||||
fail_if_no_application=fail_if_no_application,
|
||||
event_hook=litellm_params.mode,
|
||||
default_on=litellm_params.default_on,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -46,7 +46,9 @@ if TYPE_CHECKING:
|
|||
BLOCKED_BY_OVALIX_FALLBACK_MESSAGE = "This message was blocked by Ovalix"
|
||||
BLOCKED_ACTION_TYPE = "block"
|
||||
_MODIFY_ACTION_TYPES = ("anonymize", "sanitize")
|
||||
_APPLICATION_NOT_FOUND_STATUS = 404
|
||||
_ROUTING_CACHE_TTL_SECONDS = 3600
|
||||
_ROUTING_CACHE_NEGATIVE_TTL_SECONDS = 300
|
||||
_ROUTING_CACHE_MAX_SIZE = 1000
|
||||
_DEFAULT_FILE_SIZE_LIMIT = 64 * 1024 * 1024
|
||||
_FILE_BLOCK_ESCALATION_REASON = (
|
||||
|
|
@ -128,6 +130,7 @@ class OvalixGuardrail(CustomGuardrail):
|
|||
post_checkpoint_id: Optional[str] = None,
|
||||
file_checkpoint_id: str | None = None,
|
||||
enable_routing_cache: bool | None = None,
|
||||
fail_if_no_application: bool | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
self._tracker_api_base = tracker_api_base or os.environ.get("OVALIX_TRACKER_API_BASE")
|
||||
|
|
@ -143,7 +146,14 @@ class OvalixGuardrail(CustomGuardrail):
|
|||
self._enable_routing_cache = (
|
||||
True if resolved_enable_routing_cache is None else _coerce_bool(resolved_enable_routing_cache)
|
||||
)
|
||||
self._routing_cache: OrderedDict[str, tuple[float, ResolvedRouting]] = OrderedDict()
|
||||
env_fail_if_no_application = os.environ.get("OVALIX_FAIL_IF_NO_APPLICATION")
|
||||
resolved_fail_if_no_application = (
|
||||
fail_if_no_application if fail_if_no_application is not None else env_fail_if_no_application
|
||||
)
|
||||
self._fail_if_no_application = (
|
||||
True if resolved_fail_if_no_application is None else _coerce_bool(resolved_fail_if_no_application)
|
||||
)
|
||||
self._routing_cache: OrderedDict[str, tuple[float, ResolvedRouting | None]] = OrderedDict()
|
||||
self._app_name_regex: re.Pattern[str] | None = None
|
||||
|
||||
if "supported_event_hooks" not in kwargs:
|
||||
|
|
@ -313,6 +323,8 @@ class OvalixGuardrail(CustomGuardrail):
|
|||
logging_obj: Optional[Any] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
routing = await self._resolve_routing(request_data)
|
||||
if routing is None:
|
||||
return inputs
|
||||
actor = self._get_actor(request_data)
|
||||
session_id = self._get_session_id_for_application(request_data, routing.application_id)
|
||||
is_response = input_type == "response"
|
||||
|
|
@ -474,24 +486,45 @@ class OvalixGuardrail(CustomGuardrail):
|
|||
name = (match.group(1) if match.groups() else match.group(0)).strip()
|
||||
return name or None
|
||||
|
||||
def _routing_cache_get(self, name: str) -> ResolvedRouting | None:
|
||||
def _routing_cache_get(self, name: str) -> tuple[bool, ResolvedRouting | None]:
|
||||
entry = self._routing_cache.get(name)
|
||||
if entry is None:
|
||||
return None
|
||||
stored_at, routing = entry
|
||||
if time.monotonic() - stored_at >= _ROUTING_CACHE_TTL_SECONDS:
|
||||
return False, None
|
||||
expires_at, routing = entry
|
||||
if time.monotonic() >= expires_at:
|
||||
del self._routing_cache[name]
|
||||
return None
|
||||
return False, None
|
||||
self._routing_cache.move_to_end(name)
|
||||
return routing
|
||||
return True, routing
|
||||
|
||||
def _routing_cache_put(self, name: str, routing: ResolvedRouting) -> None:
|
||||
self._routing_cache[name] = (time.monotonic(), routing)
|
||||
def _routing_cache_put(self, name: str, routing: ResolvedRouting | None) -> None:
|
||||
ttl = _ROUTING_CACHE_TTL_SECONDS if routing is not None else _ROUTING_CACHE_NEGATIVE_TTL_SECONDS
|
||||
self._routing_cache[name] = (time.monotonic() + ttl, routing)
|
||||
self._routing_cache.move_to_end(name)
|
||||
while len(self._routing_cache) > _ROUTING_CACHE_MAX_SIZE:
|
||||
self._routing_cache.popitem(last=False)
|
||||
|
||||
async def _resolve_routing(self, request_data: dict) -> ResolvedRouting:
|
||||
def _no_application(self, reason: str) -> None:
|
||||
if self._fail_if_no_application:
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name,
|
||||
message=f"Ovalix guardrail error: {reason}",
|
||||
should_wrap_with_default_message=False,
|
||||
)
|
||||
verbose_proxy_logger.warning(
|
||||
"Ovalix guardrail passing the call through unguarded (fail_if_no_application=false): %s", reason
|
||||
)
|
||||
return None
|
||||
|
||||
def _routing_error(self, error: Exception) -> GuardrailRaisedException:
|
||||
verbose_proxy_logger.exception("Ovalix routing resolution failed: %s", error)
|
||||
return GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name,
|
||||
message=f"Ovalix guardrail error: routing resolution failed: {error!s}",
|
||||
should_wrap_with_default_message=False,
|
||||
)
|
||||
|
||||
async def _resolve_routing(self, request_data: dict) -> ResolvedRouting | None:
|
||||
if self._application_id:
|
||||
return ResolvedRouting(
|
||||
self._application_id,
|
||||
|
|
@ -502,29 +535,23 @@ class OvalixGuardrail(CustomGuardrail):
|
|||
)
|
||||
alias = self._get_key_alias(request_data)
|
||||
if not alias:
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name,
|
||||
message="Ovalix guardrail error: no application_id configured and no user_api_key_alias to resolve by",
|
||||
should_wrap_with_default_message=False,
|
||||
)
|
||||
return self._no_application("no application_id configured and no user_api_key_alias to resolve by")
|
||||
regex = await self._get_app_name_regex()
|
||||
name = self._extract_application_name(alias, regex)
|
||||
if not name:
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name,
|
||||
message="Ovalix guardrail error: could not extract an application name from the api key alias",
|
||||
should_wrap_with_default_message=False,
|
||||
)
|
||||
return self._no_application("could not extract an application name from the api key alias")
|
||||
if self._enable_routing_cache:
|
||||
cached = self._routing_cache_get(name)
|
||||
if cached is not None:
|
||||
return cached
|
||||
hit, cached = self._routing_cache_get(name)
|
||||
if hit:
|
||||
return cached if cached is not None else self._no_application(f"application '{name}' was not found")
|
||||
routing = await self._resolve_via_tracker(name)
|
||||
if self._enable_routing_cache:
|
||||
self._routing_cache_put(name, routing)
|
||||
if routing is None:
|
||||
return self._no_application(f"application '{name}' was not found")
|
||||
return routing
|
||||
|
||||
async def _resolve_via_tracker(self, application_name: str) -> ResolvedRouting:
|
||||
async def _resolve_via_tracker(self, application_name: str) -> ResolvedRouting | None:
|
||||
url = f"{self._tracker_api_base}/tracking/custom_application/resolve_litellm_application"
|
||||
try:
|
||||
response = await self._async_handler.post(
|
||||
|
|
@ -532,21 +559,19 @@ class OvalixGuardrail(CustomGuardrail):
|
|||
)
|
||||
response.raise_for_status()
|
||||
body = response.json()
|
||||
routing = ResolvedRouting(
|
||||
return ResolvedRouting(
|
||||
application_id=str(body["application_id"]),
|
||||
checkpoint_id_pre=body.get("checkpoint_id_pre"),
|
||||
checkpoint_id_post=body.get("checkpoint_id_post"),
|
||||
checkpoint_id_pre_file=body.get("checkpoint_id_pre_file"),
|
||||
checkpoint_id_post_file=body.get("checkpoint_id_post_file"),
|
||||
)
|
||||
except httpx.HTTPStatusError as e:
|
||||
if e.response.status_code == _APPLICATION_NOT_FOUND_STATUS:
|
||||
return None
|
||||
raise self._routing_error(e) from e
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Ovalix routing resolution failed: %s", e)
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name,
|
||||
message=f"Ovalix guardrail error: routing resolution failed: {e!s}",
|
||||
should_wrap_with_default_message=False,
|
||||
) from e
|
||||
return routing
|
||||
raise self._routing_error(e) from e
|
||||
|
||||
@staticmethod
|
||||
def get_config_model() -> Optional[Type["GuardrailConfigModel"]]:
|
||||
|
|
|
|||
|
|
@ -38,6 +38,13 @@ class OvalixGuardrailConfigModel(GuardrailConfigModel):
|
|||
default=None,
|
||||
description="Cache discovery-mode routing resolution per api-key alias for 1 hour. Default on.",
|
||||
)
|
||||
fail_if_no_application: bool | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Fail the call when no application is configured and none is discovered. Default on; "
|
||||
"set false to let such calls through unguarded instead."
|
||||
),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def ui_friendly_name() -> str:
|
||||
|
|
|
|||
|
|
@ -1194,3 +1194,259 @@ def test_initialize_guardrail_wires_new_params(monkeypatch):
|
|||
assert guardrail._file_checkpoint_id == "file-1"
|
||||
assert guardrail._enable_routing_cache is False
|
||||
assert guardrail.guardrail_name == "ovalix"
|
||||
|
||||
|
||||
def test_fail_if_no_application_defaults_true():
|
||||
g = OvalixGuardrail(
|
||||
tracker_api_base="https://t", tracker_api_key="k", guardrail_name="o", event_hook="pre_call", default_on=True
|
||||
)
|
||||
assert g._fail_if_no_application is True
|
||||
|
||||
|
||||
def test_fail_if_no_application_from_param():
|
||||
g = OvalixGuardrail(
|
||||
tracker_api_base="https://t",
|
||||
tracker_api_key="k",
|
||||
fail_if_no_application=False,
|
||||
guardrail_name="o",
|
||||
event_hook="pre_call",
|
||||
default_on=True,
|
||||
)
|
||||
assert g._fail_if_no_application is False
|
||||
|
||||
|
||||
def test_fail_if_no_application_from_env_string(monkeypatch):
|
||||
monkeypatch.setenv("OVALIX_FAIL_IF_NO_APPLICATION", "false")
|
||||
g = OvalixGuardrail(
|
||||
tracker_api_base="https://t", tracker_api_key="k", guardrail_name="o", event_hook="pre_call", default_on=True
|
||||
)
|
||||
assert g._fail_if_no_application is False
|
||||
|
||||
|
||||
def test_explicit_param_beats_env_for_fail_if_no_application(monkeypatch):
|
||||
monkeypatch.setenv("OVALIX_FAIL_IF_NO_APPLICATION", "false")
|
||||
g = OvalixGuardrail(
|
||||
tracker_api_base="https://t",
|
||||
tracker_api_key="k",
|
||||
fail_if_no_application=True,
|
||||
guardrail_name="o",
|
||||
event_hook="pre_call",
|
||||
default_on=True,
|
||||
)
|
||||
assert g._fail_if_no_application is True
|
||||
|
||||
|
||||
def test_initialize_guardrail_wires_fail_if_no_application(monkeypatch):
|
||||
import litellm
|
||||
from litellm.proxy.guardrails.guardrail_hooks.ovalix import initialize_guardrail
|
||||
|
||||
monkeypatch.setattr(litellm.logging_callback_manager, "add_litellm_callback", lambda callback: None)
|
||||
|
||||
class _Params:
|
||||
tracker_api_base = "https://t"
|
||||
tracker_api_key = "k"
|
||||
application_id = None
|
||||
pre_checkpoint_id = None
|
||||
post_checkpoint_id = None
|
||||
file_checkpoint_id = None
|
||||
enable_routing_cache = None
|
||||
fail_if_no_application = False
|
||||
mode = "pre_call"
|
||||
default_on = True
|
||||
|
||||
guardrail = initialize_guardrail(_Params(), {"guardrail_name": "ovalix"})
|
||||
assert guardrail._fail_if_no_application is False
|
||||
|
||||
|
||||
def test_config_model_declares_fail_if_no_application():
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.ovalix import OvalixGuardrailConfigModel
|
||||
|
||||
assert OvalixGuardrailConfigModel.model_fields["fail_if_no_application"].default is None
|
||||
|
||||
|
||||
def _fail_open_discovery_guardrail(enable_cache=False):
|
||||
return OvalixGuardrail(
|
||||
tracker_api_base="https://tracker.test",
|
||||
tracker_api_key="key",
|
||||
enable_routing_cache=enable_cache,
|
||||
fail_if_no_application=False,
|
||||
guardrail_name="ovalix-test",
|
||||
event_hook="pre_call",
|
||||
default_on=True,
|
||||
)
|
||||
|
||||
|
||||
def _http_status_error(status_code):
|
||||
request = httpx.Request("POST", "https://tracker.test/x")
|
||||
response = httpx.Response(status_code, request=request)
|
||||
return httpx.HTTPStatusError("boom", request=request, response=response)
|
||||
|
||||
|
||||
def _mock_handler_resolve_status(g, status_code):
|
||||
get_resp = MagicMock()
|
||||
get_resp.json.return_value = {"regex": _REGEX}
|
||||
get_resp.raise_for_status = MagicMock()
|
||||
post_resp = MagicMock()
|
||||
post_resp.raise_for_status = MagicMock(side_effect=_http_status_error(status_code))
|
||||
g._async_handler.get = AsyncMock(return_value=get_resp)
|
||||
g._async_handler.post = AsyncMock(return_value=post_resp)
|
||||
return g._async_handler.post
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fail_open_missing_alias_returns_none():
|
||||
g = _fail_open_discovery_guardrail()
|
||||
_mock_handler(g)
|
||||
assert await g._resolve_routing({"metadata": {"user_api_key_user_email": "u@x.com"}}) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fail_open_unparseable_alias_returns_none():
|
||||
g = _fail_open_discovery_guardrail()
|
||||
_mock_handler(g)
|
||||
assert await g._resolve_routing(_alias_request_data("no brackets here")) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fail_open_tracker_404_returns_none():
|
||||
g = _fail_open_discovery_guardrail()
|
||||
_mock_handler_resolve_status(g, 404)
|
||||
assert await g._resolve_routing(_alias_request_data("[Ghost App] prod")) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fail_closed_tracker_404_still_raises():
|
||||
g = _discovery_guardrail(enable_cache=False)
|
||||
_mock_handler_resolve_status(g, 404)
|
||||
with pytest.raises(GuardrailRaisedException):
|
||||
await g._resolve_routing(_alias_request_data("[Ghost App] prod"))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fail_open_tracker_500_still_raises():
|
||||
g = _fail_open_discovery_guardrail()
|
||||
_mock_handler_resolve_status(g, 500)
|
||||
with pytest.raises(GuardrailRaisedException):
|
||||
await g._resolve_routing(_alias_request_data("[Weather App] prod"))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fail_open_tracker_unreachable_still_raises():
|
||||
g = _fail_open_discovery_guardrail()
|
||||
_mock_handler(g)
|
||||
g._async_handler.post = AsyncMock(side_effect=httpx.ConnectError("boom"))
|
||||
with pytest.raises(GuardrailRaisedException):
|
||||
await g._resolve_routing(_alias_request_data())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fail_open_regex_fetch_failure_still_raises():
|
||||
g = _fail_open_discovery_guardrail()
|
||||
g._async_handler.get = AsyncMock(side_effect=httpx.ConnectError("boom"))
|
||||
with pytest.raises(GuardrailRaisedException):
|
||||
await g._resolve_routing(_alias_request_data())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fail_open_no_application_makes_no_checkpoint_call():
|
||||
g = _fail_open_discovery_guardrail()
|
||||
_mock_handler_resolve_status(g, 404)
|
||||
checkpoint = AsyncMock()
|
||||
with patch.object(g, "_call_checkpoint", new=checkpoint):
|
||||
inputs = GenericGuardrailAPIInputs(texts=["hello"])
|
||||
result = await g.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data=_alias_request_data("[Ghost App] prod"),
|
||||
input_type="request",
|
||||
logging_obj=None,
|
||||
)
|
||||
assert result == inputs
|
||||
checkpoint.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fail_open_resolved_application_without_checkpoint_still_raises():
|
||||
g = _fail_open_discovery_guardrail()
|
||||
_mock_handler(
|
||||
g,
|
||||
routing={
|
||||
"application_id": "app-9",
|
||||
"checkpoint_id_pre": None,
|
||||
"checkpoint_id_post": None,
|
||||
"checkpoint_id_pre_file": None,
|
||||
"checkpoint_id_post_file": None,
|
||||
},
|
||||
)
|
||||
with pytest.raises(GuardrailRaisedException):
|
||||
await g.apply_guardrail(
|
||||
inputs=GenericGuardrailAPIInputs(texts=["hi"]),
|
||||
request_data=_alias_request_data(),
|
||||
input_type="request",
|
||||
logging_obj=None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fail_open_still_guards_when_application_resolves():
|
||||
g = _fail_open_discovery_guardrail()
|
||||
_mock_handler(g)
|
||||
checkpoint = AsyncMock(return_value=_BLOCK)
|
||||
with patch.object(g, "_call_checkpoint", new=checkpoint):
|
||||
with pytest.raises(OvalixGuardrailBlockedException):
|
||||
await g.apply_guardrail(
|
||||
inputs=GenericGuardrailAPIInputs(texts=["hi"]),
|
||||
request_data=_alias_request_data(),
|
||||
input_type="request",
|
||||
logging_obj=None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_negative_routing_cache_hit_and_expiry(monkeypatch):
|
||||
g = _fail_open_discovery_guardrail(enable_cache=True)
|
||||
mock_post = _mock_handler_resolve_status(g, 404)
|
||||
clock = [1000.0]
|
||||
monkeypatch.setattr("litellm.proxy.guardrails.guardrail_hooks.ovalix.ovalix.time.monotonic", lambda: clock[0])
|
||||
|
||||
assert await g._resolve_routing(_alias_request_data("[Ghost App] prod")) is None
|
||||
clock[0] = 1000.0 + 299
|
||||
assert await g._resolve_routing(_alias_request_data("[Ghost App] prod")) is None
|
||||
assert mock_post.call_count == 1
|
||||
|
||||
clock[0] = 1000.0 + 301
|
||||
assert await g._resolve_routing(_alias_request_data("[Ghost App] prod")) is None
|
||||
assert mock_post.call_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_negative_cache_expires_sooner_than_positive(monkeypatch):
|
||||
g = _discovery_guardrail(enable_cache=True)
|
||||
_mock_handler(g)
|
||||
clock = [1000.0]
|
||||
monkeypatch.setattr("litellm.proxy.guardrails.guardrail_hooks.ovalix.ovalix.time.monotonic", lambda: clock[0])
|
||||
await g._resolve_routing(_alias_request_data("[Weather App] prod"))
|
||||
clock[0] = 1000.0 + 301
|
||||
hit, cached = g._routing_cache_get("Weather App")
|
||||
assert hit is True and cached is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_negative_result_not_cached_when_cache_disabled():
|
||||
g = _fail_open_discovery_guardrail(enable_cache=False)
|
||||
mock_post = _mock_handler_resolve_status(g, 404)
|
||||
await g._resolve_routing(_alias_request_data("[Ghost App] prod"))
|
||||
await g._resolve_routing(_alias_request_data("[Ghost App] prod"))
|
||||
assert mock_post.call_count == 2
|
||||
assert len(g._routing_cache) == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cached_404_still_raises_when_failing_closed(monkeypatch):
|
||||
g = _discovery_guardrail(enable_cache=True)
|
||||
mock_post = _mock_handler_resolve_status(g, 404)
|
||||
clock = [1000.0]
|
||||
monkeypatch.setattr("litellm.proxy.guardrails.guardrail_hooks.ovalix.ovalix.time.monotonic", lambda: clock[0])
|
||||
for _ in range(2):
|
||||
with pytest.raises(GuardrailRaisedException):
|
||||
await g._resolve_routing(_alias_request_data("[Ghost App] prod"))
|
||||
assert mock_post.call_count == 1
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue