rasa-pro 3.13.1a18__py3-none-any.whl → 3.13.1a20__py3-none-any.whl

rasa/builder/guardrails/models.py

@@ -0,0 +1,199 @@
+"""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)

rasa/builder/guardrails/utils.py

@@ -0,0 +1,305 @@
+import asyncio
+from functools import lru_cache
+from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple
+
+import structlog
+
+from rasa.builder.config import (
+    ASSISTANT_HISTORY_GUARDRAIL_PROJECT_ID,
+    COPILOT_HISTORY_GUARDRAIL_PROJECT_ID,
+)
+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 (
+    CopilotContext,
+    GeneratedContent,
+    ResponseCategory,
+)
+from rasa.builder.guardrails.models import (
+    GuardrailRequestKey,
+    GuardrailResponse,
+    LakeraGuardrailRequest,
+)
+from rasa.builder.llm_service import llm_service
+from rasa.builder.shared.tracker_context import (
+    AssistantConversationTurn,
+    TrackerContext,
+)
+
+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
+
+
+@lru_cache(maxsize=512)
+def _schedule_guardrails_check(
+    user_text: str,
+    hello_rasa_user_id: str,
+    hello_rasa_project_id: str,
+    lakera_project_id: str,
+) -> "asyncio.Task[GuardrailResponse]":
+    """Return a cached asyncio.Task that resolves to Lakera's response.
+
+    Args:
+        user_text: The user message text to check for policy violations.
+        hello_rasa_user_id: The user ID for the conversation.
+        hello_rasa_project_id: The project ID for the conversation.
+        lakera_project_id: The Lakera project ID to use for this check.
+
+    Returns:
+        An asyncio Task that resolves to a GuardrailResponse.
+    """
+    structlogger.debug("guardrails.cache_miss", text=user_text)
+
+    loop = asyncio.get_running_loop()
+    request = LakeraGuardrailRequest(
+        lakera_project_id=lakera_project_id,
+        hello_rasa_user_id=hello_rasa_user_id,
+        hello_rasa_project_id=hello_rasa_project_id,
+        messages=[{"role": ROLE_USER, "content": user_text}],
+    )
+
+    return loop.create_task(llm_service.guardrails.send_request(request))
+
+
+async def _detect_flagged_user_indices(
+    items: List[Tuple[int, str]],
+    *,
+    hello_rasa_user_id: Optional[str],
+    hello_rasa_project_id: Optional[str],
+    lakera_project_id: str,
+    log_prefix: str,
+) -> 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.
+        lakera_project_id: The Lakera project ID to use for this check.
+        log_prefix: Prefix for logging messages.
+
+    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 or "").strip(),
+            hello_rasa_user_id=hello_rasa_user_id or "",
+            hello_rasa_project_id=hello_rasa_project_id or "",
+            lakera_project_id=lakera_project_id,
+        )
+        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:
+        tasks_by_key[key] = _schedule_guardrails_check(
+            user_text=key.user_text,
+            hello_rasa_user_id=key.hello_rasa_user_id,
+            hello_rasa_project_id=key.hello_rasa_project_id,
+            lakera_project_id=key.lakera_project_id,
+        )
+
+    # 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
+        if response.flagged:
+            flagged.update(indices_by_key.get(key, []))
+
+    return flagged
+
+
+async def check_assistant_chat_for_policy_violations(
+    tracker_context: TrackerContext,
+    hello_rasa_user_id: Optional[str],
+    hello_rasa_project_id: Optional[str],
+) -> 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.
+
+    Returns:
+        TrackerContext with unsafe turns removed.
+    """
+    # Collect (turn_index, user_text) for all turns with a user message
+    items: List[Tuple[int, str]] = []
+    for idx, turn in enumerate(tracker_context.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))
+
+    flagged_turns = await _detect_flagged_user_indices(
+        items,
+        hello_rasa_user_id=hello_rasa_user_id,
+        hello_rasa_project_id=hello_rasa_project_id,
+        lakera_project_id=ASSISTANT_HISTORY_GUARDRAIL_PROJECT_ID,
+        log_prefix="assistant_guardrails",
+    )
+
+    if not flagged_turns:
+        return tracker_context
+
+    structlogger.info(
+        "assistant_guardrails.turns_flagged",
+        count=len(flagged_turns),
+        turn_indices=sorted(flagged_turns),
+    )
+
+    # Build a filtered TrackerContext
+    safe_turns: List[AssistantConversationTurn] = [
+        turn
+        for idx, turn in enumerate(tracker_context.conversation_turns)
+        if idx not in flagged_turns
+    ]
+
+    new_tracker_context = tracker_context.copy(deep=True)
+    new_tracker_context.conversation_turns = safe_turns
+    return new_tracker_context
+
+
+async def check_copilot_chat_for_policy_violations(
+    context: CopilotContext,
+    hello_rasa_user_id: Optional[str],
+    hello_rasa_project_id: Optional[str],
+) -> 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.
+
+    Returns:
+        Returns a default violation response if the system flags any user message,
+        otherwise return None.
+    """
+    history = context.copilot_chat_history
+
+    # Collect (index, text) for user messages; skip ones already marked as violations
+    items: List[Tuple[int, str]] = []
+    for idx, message in enumerate(history):
+        if message.response_category == ResponseCategory.GUARDRAILS_POLICY_VIOLATION:
+            continue
+        if message.role != ROLE_USER:
+            continue
+        formatted_message = message.to_openai_format()
+        text = (formatted_message.get("content") or "").strip()
+        if not text:
+            continue
+        items.append((idx, text))
+
+    flagged_user_indices = await _detect_flagged_user_indices(
+        items,
+        hello_rasa_user_id=hello_rasa_user_id,
+        hello_rasa_project_id=hello_rasa_project_id,
+        lakera_project_id=COPILOT_HISTORY_GUARDRAIL_PROJECT_ID,
+        log_prefix="copilot_guardrails",
+    )
+
+    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(history) - 1, -1, -1):
+        if getattr(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(history)
+    for uidx in flagged_user_indices:
+        indices_to_remove.add(uidx)
+        next_idx = uidx + 1
+        if (
+            next_idx < total
+            and getattr(history[next_idx], "role", None) == ROLE_COPILOT
+        ):
+            indices_to_remove.add(next_idx)
+
+    # Apply sanitization
+    filtered_history = [
+        msg for i, msg in enumerate(history) if i not in indices_to_remove
+    ]
+    if len(filtered_history) != len(history):
+        structlogger.info(
+            "copilot_guardrails.history_sanitized",
+            removed_indices=sorted(indices_to_remove),
+            removed_messages=len(history) - len(filtered_history),
+            kept_messages=len(filtered_history),
+        )
+        context.copilot_chat_history = filtered_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

rasa/builder/job_manager.py

@@ -0,0 +1,87 @@
+import asyncio
+import time
+import uuid
+from typing import AsyncGenerator, ClassVar, Dict, List, Optional
+
+from pydantic import BaseModel, Field, PrivateAttr
+
+from rasa.builder.models import JobStatusEvent, ServerSentEventType
+
+
+class JobInfo(BaseModel):
+    """Information about a job being processed by the worker."""
+
+    id: str
+    status: str = ""
+    error: Optional[str] = None
+    created_at: float = Field(default_factory=time.time)
+
+    _history: List[JobStatusEvent] = PrivateAttr(default_factory=list)
+    _queue: asyncio.Queue = PrivateAttr(default_factory=asyncio.Queue)
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    async def put(self, event: JobStatusEvent) -> None:
+        """Put an event onto the job's queue.
+
+        Args:
+            event: The JobStatusEvent to put onto the queue.
+        """
+        await self._queue.put(event)
+
+    async def event_stream(self) -> AsyncGenerator[JobStatusEvent, None]:
+        """Yield events as they are put on the queue by the worker task.
+
+        Returns:
+            An async generator that yields JobStatusEvent objects.
+        """
+        # 1) Replay history and stop if EOF recorded
+        for sse in self._history:
+            if (
+                isinstance(sse, JobStatusEvent)
+                and sse.event == ServerSentEventType._EOF.value
+            ):
+                return
+            yield sse
+
+        # 2) Stream live events and stop on EOF
+        while True:
+            sse = await self._queue.get()
+            self._history.append(sse)
+            if (
+                isinstance(sse, JobStatusEvent)
+                and sse.event == ServerSentEventType._EOF.value
+            ):
+                return
+            yield sse
+
+
+class JobManager:
+    """Very small in-memory job registry (single-process only)."""
+
+    _jobs: ClassVar[Dict[str, JobInfo]] = {}
+
+    def create_job(self) -> JobInfo:
+        job_id = uuid.uuid4().hex
+        job = JobInfo(id=job_id)
+        self._jobs[job_id] = job
+        return job
+
+    def get_job(self, job_id: str) -> Optional[JobInfo]:
+        return self._jobs.get(job_id)
+
+    @staticmethod
+    def mark_done(job: JobInfo, *, error: Optional[str] = None) -> None:
+        """Mark a job as done.
+
+        Args:
+            job: The JobInfo instance to mark as done.
+            error: Optional error message if the job failed.
+        """
+        job.error = error
+        eof = JobStatusEvent.eof()
+        job._queue.put_nowait(eof)
+
+
+job_manager = JobManager()