mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
fix(policy_engine): keep a policy-added guardrail's other stages when a pipeline steps it
A policy that both adds a guardrail and steps it in a post_call pipeline used to drop the guardrail from the request's guardrail list outright, so its pre_call stage never ran. The per-hook loops already skip guardrails by pipeline mode, so the mode-agnostic subtraction only lost coverage
This commit is contained in:
parent
91fc1b2010
commit
d08a177bc7
2 changed files with 44 additions and 5 deletions
|
|
@ -3079,10 +3079,9 @@ def _apply_resolved_guardrails_to_metadata(
|
|||
if metadata_variable_name not in data:
|
||||
data[metadata_variable_name] = {}
|
||||
|
||||
# Track pipeline-managed guardrails to exclude from independent execution
|
||||
pipeline_managed_guardrails: set = set()
|
||||
# Record the pipelines and the guardrails they step; the hook loops skip those per pipeline mode
|
||||
if pipelines:
|
||||
pipeline_managed_guardrails = PolicyResolver.get_pipeline_managed_guardrails(pipelines)
|
||||
pipeline_managed_guardrails: Final = PolicyResolver.get_pipeline_managed_guardrails(pipelines)
|
||||
data[metadata_variable_name]["_guardrail_pipelines"] = pipelines
|
||||
data[metadata_variable_name]["_pipeline_managed_guardrails"] = pipeline_managed_guardrails
|
||||
verbose_proxy_logger.debug(
|
||||
|
|
@ -3099,10 +3098,8 @@ def _apply_resolved_guardrails_to_metadata(
|
|||
existing_guardrails = []
|
||||
|
||||
# Combine existing guardrails with policy-resolved guardrails (no duplicates)
|
||||
# Exclude pipeline-managed guardrails from the flat list
|
||||
combined = set(existing_guardrails)
|
||||
combined.update(resolved_guardrails)
|
||||
combined -= pipeline_managed_guardrails
|
||||
data[metadata_variable_name]["guardrails"] = list(combined)
|
||||
|
||||
verbose_proxy_logger.debug("Policy engine: added guardrails to request metadata: %s", list(combined))
|
||||
|
|
|
|||
|
|
@ -4148,6 +4148,48 @@ async def test_add_guardrails_from_policy_engine():
|
|||
attachment_registry._initialized = False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_guardrails_from_policy_engine_keeps_a_policy_added_guardrail_its_pipeline_also_steps():
|
||||
from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry
|
||||
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
|
||||
from litellm.types.proxy.policy_engine import (
|
||||
GuardrailPipeline,
|
||||
PipelineStep,
|
||||
Policy,
|
||||
PolicyAttachment,
|
||||
PolicyGuardrails,
|
||||
)
|
||||
|
||||
data = {"model": "gpt-4", "messages": [{"role": "user", "content": "Hello"}], "metadata": {}}
|
||||
policy_registry = get_policy_registry()
|
||||
policy_registry._policies = {
|
||||
"response-governance": Policy(
|
||||
guardrails=PolicyGuardrails(add=["pii_blocker"]),
|
||||
pipeline=GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="pii_blocker")]),
|
||||
),
|
||||
}
|
||||
policy_registry._initialized = True
|
||||
attachment_registry = get_attachment_registry()
|
||||
attachment_registry._attachments = [PolicyAttachment(policy="response-governance", scope="*")]
|
||||
attachment_registry._initialized = True
|
||||
|
||||
try:
|
||||
await add_guardrails_from_policy_engine(
|
||||
data=data,
|
||||
metadata_variable_name="metadata",
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test-key"),
|
||||
)
|
||||
finally:
|
||||
policy_registry._policies = {}
|
||||
policy_registry._initialized = False
|
||||
attachment_registry._attachments = []
|
||||
attachment_registry._initialized = False
|
||||
|
||||
assert data["metadata"]["guardrails"] == ["pii_blocker"]
|
||||
assert data["metadata"]["_pipeline_managed_guardrails"] == {"pii_blocker"}
|
||||
assert [pipeline.mode for _policy_name, pipeline in data["metadata"]["_guardrail_pipelines"]] == ["post_call"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_guardrails_from_policy_engine_accepts_dynamic_policies_and_pops_from_data():
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue