mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
fix(policy_engine): skip pipeline-managed guardrails in the response-path guardrail loop
This commit is contained in:
parent
e6edd62f5d
commit
aeac6a412c
2 changed files with 34 additions and 1 deletions
|
|
@ -2825,6 +2825,7 @@ class ProxyLogging:
|
|||
if pipeline_response is not None:
|
||||
response = pipeline_response # rebind-ok: adopt the pipeline's replacement response, same contract as the callback loops below
|
||||
|
||||
pipeline_managed: Final = _pipeline_managed_guardrail_names(data)
|
||||
guardrail_callbacks: Final[list[CustomGuardrail]] = []
|
||||
other_callbacks: Final[list[CustomLogger]] = []
|
||||
try:
|
||||
|
|
@ -2849,12 +2850,18 @@ class ProxyLogging:
|
|||
guardrail_data: Final = _check_and_merge_model_level_guardrails(data=data, llm_router=llm_router)
|
||||
|
||||
parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = tuple(
|
||||
callback for callback in guardrail_callbacks if getattr(callback, "run_in_parallel", False)
|
||||
callback
|
||||
for callback in guardrail_callbacks
|
||||
if getattr(callback, "run_in_parallel", False)
|
||||
and not (callback.guardrail_name and callback.guardrail_name in pipeline_managed)
|
||||
)
|
||||
|
||||
for callback in guardrail_callbacks:
|
||||
# Main - V2 Guardrails implementation
|
||||
|
||||
if callback.guardrail_name and callback.guardrail_name in pipeline_managed:
|
||||
continue
|
||||
|
||||
if getattr(callback, "run_in_parallel", False):
|
||||
continue
|
||||
|
||||
|
|
|
|||
|
|
@ -951,6 +951,32 @@ async def test_post_call_pipeline_pass_runs_once_and_leaves_request_data_untouch
|
|||
assert "guardrails" not in data["metadata"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_pipeline_managed_default_on_guardrail_runs_exactly_once(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch
|
||||
):
|
||||
seen: Dict[str, Any] = {"count": 0}
|
||||
|
||||
class CountingGuardrail(CustomGuardrail):
|
||||
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
|
||||
seen["count"] += 1
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"callbacks",
|
||||
[CountingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True)],
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
data = _post_call_pipeline_data()
|
||||
|
||||
await proxy_logging.post_call_success_hook(
|
||||
data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert seen["count"] == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_pipeline_replacement_response_reaches_caller(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue