mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
fix(guardrails): don't add post_call output scan for MCP-only Presidio modes
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
fbed17d567
commit
7dfda63bb7
2 changed files with 62 additions and 1 deletions
|
|
@ -86,12 +86,32 @@ def initialize_lakera_v2(litellm_params: LitellmParams, guardrail: Guardrail):
|
|||
return _lakera_v2_callback
|
||||
|
||||
|
||||
_MCP_EVENT_HOOKS: Final = frozenset(
|
||||
{
|
||||
GuardrailEventHooks.pre_mcp_call.value,
|
||||
GuardrailEventHooks.during_mcp_call.value,
|
||||
GuardrailEventHooks.post_mcp_call.value,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _is_mcp_only_mode(mode: str | list[str] | Mode) -> bool:
|
||||
match mode:
|
||||
case str():
|
||||
return mode in _MCP_EVENT_HOOKS
|
||||
case list():
|
||||
return bool(mode) and all(m in _MCP_EVENT_HOOKS for m in mode)
|
||||
case Mode():
|
||||
return False
|
||||
|
||||
|
||||
def initialize_presidio(litellm_params: LitellmParams, guardrail: Guardrail) -> tuple[CustomGuardrail, ...]:
|
||||
from litellm.proxy.guardrails.guardrail_hooks.presidio import (
|
||||
_OPTIONAL_PresidioPIIMasking,
|
||||
)
|
||||
|
||||
filter_scope: Final = getattr(litellm_params, "presidio_filter_scope", None) or "both"
|
||||
explicit_filter_scope: Final = getattr(litellm_params, "presidio_filter_scope", None)
|
||||
filter_scope: Final = explicit_filter_scope or ("input" if _is_mcp_only_mode(litellm_params.mode) else "both")
|
||||
run_input: Final = filter_scope in ("input", "both")
|
||||
run_output: Final = filter_scope in ("output", "both")
|
||||
|
||||
|
|
|
|||
|
|
@ -156,6 +156,47 @@ def test_initialize_presidio_forwards_analyze_chunk_size_bytes():
|
|||
assert initialized[-1].presidio_analyze_chunk_size_bytes == 250_000
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"mode, filter_scope, expected_hooks",
|
||||
[
|
||||
("pre_mcp_call", None, {"pre_mcp_call"}),
|
||||
(["pre_mcp_call", "post_mcp_call"], None, {"pre_mcp_call", "post_mcp_call"}),
|
||||
("pre_mcp_call", "both", {"pre_mcp_call", "post_call"}),
|
||||
("pre_call", None, {"pre_call", "post_call"}),
|
||||
],
|
||||
)
|
||||
def test_initialize_presidio_mcp_mode_does_not_add_post_call_scan(mode, filter_scope, expected_hooks):
|
||||
"""Regression: a `pre_mcp_call` Presidio guardrail used to also register a
|
||||
`post_call` output scanner, so a blocked tool call that the model mentioned in
|
||||
its answer turned the whole request into an HTTP 400 instead of a 200."""
|
||||
import litellm
|
||||
from litellm.proxy.guardrails.guardrail_hooks.presidio import (
|
||||
_OPTIONAL_PresidioPIIMasking,
|
||||
)
|
||||
|
||||
guardrail_name = f"test_presidio_mcp_scope_{id(mode)}_{filter_scope}"
|
||||
litellm_params = {
|
||||
"guardrail": SupportedGuardrailIntegrations.PRESIDIO.value,
|
||||
"mode": mode,
|
||||
"presidio_analyzer_api_base": "https://fakelink.com/v1/presidio/analyze",
|
||||
"presidio_anonymizer_api_base": "https://fakelink.com/v1/presidio/anonymize",
|
||||
}
|
||||
if filter_scope is not None:
|
||||
litellm_params["presidio_filter_scope"] = filter_scope
|
||||
|
||||
InMemoryGuardrailHandler().initialize_guardrail(
|
||||
guardrail={"guardrail_name": guardrail_name, "litellm_params": litellm_params}
|
||||
)
|
||||
|
||||
registered_hooks = {
|
||||
hook
|
||||
for callback in litellm.callbacks
|
||||
if isinstance(callback, _OPTIONAL_PresidioPIIMasking) and callback.guardrail_name == guardrail_name
|
||||
for hook in ([callback.event_hook] if isinstance(callback.event_hook, str) else callback.event_hook)
|
||||
}
|
||||
assert registered_hooks == expected_hooks
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"config_value, expected",
|
||||
[(True, True), (False, False), (None, False)],
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue