rasa-pro 3.14.0.dev20250922__py3-none-any.whl → 3.14.0rc2__py3-none-any.whl

rasa/builder/guardrails/utils.py

@@ -1,31 +1,13 @@
-import asyncio
-from functools import lru_cache
-from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Type
 
 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 (
-    CopilotChatMessage,
-    CopilotContext,
-    GeneratedContent,
-    ResponseCategory,
-)
+from rasa.builder.guardrails.clients import GuardrailsClient, LakeraAIGuardrails
 from rasa.builder.guardrails.models import (
-    GuardrailRequestKey,
-    GuardrailResponse,
+    GuardrailRequest,
     LakeraGuardrailRequest,
 )
-from rasa.builder.llm_service import llm_service
-from rasa.builder.shared.tracker_context import (
-    AssistantConversationTurn,
-    TrackerContext,
-)
+from rasa.shared.constants import ROLE_USER
 
 if TYPE_CHECKING:
     from rasa.builder.guardrails.models import GuardrailType
@@ -64,265 +46,49 @@ def map_lakera_detector_type_to_guardrail_type(
     return GuardrailType.OTHER
 
 
-@lru_cache(maxsize=512)
-def _schedule_guardrails_check(
+def create_guardrail_request(
+    client_type: Type[GuardrailsClient],
     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,
+    **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,
         )
-        if not key.user_text:
-            continue
-        indices_by_key.setdefault(key, []).append(idx)
 
-    if not indices_by_key:
-        return set()
+    map_client_to_request: Dict[
+        Type[GuardrailsClient], Callable[..., GuardrailRequest]
+    ] = {
+        LakeraAIGuardrails: _create_lakera_guardrail_request,
+    }
 
-    # 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,
+    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,
         )
-
-    # 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
-
-
-def _annotate_flagged_user_messages(
-    history: List[CopilotChatMessage], flagged_user_indices: Set[int]
-) -> None:
-    """Mark flagged user messages in-place on the original history.
-
-    Args:
-        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(history)
-    for uidx in flagged_user_indices:
-        if 0 <= uidx < total and history[uidx].role == ROLE_USER:
-            history[
-                uidx
-            ].response_category = ResponseCategory.GUARDRAILS_POLICY_VIOLATION
-
-
-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",
-    )
-
-    _annotate_flagged_user_messages(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(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),
+    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,
         )
-        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
+        raise ValueError(message)

rasa/builder/jobs.py

@@ -1,9 +1,17 @@
-from typing import Any, Optional
+from typing import Any, Dict, Optional
 
 import structlog
 from sanic import Sanic
 
 from rasa.builder import config
+from rasa.builder.copilot.models import (
+    CopilotContext,
+    FileContent,
+    InternalCopilotRequestChatMessage,
+    LogContent,
+    ResponseCategory,
+    TrainingErrorLog,
+)
 from rasa.builder.exceptions import (
     LLMGenerationError,
     ProjectGenerationError,
@@ -11,6 +19,7 @@ from rasa.builder.exceptions import (
     ValidationError,
 )
 from rasa.builder.job_manager import JobInfo, job_manager
+from rasa.builder.llm_service import llm_service
 from rasa.builder.models import (
     JobStatus,
     JobStatusEvent,
@@ -28,9 +37,14 @@ structlogger = structlog.get_logger()
 
 
 async def push_job_status_event(
-    job: JobInfo, status: JobStatus, message: Optional[str] = None
+    job: JobInfo,
+    status: JobStatus,
+    message: Optional[str] = None,
+    payload: Optional[Dict[str, Any]] = None,
 ) -> None:
-    event = JobStatusEvent.from_status(status=status.value, message=message)
+    event = JobStatusEvent.from_status(
+        status=status.value, message=message, payload=payload
+    )
     job.status = status.value
     await job.put(event)
 
@@ -210,7 +224,7 @@ async def run_template_to_bot_job(
 async def run_replace_all_files_job(
     app: "Sanic",
     job: JobInfo,
-    bot_files: dict,
+    bot_files: Dict[str, Any],
 ) -> None:
     """Run the replace-all-files job in the background.
 
@@ -228,7 +242,7 @@ async def run_replace_all_files_job(
     try:
         project_generator.replace_all_bot_files(bot_files)
 
-        # 1. Validating
+        # Validating
         await push_job_status_event(job, JobStatus.validating)
         training_input = project_generator.get_training_input()
         validation_error = await validate_project(training_input.importer)
@@ -236,12 +250,13 @@ async def run_replace_all_files_job(
             raise ValidationError(validation_error)
         await push_job_status_event(job, JobStatus.validation_success)
 
-        # 2. Training
+        # Training
         await push_job_status_event(job, JobStatus.training)
         agent = await train_and_load_agent(training_input)
         update_agent(agent, app)
         await push_job_status_event(job, JobStatus.train_success)
 
+        # Send final done event
         await push_job_status_event(job, JobStatus.done)
         job_manager.mark_done(job)
 
@@ -257,26 +272,181 @@ async def run_replace_all_files_job(
             included_log_levels=log_levels,
         )
         error_message = exc.get_error_message_with_logs(log_levels=log_levels)
-        await push_job_status_event(
-            job, JobStatus.validation_error, message=error_message
+        # Push error event and start copilot analysis job
+        await push_error_and_start_copilot_analysis(
+            app,
+            job,
+            JobStatus.validation_error,
+            error_message,
+            bot_files,
         )
+
+        # After error mark job as done
         job_manager.mark_done(job, error=error_message)
 
     except TrainingError as exc:
+        error_message = str(exc)
         structlogger.debug(
             "replace_all_files_job.train_error",
             job_id=job.id,
-            error=str(exc),
+            error=error_message,
         )
-        await push_job_status_event(job, JobStatus.train_error, message=str(exc))
-        job_manager.mark_done(job, error=str(exc))
+        # Push error event and start copilot analysis job
+        await push_error_and_start_copilot_analysis(
+            app,
+            job,
+            JobStatus.train_error,
+            error_message,
+            bot_files,
+        )
+
+        # After error mark job as done
+        job_manager.mark_done(job, error=error_message)
 
     except Exception as exc:
         # Capture full traceback for anything truly unexpected
+        error_message = str(exc)
         structlogger.exception(
             "replace_all_files_job.unexpected_error",
             job_id=job.id,
+            error=error_message,
+        )
+
+        # Push error event and start copilot analysis job
+        await push_error_and_start_copilot_analysis(
+            app,
+            job,
+            JobStatus.error,
+            error_message,
+            bot_files,
+        )
+
+        # After error mark job as done
+        job_manager.mark_done(job, error=str(exc))
+
+
+async def push_error_and_start_copilot_analysis(
+    app: "Sanic",
+    original_job: JobInfo,
+    original_job_status: JobStatus,
+    error_message: str,
+    bot_files: Dict[str, Any],
+) -> None:
+    """Start a copilot analysis job and notify the client.
+
+    Creates a copilot analysis job and sends the new job ID to the client. The new
+    job runs in the background.
+
+    Args:
+        app: The Sanic application instance
+        original_job: The original job that failed
+        original_job_status: The status of the job that failed
+        error_message: The error message to analyze
+        bot_files: The bot files to include in analysis
+    """
+    # Create a copilot analysis job. Send the new job ID to the client and
+    # run the Copilot Analysis job in the background.
+    message = "Failed to train the assistant. Starting copilot analysis."
+
+    copilot_job = job_manager.create_job()
+    # Push the error status event for the original job
+    await push_job_status_event(
+        original_job,
+        original_job_status,
+        message=message,
+        payload={"copilot_job_id": copilot_job.id},
+    )
+    # Run the copilot analysis job in the background
+    app.add_task(
+        run_copilot_training_error_analysis_job(
+            app, copilot_job, error_message, bot_files
+        )
+    )
+    structlogger.debug(
+        f"update_files_job.{original_job_status.value}.copilot_analysis_start",
+        event_info=message,
+        job_id=original_job.id,
+        error=error_message,
+        copilot_job_id=copilot_job.id,
+    )
+
+
+async def run_copilot_training_error_analysis_job(
+    app: "Sanic",
+    job: JobInfo,
+    training_error_message: str,
+    bot_files: Dict[str, Any],
+) -> None:
+    """Run copilot training error analysis job."""
+    await push_job_status_event(job, JobStatus.received)
+
+    try:
+        # Create message content blocks with log content and available files
+        log_content_block = LogContent(
+            type="log", content=training_error_message, context="training_error"
+        )
+        file_content_blocks = [
+            FileContent(type="file", file_path=file_path, file_content=file_content)
+            for file_path, file_content in bot_files.items()
+        ]
+        context = CopilotContext(
+            tracker_context=None,  # No conversation context needed
+            copilot_chat_history=[
+                InternalCopilotRequestChatMessage(
+                    role="internal_copilot_request",
+                    content=[log_content_block, *file_content_blocks],
+                    response_category=ResponseCategory.TRAINING_ERROR_LOG_ANALYSIS,
+                )
+            ],
+        )
+
+        # Generate copilot response
+        copilot_client = llm_service.instantiate_copilot()
+        (
+            original_stream,
+            generation_context,
+        ) = await copilot_client.generate_response(context)
+
+        copilot_response_handler = llm_service.instantiate_handler(
+            config.COPILOT_HANDLER_ROLLING_BUFFER_SIZE
+        )
+        intercepted_stream = copilot_response_handler.handle_response(original_stream)
+
+        # Stream the copilot response as job events
+        async for token in intercepted_stream:
+            # Send each token as a job event using the same format as /copilot endpoint
+            await push_job_status_event(
+                job, JobStatus.copilot_analyzing, payload=token.sse_data
+            )
+
+        # Send references (if any) as part of copilot_analyzing stream
+        if generation_context.relevant_documents:
+            reference_section = copilot_response_handler.extract_references(
+                generation_context.relevant_documents
+            )
+            await push_job_status_event(
+                job, JobStatus.copilot_analyzing, payload=reference_section.sse_data
+            )
+
+        # Send original error log as part of copilot_analyzing stream
+        training_error_log = TrainingErrorLog(logs=[log_content_block])
+        await push_job_status_event(
+            job, JobStatus.copilot_analyzing, payload=training_error_log.sse_data
+        )
+
+        # Send success status
+        await push_job_status_event(job, JobStatus.copilot_analysis_success)
+
+        await push_job_status_event(job, JobStatus.done)
+        job_manager.mark_done(job)
+
+    except Exception as exc:
+        structlogger.exception(
+            "copilot_training_error_analysis_job.error",
+            job_id=job.id,
             error=str(exc),
         )
-        await push_job_status_event(job, JobStatus.error, message=str(exc))
+        await push_job_status_event(
+            job, JobStatus.copilot_analysis_error, message=str(exc)
+        )
         job_manager.mark_done(job, error=str(exc))

rasa/builder/llm_service.py

@@ -19,7 +19,11 @@ from rasa.builder.copilot.copilot_templated_message_provider import (
     load_copilot_internal_message_templates,
 )
 from rasa.builder.exceptions import LLMGenerationError
-from rasa.builder.guardrails.lakera import LakeraAIGuardrails
+from rasa.builder.guardrails.clients import (
+    GuardrailsClient,
+    LakeraAIGuardrails,
+)
+from rasa.builder.guardrails.policy_checker import GuardrailsPolicyChecker
 from rasa.constants import PACKAGE_NAME
 from rasa.shared.constants import DOMAIN_SCHEMA_FILE, RESPONSES_SCHEMA_FILE
 from rasa.shared.core.flows.yaml_flows_io import FLOWS_SCHEMA_FILE
@@ -37,7 +41,8 @@ class LLMService:
         self._domain_schema: Optional[Dict[str, Any]] = None
         self._flows_schema: Optional[Dict[str, Any]] = None
         self._copilot: Optional[Copilot] = None
-        self._guardrails: Optional[LakeraAIGuardrails] = None
+        self._guardrails: Optional[GuardrailsClient] = None
+        self._guardrails_policy_checker: Optional[GuardrailsPolicyChecker] = None
         self._copilot_response_handler: Optional[CopilotResponseHandler] = None
         self._copilot_internal_message_templates: Optional[Dict[str, str]] = None
 
@@ -77,11 +82,14 @@ class LLMService:
             raise
 
     @property
-    def guardrails(self) -> LakeraAIGuardrails:
+    def guardrails(self) -> Optional[GuardrailsClient]:
         """Get or lazy create guardrails instance."""
-        if self._guardrails is None:
-            self._guardrails = LakeraAIGuardrails()
+        if not config.ENABLE_GUARDRAILS:
+            return None
+        # TODO: Replace with Open Source guardrails implementation once it's ready
         try:
+            if self._guardrails is None:
+                self._guardrails = LakeraAIGuardrails()
             return self._guardrails
         except Exception as e:
             structlogger.error(
@@ -91,6 +99,25 @@ class LLMService:
             )
             raise
 
+    @property
+    def guardrails_policy_checker(self) -> Optional[GuardrailsPolicyChecker]:
+        """Get or lazy create guardrails policy checker instance."""
+        try:
+            if self._guardrails_policy_checker is None and self.guardrails is not None:
+                self._guardrails_policy_checker = GuardrailsPolicyChecker(
+                    self.guardrails
+                )
+            return self._guardrails_policy_checker
+        except Exception as e:
+            structlogger.error(
+                "llm_service.guardrails_policy_checker.error",
+                event_info=(
+                    "LLM Service: Error getting guardrails policy checker instance."
+                ),
+                error=str(e),
+            )
+            raise
+
     @property
     def copilot_internal_message_templates(self) -> Dict[str, str]:
         """Get or lazy load copilot internal message templates."""