rasa-pro 3.14.0.dev7__py3-none-any.whl → 3.14.0.dev8__py3-none-any.whl

rasa/builder/guardrails/lakera.py

@@ -0,0 +1,206 @@
+"""Lakera AI guardrails provider implementation."""
+
+import asyncio
+import os
+import time
+from contextlib import asynccontextmanager
+from typing import Any, AsyncGenerator, Dict, Optional
+
+import aiohttp
+import structlog
+
+from rasa.builder import config
+from rasa.builder.guardrails.constants import (
+    LAKERA_API_KEY_ENV_VAR,
+    LAKERA_GUARD_ENDPOINT,
+    LAKERA_GUARD_RESULTS_ENDPOINT,
+)
+from rasa.builder.guardrails.exceptions import GuardrailsError
+from rasa.builder.guardrails.models import (
+    GuardrailResponse,
+    LakeraGuardrailRequest,
+    LakeraGuardrailResponse,
+)
+
+structlogger = structlog.get_logger()
+
+
+class LakeraAIGuardrails:
+    """Guardrails provider using Lakera AI."""
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        base_url: Optional[str] = None,
+    ):
+        """Initialize Lakera guardrails provider.
+
+        Args:
+            severity_threshold: The minimum severity level for the guardrail to flag
+                the user request as unsafe.
+            api_key: Lakera AI API key.
+            base_url: Optional base URL for the API. If not provided, the default
+                Lakera API URL (https://api.lakera.ai/v2) will be used.
+        """
+        self.base_url: str = base_url or config.LAKERA_BASE_URL
+        self._api_key: Optional[str] = api_key or os.getenv(LAKERA_API_KEY_ENV_VAR)
+        self._session: Optional[aiohttp.ClientSession] = None
+
+    @property
+    def guard_endpoint(self) -> str:
+        """Get the guard endpoint for the Lakera API."""
+        return f"{self.base_url}/{LAKERA_GUARD_ENDPOINT}"
+
+    @property
+    def guard_results_endpoint(self) -> str:
+        """Get the guard results endpoint for the Lakera API."""
+        return f"{self.base_url}/{LAKERA_GUARD_RESULTS_ENDPOINT}"
+
+    @asynccontextmanager
+    async def _get_session(self) -> AsyncGenerator[aiohttp.ClientSession, None]:
+        """Create a fresh ClientSession, yield it, and always close it."""
+        session = aiohttp.ClientSession(headers=self._get_headers())
+        structlogger.debug("lakera_guardrails._get_session", base_url=self.base_url)
+
+        try:
+            yield session
+        except Exception as e:
+            structlogger.error("lakera_guardrails.session_error", error=str(e))
+            raise
+        finally:
+            try:
+                await session.close()
+            except Exception as exc:
+                structlogger.warning(
+                    "lakera_guardrails.session_close_error",
+                    event_info="Failed to close aiohttp client session cleanly.",
+                    error=str(exc),
+                )
+
+    def _get_headers(self) -> Dict[str, str]:
+        """Get the headers for the Lakera API request.
+
+        Returns:
+            A dictionary containing the Authorization header with the API key.
+        """
+        using_proxy = bool(config.HELLO_LLM_PROXY_BASE_URL)
+
+        if using_proxy:
+            if not config.RASA_PRO_LICENSE:
+                raise GuardrailsError(
+                    "HELLO_LLM_PROXY_BASE_URL is set but RASA_PRO_LICENSE is missing. "
+                    "Proxy requires a Rasa Pro license token for authentication."
+                )
+            return {"Authorization": f"Bearer {config.RASA_PRO_LICENSE}"}
+
+        if not self._api_key:
+            raise GuardrailsError(
+                "LAKERA_API_KEY is missing. Provide it via env LAKERA_API_KEY or "
+                "pass api_key= to LakeraAIGuardrails."
+            )
+        return {"Authorization": f"Bearer {self._api_key}"}
+
+    async def send_request(self, request: LakeraGuardrailRequest) -> GuardrailResponse:
+        """Send a request to the Lakera API.
+
+        Args:
+            request: The guardrail request to send to the Lakera API.
+
+        Returns:
+            GuardrailResponse with the results of the check.
+
+        Raises:
+            GuardrailsError: If the request times out or returns a non-200 status code.
+            Exception: If the request fails for any other reason.
+        """
+        start_time = time.time()
+        try:
+            async with self._get_session() as session:
+                raw_response = await self._send_http_request(session, request)
+                response = LakeraGuardrailResponse.from_raw_response(
+                    raw_response,
+                    hello_rasa_user_id=request.hello_rasa_user_id,
+                    hello_rasa_project_id=request.hello_rasa_project_id,
+                )
+                processing_time_ms = (time.time() - start_time) * 1000
+                response.processing_time_ms = processing_time_ms
+                return response
+
+        # Propagate the GuardrailsError if it occurs.
+        except GuardrailsError as e:
+            raise e
+
+        except asyncio.TimeoutError:
+            message = "Lakera API request timed out."
+            structlogger.error(
+                "lakera_guardrails.send_request.timeout_error",
+                event_info=message,
+                processing_time_ms=(time.time() - start_time) * 1000,
+            )
+            raise GuardrailsError(message)
+
+        # Propagate the unexpected exceptions.
+        except Exception as e:
+            message = "Lakera API request failed."
+            structlogger.error(
+                "lakera_guardrails.send_request.unexpected_error",
+                event_info="Lakera API request failed.",
+                request=request,
+                error=e,
+                processing_time_ms=(time.time() - start_time) * 1000,
+            )
+            raise e
+
+    async def _send_http_request(
+        self, session: aiohttp.ClientSession, request: LakeraGuardrailRequest
+    ) -> Dict[str, Any]:
+        """Make an HTTP request to the Lakera API.
+
+        Args:
+            session: The aiohttp session to use for the request.
+            request: The guardrail request to send.
+
+        Returns:
+            The raw JSON response from the API.
+
+        Raises:
+            GuardrailsError: If the request fails or response parsing fails.
+        """
+        # Log the request details for debugging
+        json_payload = request.to_json_payload()
+        structlogger.debug(
+            "lakera_guardrails.send_request.request",
+            url=self.guard_endpoint,
+            method="POST",
+            request_body=json_payload,
+        )
+
+        async with session.post(
+            self.guard_endpoint,
+            json=json_payload,
+        ) as client_response:
+            # Check if the response is successful. If not, raise an error.
+            if client_response.status >= 400:
+                error_text = await client_response.text()
+                message = (
+                    f"Lakera API request failed with status "
+                    f"`{client_response.status}`. Error: "
+                    f"`{error_text}`."
+                )
+                structlogger.error(
+                    "lakera_guardrails.send_request.http_error",
+                    event_info=message,
+                    url=self.guard_endpoint,
+                    status=client_response.status,
+                    error=error_text,
+                    request_body=json_payload,
+                )
+                raise GuardrailsError(message)
+
+            # Parse the response as a dictionary.
+            raw_response = await client_response.json()
+            structlogger.debug(
+                "lakera_guardrails.send_request.response",
+                response_body=raw_response,
+            )
+            return raw_response

rasa/builder/guardrails/models.py

@@ -0,0 +1,231 @@
+"""Models for guardrails system."""
+
+import copy
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class GuardrailType(Enum):
+    """Types of guardrails that can be applied with Lakera AI."""
+
+    PROMPT_ATTACK = "prompt_attack"
+    CONTENT_VIOLATION = "content_violation"
+    DATA_LEAKAGE = "data_leakage"
+    MALICIOUS_LINKS = "malicious_content"
+    CUSTOM = "custom"
+    OTHER = "other"
+
+
+class GuardrailRequest(BaseModel, ABC):
+    """Request for guardrails check."""
+
+    hello_rasa_user_id: Optional[str] = Field(
+        default=None,
+        description="Required. User identifier for the Hello Rasa project. ",
+    )
+    hello_rasa_project_id: Optional[str] = Field(
+        default=None,
+        description="Required. Project identifier for the Hello Rasa project. ",
+    )
+    metadata: Optional[Dict[str, Any]] = Field(
+        default=None, description="Additional metadata for the guardrails endpoint."
+    )
+
+    @abstractmethod
+    def to_json_payload(self) -> Dict[str, Any]:
+        """Convert the request to a JSON payload."""
+        ...
+
+
+class LakeraGuardrailRequest(GuardrailRequest):
+    """Request for Lakera guardrails check."""
+
+    lakera_project_id: str = Field(
+        description="Required. Project identifier for the Lakera AI project."
+    )
+    payload: bool = Field(
+        default=True,
+        description=(
+            "From Lakera AI: When true the response will return a payload object "
+            "containing any PII, profanity or custom detector regex matches detected, "
+            "along with their location within the contents."
+        ),
+    )
+    breakdown: bool = Field(
+        default=True,
+        description=(
+            "From Lakera AI: When true the response will return a breakdown list of "
+            "the detectors that were run, as defined in the policy, and whether each "
+            "of them detected something or not."
+        ),
+    )
+
+    messages: List[Dict[str, Any]] = Field(
+        description=(
+            "Required. From Lakera AI: List of messages comprising the interaction "
+            "history with the LLM in OpenAI API Chat Completions format. Can be "
+            "multiple messages of any role: user, assistant, system, tool, or "
+            "developer."
+        ),
+    )
+
+    def to_json_payload(self) -> Dict[str, Any]:
+        """Convert the request to a JSON payload to be sent to the Lakera endpoint."""
+        metadata = self.metadata or {}
+        metadata["hello_rasa_project_id"] = self.hello_rasa_project_id
+        metadata["hello_rasa_user_id"] = self.hello_rasa_user_id
+
+        json_payload: Dict[str, Any] = {
+            "messages": self.messages,
+            "project_id": self.lakera_project_id,
+            "metadata": metadata,
+        }
+
+        if self.payload:
+            json_payload["payload"] = self.payload
+        if self.breakdown:
+            json_payload["breakdown"] = self.breakdown
+
+        return json_payload
+
+
+class GuardrailDetection(BaseModel):
+    """Represents a single guardrail detection."""
+
+    type: GuardrailType = Field(description="Type of guardrail detection.")
+    original_type: str = Field(description="Original detection from the provider.")
+    metadata: Optional[Dict[str, Any]] = Field(
+        default=None, description="Additional metadata about the detection itself."
+    )
+
+
+class GuardrailResponse(BaseModel):
+    """Response from guardrails system."""
+
+    hello_rasa_user_id: Optional[str] = Field(
+        default=None,
+        description="Required. User identifier for the Hello Rasa project. ",
+    )
+    hello_rasa_project_id: Optional[str] = Field(
+        default=None,
+        description="Required. Project identifier for the Hello Rasa project. ",
+    )
+    flagged: bool = Field(description="Whether any policy violations were detected.")
+    detections: List[GuardrailDetection] = Field(
+        default_factory=list, description="List of detected policy violations."
+    )
+    processing_time_ms: Optional[float] = Field(
+        default=None, description="Processing time in milliseconds."
+    )
+    metadata: Optional[Dict[str, Any]] = Field(
+        default=None, description="Additional metadata from the provider."
+    )
+
+
+class LakeraGuardrailResponse(GuardrailResponse):
+    """Response from Lakera AI `/guard` endpoint."""
+
+    @classmethod
+    def from_raw_response(
+        cls,
+        raw_response: Dict[str, Any],
+        hello_rasa_user_id: str,
+        hello_rasa_project_id: str,
+    ) -> "LakeraGuardrailResponse":
+        """Create a LakeraGuardrailResponse from a response."""
+        from rasa.builder.guardrails.utils import (
+            map_lakera_detector_type_to_guardrail_type,
+        )
+
+        # Get the basic information from the response and create the response object.
+        flagged = raw_response.get("flagged", False)
+        metadata = raw_response.get("metadata")
+        response = cls(
+            flagged=flagged,
+            metadata=metadata,
+            hello_rasa_user_id=hello_rasa_user_id,
+            hello_rasa_project_id=hello_rasa_project_id,
+        )
+
+        # If the response is not flagged, return the response object.
+        if not flagged:
+            return response
+
+        # If the response is flagged, parse the breakdown section.
+        breakdown = raw_response.get("breakdown", [])
+
+        # Parse the breakdown.
+        detections: List[GuardrailDetection] = []
+        for detector in breakdown:
+            if detector.get("detected", True):
+                detector_type = detector.get("detector_type")
+                rasa_detection_type = map_lakera_detector_type_to_guardrail_type(
+                    detector_type,
+                )
+                if not rasa_detection_type:
+                    continue
+                # Remove the detector type and the detected flag from the detector.
+                # And keep the rest as part of the metadata. In case of a Lakera this
+                # will include the message_id, project_id, policy_id, detector_id, etc.
+                metadata = copy.deepcopy(detector)
+                metadata.pop("detected")
+                metadata.pop("detector_type")
+
+                detections.append(
+                    GuardrailDetection(
+                        type=rasa_detection_type,
+                        original_type=detector_type,
+                        metadata=metadata,
+                    )
+                )
+
+        # If there are detections, add them to the response.
+        if detections:
+            response.detections = detections
+
+        return response
+
+
+class GuardrailRequestKey(BaseModel):
+    user_text: str
+    hello_rasa_user_id: str = ""
+    hello_rasa_project_id: str = ""
+    lakera_project_id: str
+
+    # hashable by value
+    model_config = ConfigDict(frozen=True)
+
+
+class ScopeState(BaseModel):
+    blocked_until: Optional[float] = Field(
+        default=None,
+        description="UNIX timestamp in seconds until this scope is blocked.",
+    )
+    violations: List[float] = Field(
+        default_factory=list,
+        description="UNIX timestamps in seconds when violations occurred.",
+    )
+
+    def is_blocked(self, now: float) -> bool:
+        """Return whether the scope is currently blocked.
+
+        Args:
+            now: Current time as a UNIX timestamp.
+
+        Returns:
+            True if blocked_until is set and in the future, else False.
+        """
+        return self.blocked_until is not None and now < self.blocked_until
+
+
+class ProjectState(BaseModel):
+    project: ScopeState = Field(default_factory=ScopeState)
+    users: Dict[str, ScopeState] = Field(default_factory=dict)
+
+
+class BlockResult(BaseModel):
+    user_blocked_now: bool = False
+    project_blocked_now: bool = False

rasa/builder/guardrails/store.py

@@ -0,0 +1,238 @@
+import asyncio
+import time
+from typing import Optional
+
+import structlog
+
+from rasa.builder import config
+from rasa.builder.guardrails.constants import (
+    BLOCK_SCOPE_PROJECT,
+    BLOCK_SCOPE_USER,
+    BlockScope,
+)
+from rasa.builder.guardrails.models import BlockResult, ProjectState, ScopeState
+
+structlogger = structlog.get_logger()
+
+
+class GuardrailsInMemoryStore:
+    """In-memory strike and block tracking for a single project instance."""
+
+    def __init__(self) -> None:
+        self._state: ProjectState = ProjectState()
+        self._lock = asyncio.Lock()
+
+    @staticmethod
+    def _normalize_user_id(user_id: Optional[str]) -> str:
+        """Normalize and validate user_id.
+
+        Args:
+            user_id: User identifier.
+
+        Returns:
+            Validated, non-empty user id.
+        """
+        uid = (user_id or "").strip()
+        if not uid:
+            raise ValueError("user_id is required for guardrails tracking.")
+        return uid
+
+    def _get_or_create_user(self, user_id: str) -> ScopeState:
+        """Return (and create if needed) the per-user scope state.
+
+        Args:
+            user_id: User identifier.
+
+        Returns:
+            Mutable ScopeState for the user.
+        """
+        uid = self._normalize_user_id(user_id)
+        return self._state.users.setdefault(uid, ScopeState())
+
+    def _clear_expired_project_block(self, now: float) -> None:
+        """Clear project-level block if expired.
+
+        Args:
+            now: Current time as a UNIX timestamp.
+        """
+        project = self._state.project
+        if project.blocked_until and now >= project.blocked_until:
+            project.blocked_until = None
+
+    def _clear_expired_user_block(self, user_id: str, now: float) -> None:
+        """Clear user-level block if expired.
+
+        Args:
+            user_id: User identifier.
+            now: Current time as a UNIX timestamp.
+        """
+        try:
+            uid = self._normalize_user_id(user_id)
+        except ValueError:
+            return
+
+        user_state = self._state.users.get(uid)
+        if user_state and user_state.blocked_until and now >= user_state.blocked_until:
+            user_state.blocked_until = None
+
+    @staticmethod
+    def _apply_user_block_if_needed(user_id: str, user_state: ScopeState) -> bool:
+        """Apply a user-level block if the threshold is reached.
+
+        Args:
+            user_id: User identifier.
+            user_state: Current user's scope state.
+
+        Returns:
+            True if a block was applied during this call; otherwise False.
+        """
+        now = time.time()
+        if user_state.is_blocked(now):
+            return False
+
+        if len(user_state.violations) < config.GUARDRAILS_USER_MAX_STRIKES:
+            return False
+
+        duration_seconds = config.GUARDRAILS_BLOCK_DURATION_SECONDS
+        block_until = float("inf") if duration_seconds <= 0 else now + duration_seconds
+        user_state.blocked_until = block_until
+        structlogger.info(
+            "guardrails.store.user_blocked",
+            project_id=config.HELLO_RASA_PROJECT_ID,
+            user_id=user_id,
+            strikes=len(user_state.violations),
+            duration_sec=duration_seconds,
+        )
+        return True
+
+    def _apply_project_block_if_needed(self) -> bool:
+        """Apply a project-level block if the threshold is reached.
+
+        Returns:
+            True if a block was applied during this call; otherwise False.
+        """
+        now = time.time()
+        project_state = self._state.project
+        if project_state.is_blocked(now):
+            return False
+
+        if len(project_state.violations) < config.GUARDRAILS_PROJECT_MAX_STRIKES:
+            return False
+
+        duration_seconds = config.GUARDRAILS_BLOCK_DURATION_SECONDS
+        block_until = float("inf") if duration_seconds <= 0 else now + duration_seconds
+        project_state.blocked_until = block_until
+        structlogger.info(
+            "guardrails.store.project_blocked",
+            project_id=config.HELLO_RASA_PROJECT_ID,
+            strikes=len(project_state.violations),
+            duration_sec=duration_seconds,
+        )
+        return True
+
+    async def check_block_scope(self, user_id: str) -> Optional[BlockScope]:
+        """Return current block scope if blocked.
+
+        Args:
+            user_id: User identifier.
+
+        Returns:
+            'user' if user blocked, 'project' if project blocked, otherwise None.
+        """
+        uid = self._normalize_user_id(user_id)
+
+        async with self._lock:
+            now = time.time()
+
+            # Clear expired blocks
+            self._clear_expired_project_block(now)
+            self._clear_expired_user_block(uid, now)
+
+            user_state = self._state.users.get(uid)
+
+            if user_state and user_state.is_blocked(now):
+                return BLOCK_SCOPE_USER
+            if self._state.project.is_blocked(now):
+                return BLOCK_SCOPE_PROJECT
+            return None
+
+    async def record_violation(self, user_id: str) -> BlockResult:
+        """Record a violation and apply thresholds/blocks.
+
+        Args:
+            user_id: User identifier.
+
+        Returns:
+            BlockResult indicating whether a new user or project block was applied.
+        """
+        uid = self._normalize_user_id(user_id)
+        result = BlockResult()
+
+        async with self._lock:
+            now = time.time()
+            user_state = self._get_or_create_user(uid)
+
+            # Accumulate violations (no sliding window)
+            self._state.project.violations.append(now)
+            user_state.violations.append(now)
+
+            # Apply user and project blocks
+            result.user_blocked_now = self._apply_user_block_if_needed(
+                user_id=uid,
+                user_state=user_state,
+            )
+            result.project_blocked_now = self._apply_project_block_if_needed()
+
+        return result
+
+    async def unblock_user(self, user_id: str) -> None:
+        """Unblock a user without altering strike history.
+
+        Args:
+            user_id: User identifier.
+        """
+        uid = self._normalize_user_id(user_id)
+
+        async with self._lock:
+            if state := self._state.users.get(uid):
+                state.blocked_until = None
+
+    async def reset_user(self, user_id: str) -> None:
+        """Unblock a user and clear their strikes.
+
+        Args:
+            user_id: User identifier.
+        """
+        uid = self._normalize_user_id(user_id)
+
+        async with self._lock:
+            if state := self._state.users.get(uid):
+                state.blocked_until = None
+                state.violations.clear()
+
+    async def unblock_project(self) -> None:
+        """Unblock the project without altering strike history."""
+        async with self._lock:
+            self._state.project.blocked_until = None
+
+    async def reset_project(self) -> None:
+        """Unblock the project and clear project-wide strikes."""
+        async with self._lock:
+            self._state.project.blocked_until = None
+            self._state.project.violations.clear()
+
+    async def reset_all(self) -> None:
+        """Reset all guardrail state (project and users)."""
+        async with self._lock:
+            # Clear project scope
+            self._state.project.blocked_until = None
+            self._state.project.violations.clear()
+
+            # Clear all users
+            for state in self._state.users.values():
+                state.blocked_until = None
+                state.violations.clear()
+
+
+# Singleton instance for application-wide use
+guardrails_store = GuardrailsInMemoryStore()