rasa-pro 3.13.11__py3-none-any.whl → 3.14.0__py3-none-any.whl

rasa/builder/guardrails/policy_checker.py

@@ -0,0 +1,324 @@
+import asyncio
+import copy
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, cast
+
+import structlog
+
+from rasa.builder.copilot.constants import ROLE_COPILOT, ROLE_USER
+from rasa.builder.copilot.copilot_response_handler import CopilotResponseHandler
+from rasa.builder.copilot.models import (
+    CopilotChatMessage,
+    CopilotContext,
+    GeneratedContent,
+    ResponseCategory,
+)
+from rasa.builder.guardrails.clients import GuardrailsClient
+from rasa.builder.guardrails.models import (
+    GuardrailRequestKey,
+    GuardrailResponse,
+)
+from rasa.builder.guardrails.utils import (
+    create_guardrail_request,
+)
+from rasa.builder.shared.tracker_context import (
+    AssistantConversationTurn,
+    TrackerContext,
+)
+
+if TYPE_CHECKING:
+    pass
+
+
+structlogger = structlog.get_logger()
+
+
+class GuardrailsPolicyChecker:
+    def __init__(self, guardrail_client: GuardrailsClient):
+        self.guardrail_client = guardrail_client
+
+    async def check_assistant_chat_for_policy_violations(
+        self,
+        tracker_context: TrackerContext,
+        hello_rasa_user_id: Optional[str],
+        hello_rasa_project_id: Optional[str],
+        **kwargs: Any,
+    ) -> TrackerContext:
+        """Return a sanitised TrackerContext with unsafe turns removed.
+
+        Only user messages are moderated - assistant messages are assumed safe.
+        LRU cache is used, so each unique user text is checked once.
+
+        Args:
+            tracker_context: The TrackerContext containing conversation turns.
+            hello_rasa_user_id: The user ID for the conversation.
+            hello_rasa_project_id: The project ID for the conversation.
+            **kwargs: Additional parameters for the guardrail request.
+
+        Returns:
+            TrackerContext with unsafe turns removed.
+        """
+        # Collect (turn_index, user_text) for all turns with a user message
+        items = self._format_user_messages_to_assistant(
+            tracker_context.conversation_turns
+        )
+
+        flagged_turns = await self._check_user_messages_for_violations(
+            items,
+            hello_rasa_user_id=hello_rasa_user_id,
+            hello_rasa_project_id=hello_rasa_project_id,
+            log_prefix="assistant_guardrails",
+            **kwargs,
+        )
+
+        if not flagged_turns:
+            return tracker_context
+
+        structlogger.info(
+            "guardrails_policy_checker.assistant_guardrails.turns_flagged",
+            count=len(flagged_turns),
+            turn_indices=sorted(flagged_turns),
+        )
+
+        # Build a TrackerContext with safe turns
+        safe_turns: List[AssistantConversationTurn] = [
+            turn
+            for idx, turn in enumerate(tracker_context.conversation_turns)
+            if idx not in flagged_turns
+        ]
+
+        new_tracker_context = copy.deepcopy(tracker_context)
+        new_tracker_context.conversation_turns = safe_turns
+        return new_tracker_context
+
+    async def check_copilot_chat_for_policy_violations(
+        self,
+        context: CopilotContext,
+        hello_rasa_user_id: Optional[str],
+        hello_rasa_project_id: Optional[str],
+        **kwargs: Any,
+    ) -> Optional[GeneratedContent]:
+        """Check the copilot chat history for guardrail policy violations.
+
+        Only user messages are moderated – assistant messages are assumed safe.
+        LRU cache is used, so each unique user text is checked once.
+
+        Args:
+            context: The CopilotContext containing the copilot chat history.
+            hello_rasa_user_id: The user ID for the conversation.
+            hello_rasa_project_id: The project ID for the conversation.
+            **kwargs: Additional parameters for the guardrail request.
+
+        Returns:
+            Returns a default violation response if the system flags any user message,
+            otherwise return None.
+        """
+        # Collect (index, text) for user messages; skip ones already marked as
+        # violations
+        items = self._format_user_messages_to_copilot(context.copilot_chat_history)
+
+        flagged_user_indices = await self._check_user_messages_for_violations(
+            items,
+            hello_rasa_user_id=hello_rasa_user_id,
+            hello_rasa_project_id=hello_rasa_project_id,
+            log_prefix="copilot_guardrails",
+            **kwargs,
+        )
+
+        self._annotate_flagged_user_messages(
+            context.copilot_chat_history, flagged_user_indices
+        )
+
+        if not flagged_user_indices:
+            return None
+
+        # Identify the latest user message index in the current request
+        last_user_idx: Optional[int] = None
+        for i in range(len(context.copilot_chat_history) - 1, -1, -1):
+            if getattr(context.copilot_chat_history[i], "role", None) == ROLE_USER:
+                last_user_idx = i
+                break
+
+        # Remove flagged user messages and their next copilot messages
+        indices_to_remove: Set[int] = set()
+        total = len(context.copilot_chat_history)
+        for uidx in flagged_user_indices:
+            indices_to_remove.add(uidx)
+            next_idx = uidx + 1
+            if (
+                next_idx < total
+                and getattr(context.copilot_chat_history[next_idx], "role", None)
+                == ROLE_COPILOT
+            ):
+                indices_to_remove.add(next_idx)
+
+        # Apply sanitization
+        filtered_copilot_chat_history = [
+            msg
+            for i, msg in enumerate(context.copilot_chat_history)
+            if i not in indices_to_remove
+        ]
+        if len(filtered_copilot_chat_history) != len(context.copilot_chat_history):
+            structlogger.info(
+                "guardrails_policy_checker"
+                ".copilot_guardrails"
+                ".copilot_chat_history_sanitized",
+                removed_indices=sorted(indices_to_remove),
+                removed_messages=(
+                    len(context.copilot_chat_history)
+                    - len(filtered_copilot_chat_history)
+                ),
+                kept_messages=len(filtered_copilot_chat_history),
+            )
+            context.copilot_chat_history = filtered_copilot_chat_history
+
+        # Block only if the latest user message in this request was flagged
+        if last_user_idx is not None and last_user_idx in flagged_user_indices:
+            return CopilotResponseHandler.respond_to_guardrail_policy_violations()
+
+        # Otherwise proceed (following messages are respected)
+        return None
+
+    async def _check_user_messages_for_violations(
+        self,
+        items: List[Tuple[int, str]],
+        hello_rasa_user_id: Optional[str],
+        hello_rasa_project_id: Optional[str],
+        log_prefix: str,
+        **kwargs: Any,
+    ) -> Set[int]:
+        """Run guardrail checks for provided (index, user_text) pairs.
+
+        Args:
+            items: List of tuples containing (index, user_text) to check.
+            hello_rasa_user_id: The user ID for the conversation.
+            hello_rasa_project_id: The project ID for the conversation.
+            log_prefix: Prefix for logging messages.
+            **kwargs: Additional parameters for the guardrail request.
+
+        Returns:
+            A set of indices that were flagged by the guardrails.
+        """
+        if not items:
+            return set()
+
+        # 1) Group indices by logical request key (hashable by value)
+        indices_by_key: Dict[GuardrailRequestKey, List[int]] = {}
+        for idx, text in items:
+            key = GuardrailRequestKey(
+                user_text=text,
+                hello_rasa_user_id=hello_rasa_user_id or "",
+                hello_rasa_project_id=hello_rasa_project_id or "",
+                # The client-specific request parameters go in the metadata
+                # so that they can be used to create the request
+                metadata=kwargs,
+            )
+            if not key.user_text:
+                continue
+            indices_by_key.setdefault(key, []).append(idx)
+
+        if not indices_by_key:
+            return set()
+
+        # 2) Create one task per logical key
+        tasks_by_key: Dict[GuardrailRequestKey, asyncio.Task[GuardrailResponse]] = {}
+        for key in indices_by_key:
+            request = create_guardrail_request(
+                client_type=type(self.guardrail_client),
+                user_text=key.user_text,
+                hello_rasa_user_id=key.hello_rasa_user_id,
+                hello_rasa_project_id=key.hello_rasa_project_id,
+                **kwargs,
+            )
+            tasks_by_key[key] = self.guardrail_client.schedule_check(request)
+
+        # 3) Await unique tasks once
+        keys = list(tasks_by_key.keys())
+        tasks = [tasks_by_key[k] for k in keys]
+        responses = await asyncio.gather(*tasks, return_exceptions=True)
+
+        # 4) Map results back to all corresponding indices
+        flagged: Set[int] = set()
+        for key, response in zip(keys, responses):
+            if isinstance(response, Exception):
+                structlogger.warning(
+                    f"{log_prefix}.request_failed", error=str(response)
+                )
+                continue
+            # At this point, response is guaranteed to be GuardrailResponse
+            # Use typing.cast to explicitly cast the type
+            guardrail_response = cast(GuardrailResponse, response)
+            if guardrail_response.flagged:
+                flagged.update(indices_by_key.get(key, []))
+
+        return flagged
+
+    def _annotate_flagged_user_messages(
+        self,
+        copilot_chat_history: List[CopilotChatMessage],
+        flagged_user_indices: Set[int],
+    ) -> None:
+        """Mark flagged user messages in-place on the original copilot chat history.
+
+        Args:
+            copilot_chat_history: The copilot chat history containing messages.
+            flagged_user_indices: Set of indices of user messages that were flagged.
+        """
+        if not flagged_user_indices:
+            return
+
+        total = len(copilot_chat_history)
+        for uidx in flagged_user_indices:
+            if 0 <= uidx < total and copilot_chat_history[uidx].role == ROLE_USER:
+                copilot_chat_history[
+                    uidx
+                ].response_category = ResponseCategory.GUARDRAILS_POLICY_VIOLATION
+
+    def _format_user_messages_to_assistant(
+        self, conversation_turns: List[AssistantConversationTurn]
+    ) -> List[Tuple[int, str]]:
+        """Collect (turn_index, user_text) tuples for all turns with user messages.
+
+        Args:
+            conversation_turns: The list of conversation turns.
+
+        Returns:
+            List of tuples containing (turn_index, user_text) for valid user messages.
+        """
+        items: List[Tuple[int, str]] = []
+        for idx, turn in enumerate(conversation_turns):
+            user_message = turn.user_message
+            if not user_message:
+                continue
+
+            text = (user_message.text or "").strip()
+            if not text:
+                continue
+
+            items.append((idx, text))
+
+        return items
+
+    def _format_user_messages_to_copilot(
+        self, copilot_chat_history: List[CopilotChatMessage]
+    ) -> List[Tuple[int, str]]:
+        """Collect (index, user_text) tuples for all messages with user messages.
+
+        Args:
+            copilot_chat_history: The list of messages.
+        """
+        items: List[Tuple[int, str]] = []
+        for idx, message in enumerate(copilot_chat_history):
+            if (
+                message.response_category
+                == ResponseCategory.GUARDRAILS_POLICY_VIOLATION
+            ):
+                continue
+            if message.role != ROLE_USER:
+                continue
+            formatted_message = message.build_openai_message()
+            text = (formatted_message.get("content") or "").strip()
+            if not text:
+                continue
+            items.append((idx, text))
+
+        return items

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()

rasa/builder/guardrails/utils.py

@@ -0,0 +1,94 @@
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Type
+
+import structlog
+
+from rasa.builder.guardrails.clients import GuardrailsClient, LakeraAIGuardrails
+from rasa.builder.guardrails.models import (
+    GuardrailRequest,
+    LakeraGuardrailRequest,
+)
+from rasa.shared.constants import ROLE_USER
+
+if TYPE_CHECKING:
+    from rasa.builder.guardrails.models import GuardrailType
+
+
+structlogger = structlog.get_logger()
+
+
+def map_lakera_detector_type_to_guardrail_type(
+    lakera_detector_type: str,
+) -> Optional["GuardrailType"]:
+    """Map a Lakera detector type to a Rasa guardrail type."""
+    from rasa.builder.guardrails.models import GuardrailType
+
+    # Check for exact matches first
+    LAKERA_DETECTOR_TYPES_2_RASA_GUARDRAIL_TYPES_MAPPING = {
+        "prompt_attack": GuardrailType.PROMPT_ATTACK,
+        "unknown_links": GuardrailType.MALICIOUS_LINKS,
+        "custom": GuardrailType.CUSTOM,
+    }
+
+    # Check for exact match first
+    if lakera_detector_type in LAKERA_DETECTOR_TYPES_2_RASA_GUARDRAIL_TYPES_MAPPING:
+        return LAKERA_DETECTOR_TYPES_2_RASA_GUARDRAIL_TYPES_MAPPING[
+            lakera_detector_type
+        ]
+
+    # Check for subtypes that start with specific prefixes
+    # https://docs.lakera.ai/docs/policies/self-hosted-policies#detectors-section
+    if lakera_detector_type.startswith("moderated_content"):
+        return GuardrailType.CONTENT_VIOLATION
+    if lakera_detector_type.startswith("pii"):
+        return GuardrailType.DATA_LEAKAGE
+
+    # If no match found, return OTHER
+    return GuardrailType.OTHER
+
+
+def create_guardrail_request(
+    client_type: Type[GuardrailsClient],
+    user_text: str,
+    hello_rasa_user_id: str,
+    hello_rasa_project_id: str,
+    **kwargs: Any,
+) -> GuardrailRequest:
+    """Create a guardrail request."""
+
+    def _create_lakera_guardrail_request(
+        user_text: str,
+        hello_rasa_user_id: str,
+        hello_rasa_project_id: str,
+        **kwargs: Any,
+    ) -> LakeraGuardrailRequest:
+        """Create a Lakera guardrail request."""
+        return LakeraGuardrailRequest(
+            hello_rasa_user_id=hello_rasa_user_id,
+            hello_rasa_project_id=hello_rasa_project_id,
+            messages=[{"role": ROLE_USER, "content": user_text}],
+            **kwargs,
+        )
+
+    map_client_to_request: Dict[
+        Type[GuardrailsClient], Callable[..., GuardrailRequest]
+    ] = {
+        LakeraAIGuardrails: _create_lakera_guardrail_request,
+    }
+
+    if client_type in map_client_to_request:
+        return map_client_to_request[client_type](
+            user_text,
+            hello_rasa_user_id,
+            hello_rasa_project_id,
+            **kwargs,
+        )
+    else:
+        message = f"Unsupported guardrail client: {type(client_type)}"
+        structlogger.error(
+            "guardrails_policy_checker"
+            ".create_guardrail_request"
+            ".unsupported_guardrail_client",
+            message=message,
+            guardrail_client=client_type,
+        )
+        raise ValueError(message)