Guardrails implementation

examples/tracing/programmatic_configuration.py

@@ -123,7 +123,9 @@ def mixed_config_function(query: str) -> str:
     print("=" * 50)
 
     # Note: Replace the placeholder API keys and IDs with real values
-    print("Note: Replace placeholder API keys and pipeline IDs with real values before running.")
+    print(
+        "Note: Replace placeholder API keys and pipeline IDs with real values before running."
+    )
     print()
 
     try:

src/openlayer/lib/guardrails/__init__.py

@@ -0,0 +1,17 @@
+"""Guardrails module for Openlayer tracing."""
+
+from .base import (
+    GuardrailAction,
+    BlockStrategy,
+    GuardrailResult,
+    BaseGuardrail,
+    GuardrailBlockedException,
+)
+
+__all__ = [
+    "GuardrailAction",
+    "BlockStrategy",
+    "GuardrailResult",
+    "BaseGuardrail",
+    "GuardrailBlockedException",
+]

src/openlayer/lib/guardrails/base.py

@@ -0,0 +1,118 @@
+"""Base classes and interfaces for guardrails system."""
+
+import abc
+import enum
+import logging
+from typing import Any, Dict, Optional
+from dataclasses import dataclass
+
+logger = logging.getLogger(__name__)
+
+
+class GuardrailAction(enum.Enum):
+    """Actions that a guardrail can take."""
+
+    ALLOW = "allow"
+    BLOCK = "block"
+    MODIFY = "modify"
+
+
+class BlockStrategy(enum.Enum):
+    """Strategies for handling blocked requests."""
+
+    RAISE_EXCEPTION = (
+        "raise_exception"  # Raise GuardrailBlockedException (breaks pipeline)
+    )
+    RETURN_EMPTY = "return_empty"  # Return empty/None response (graceful)
+    RETURN_ERROR_MESSAGE = "return_error_message"  # Return error message (graceful)
+    SKIP_FUNCTION = "skip_function"  # Skip function execution, return None (graceful)
+
+
+@dataclass
+class GuardrailResult:
+    """Result of applying a guardrail."""
+
+    action: GuardrailAction
+    modified_data: Optional[Any] = None
+    metadata: Optional[Dict[str, Any]] = None
+    reason: Optional[str] = None
+    block_strategy: Optional[BlockStrategy] = None
+    error_message: Optional[str] = None
+
+    def __post_init__(self):
+        """Validate the result after initialization."""
+        if self.action == GuardrailAction.MODIFY and self.modified_data is None:
+            raise ValueError("modified_data must be provided when action is MODIFY")
+        if self.action == GuardrailAction.BLOCK and self.block_strategy is None:
+            self.block_strategy = (
+                BlockStrategy.RAISE_EXCEPTION
+            )  # Default to existing behavior
+
+
+class GuardrailBlockedException(Exception):
+    """Exception raised when a guardrail blocks execution."""
+
+    def __init__(
+        self,
+        guardrail_name: str,
+        reason: str,
+        metadata: Optional[Dict[str, Any]] = None,
+    ):
+        self.guardrail_name = guardrail_name
+        self.reason = reason
+        self.metadata = metadata or {}
+        super().__init__(f"Guardrail '{guardrail_name}' blocked execution: {reason}")
+
+
+class BaseGuardrail(abc.ABC):
+    """Base class for all guardrails."""
+
+    def __init__(self, name: str, enabled: bool = True, **config):
+        """Initialize the guardrail.
+
+        Args:
+            name: Human-readable name for this guardrail
+            enabled: Whether this guardrail is active
+            **config: Guardrail-specific configuration
+        """
+        self.name = name
+        self.enabled = enabled
+        self.config = config
+
+    @abc.abstractmethod
+    def check_input(self, inputs: Dict[str, Any]) -> GuardrailResult:
+        """Check and potentially modify function inputs.
+
+        Args:
+            inputs: Dictionary of function inputs (parameter_name -> value)
+
+        Returns:
+            GuardrailResult indicating the action to take
+        """
+        pass
+
+    @abc.abstractmethod
+    def check_output(self, output: Any, inputs: Dict[str, Any]) -> GuardrailResult:
+        """Check and potentially modify function output.
+
+        Args:
+            output: The function's output
+            inputs: Dictionary of function inputs for context
+
+        Returns:
+            GuardrailResult indicating the action to take
+        """
+        pass
+
+    def is_enabled(self) -> bool:
+        """Check if this guardrail is enabled."""
+        return self.enabled
+
+    def get_metadata(self) -> Dict[str, Any]:
+        """Get metadata about this guardrail for trace logging."""
+        return {
+            "name": self.name,
+            "type": self.__class__.__name__,
+            "enabled": self.enabled,
+            "config": self.config,
+        }

src/openlayer/lib/tracing/enums.py

@@ -4,9 +4,10 @@
 
 
 class StepType(enum.Enum):
-    USER_CALL = "user_call"
-    CHAT_COMPLETION = "chat_completion"
     AGENT = "agent"
+    CHAT_COMPLETION = "chat_completion"
+    GUARDRAIL = "guardrail"
+    HANDOFF = "handoff"
     RETRIEVER = "retriever"
     TOOL = "tool"
-    HANDOFF = "handoff"
+    USER_CALL = "user_call"

src/openlayer/lib/tracing/steps.py

@@ -2,7 +2,7 @@
 
 import time
 import uuid
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, Optional, List
 
 from .. import utils
 from . import enums
@@ -229,6 +229,39 @@ def to_dict(self) -> Dict[str, Any]:
         return step_dict
 
 
+class GuardrailStep(Step):
+    """Step for tracking guardrail execution."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.step_type = enums.StepType.GUARDRAIL
+        self.action: Optional[str] = None
+        self.blocked_entities: Optional[List[str]] = None
+        self.confidence_threshold: float = None
+        self.reason: Optional[str] = None
+        self.detected_entities: Optional[List[str]] = None
+        self.redacted_entities: Optional[List[str]] = None
+        self.block_strategy: Optional[str] = None
+        self.data_type: Optional[str] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Dictionary representation of the GuardrailStep."""
+        step_dict = super().to_dict()
+        step_dict.update(
+            {
+                "action": self.action,
+                "blockedEntities": self.blocked_entities,
+                "confidenceThreshold": self.confidence_threshold,
+                "reason": self.reason,
+                "detectedEntities": self.detected_entities,
+                "blockStrategy": self.block_strategy,
+                "redactedEntities": self.redacted_entities,
+                "dataType": self.data_type,
+            }
+        )
+        return step_dict
+
+
 # ----------------------------- Factory function ----------------------------- #
 def step_factory(step_type: enums.StepType, *args, **kwargs) -> Step:
     """Factory function to create a step based on the step_type."""
@@ -241,5 +274,6 @@ def step_factory(step_type: enums.StepType, *args, **kwargs) -> Step:
         enums.StepType.RETRIEVER: RetrieverStep,
         enums.StepType.TOOL: ToolStep,
         enums.StepType.HANDOFF: HandoffStep,
+        enums.StepType.GUARDRAIL: GuardrailStep,
     }
     return step_type_mapping[step_type](*args, **kwargs)