Add pipeline executor for sequential guardrail execution

litellm/proxy/policy_engine/pipeline_executor.py

@@ -0,0 +1,208 @@
+"""
+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.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()
+
+        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:
+            # Use unified_guardrail path if callback implements apply_guardrail
+            target = 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)