"""
Pipeline Executor - Executes guardrail pipelines with conditional step logic.
Runs guardrails sequentially per pipeline step definitions, handling
pass/fail actions (allow, block, next, modify_response) and data forwarding.
"""
import time
from typing import Any, List, Optional
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
ModifyResponseException,
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
UnifiedLLMGuardrails,
)
from litellm.types.proxy.policy_engine.pipeline_types import (
PipelineExecutionResult,
PipelineStep,
PipelineStepResult,
)
try:
from fastapi.exceptions import HTTPException
except ImportError:
HTTPException = None # type: ignore
class PipelineExecutor:
"""Executes guardrail pipelines with ordered, conditional step logic."""
@staticmethod
async def execute_steps(
steps: List[PipelineStep],
mode: str,
data: dict,
user_api_key_dict: Any,
call_type: str,
policy_name: str,
) -> PipelineExecutionResult:
"""
Execute pipeline steps sequentially with conditional actions.
Args:
steps: Ordered list of pipeline steps
mode: Event hook mode (pre_call, post_call)
data: Request data dict
user_api_key_dict: User API key auth
call_type: Type of call (completion, etc.)
policy_name: Name of the owning policy (for logging)
Returns:
PipelineExecutionResult with terminal action and step results
"""
step_results: List[PipelineStepResult] = []
working_data = data.copy()
if "metadata" in working_data:
working_data["metadata"] = working_data["metadata"].copy()
for i, step in enumerate(steps):
start_time = time.perf_counter()
outcome, modified_data, error_detail = await PipelineExecutor._run_step(
step=step,
mode=mode,
data=working_data,
user_api_key_dict=user_api_key_dict,
call_type=call_type,
)
duration = time.perf_counter() - start_time
action = step.on_pass if outcome == "pass" else step.on_fail
step_result = PipelineStepResult(
guardrail_name=step.guardrail,
outcome=outcome,
action_taken=action,
modified_data=modified_data,
error_detail=error_detail,
duration_seconds=round(duration, 4),
)
step_results.append(step_result)
verbose_proxy_logger.debug(
f"Pipeline '{policy_name}' step {i}: guardrail={step.guardrail}, "
f"outcome={outcome}, action={action}"
)
# Forward modified data to next step if pass_data is True
if step.pass_data and modified_data is not None:
working_data = {**working_data, **modified_data}
# Handle terminal actions
if action == "allow":
return PipelineExecutionResult(
terminal_action="allow",
step_results=step_results,
modified_data=working_data if working_data != data else None,
)
if action == "block":
return PipelineExecutionResult(
terminal_action="block",
step_results=step_results,
error_message=error_detail,
)
if action == "modify_response":
return PipelineExecutionResult(
terminal_action="modify_response",
step_results=step_results,
modify_response_message=step.modify_response_message or error_detail,
)
# action == "next" → continue to next step
# Ran out of steps without a terminal action → default allow
return PipelineExecutionResult(
terminal_action="allow",
step_results=step_results,
modified_data=working_data if working_data != data else None,
)
@staticmethod
async def _run_step(
step: PipelineStep,
mode: str,
data: dict,
user_api_key_dict: Any,
call_type: str,
) -> tuple:
"""
Run a single pipeline step's guardrail.
Returns:
Tuple of (outcome, modified_data, error_detail) where:
- outcome: "pass", "fail", or "error"
- modified_data: dict if guardrail returned modified data, else None
- error_detail: error message string if fail/error, else None
"""
callback = PipelineExecutor._find_guardrail_callback(step.guardrail)
if callback is None:
verbose_proxy_logger.warning(
f"Pipeline: guardrail '{step.guardrail}' not found in callbacks"
)
return ("error", None, f"Guardrail '{step.guardrail}' not found")
try:
# Inject guardrail name into metadata so should_run_guardrail() allows it
if "metadata" not in data:
data["metadata"] = {}
data["metadata"]["guardrails"] = [step.guardrail]
# Use unified_guardrail path if callback implements apply_guardrail
target: CustomLogger = callback
use_unified = "apply_guardrail" in type(callback).__dict__
if use_unified:
data["guardrail_to_apply"] = callback
target = UnifiedLLMGuardrails()
if mode == "pre_call":
response = await target.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=None, # type: ignore
data=data,
call_type=call_type, # type: ignore
)
elif mode == "post_call":
response = await target.async_post_call_success_hook(
user_api_key_dict=user_api_key_dict,
data=data,
response=data.get("response"), # type: ignore
)
else:
return ("error", None, f"Unsupported pipeline mode: {mode}")
# Normal return means pass
modified_data = None
if response is not None and isinstance(response, dict):
modified_data = response
return ("pass", modified_data, None)
except Exception as e:
if CustomGuardrail._is_guardrail_intervention(e):
error_msg = _extract_error_message(e)
return ("fail", None, error_msg)
else:
verbose_proxy_logger.error(
f"Pipeline: unexpected error from guardrail '{step.guardrail}': {e}"
)
return ("error", None, str(e))
@staticmethod
def _find_guardrail_callback(guardrail_name: str) -> Optional[CustomGuardrail]:
"""Look up an initialized guardrail callback by name from litellm.callbacks."""
for callback in litellm.callbacks:
if isinstance(callback, CustomGuardrail):
if callback.guardrail_name == guardrail_name:
return callback
return None
def _extract_error_message(e: Exception) -> str:
"""Extract a human-readable error message from a guardrail exception."""
if isinstance(e, ModifyResponseException):
return str(e)
if HTTPException is not None and isinstance(e, HTTPException):
detail = getattr(e, "detail", None)
if detail:
return str(detail)
return str(e)