rasa-pro 3.14.0.dev20250922__py3-none-any.whl → 3.14.0rc1__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/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."""

rasa/builder/models.py

@@ -206,6 +206,7 @@ class TrainingInput(BaseModel):
 
     importer: TrainingDataImporter = Field(..., description="Training data importer")
     endpoints_file: Path = Field(..., description="Path to the endpoints file")
+    config_file: Path = Field(..., description="Path to the config file")
 
 
 class AgentStatus(str, Enum):

rasa/builder/project_generator.py

@@ -192,11 +192,16 @@ class ProjectGenerator:
         """Get the endpoints file."""
         return self.project_folder / "endpoints.yml"
 
+    def _get_config_file(self) -> Path:
+        """Get the config file."""
+        return self.project_folder / "config.yml"
+
     def get_training_input(self) -> TrainingInput:
         """Get the training input."""
         return TrainingInput(
             importer=self._create_importer(),
             endpoints_file=self._get_endpoints_file(),
+            config_file=self._get_config_file(),
         )
 
     def _create_importer(self) -> TrainingDataImporter:
@@ -208,7 +213,7 @@ class ProjectGenerator:
                 domain_path = self.project_folder / "domain"
 
             return TrainingDataImporter.load_from_config(
-                config_path=str(self.project_folder / "config.yml"),
+                config_path=str(self._get_config_file()),
                 domain_path=str(domain_path),
                 training_data_paths=[str(self.project_folder / "data")],
                 args={},

rasa/builder/service.py

@@ -17,6 +17,8 @@ from rasa.builder.config import (
     COPILOT_HANDLER_ROLLING_BUFFER_SIZE,
     GUARDRAILS_ENABLE_BLOCKING,
     HELLO_RASA_PROJECT_ID,
+    LAKERA_ASSISTANT_HISTORY_GUARDRAIL_PROJECT_ID,
+    LAKERA_COPILOT_HISTORY_GUARDRAIL_PROJECT_ID,
 )
 from rasa.builder.copilot.constants import ROLE_USER, SIGNATURE_VERSION_V1
 from rasa.builder.copilot.copilot_response_handler import CopilotResponseHandler
@@ -47,10 +49,6 @@ from rasa.builder.guardrails.constants import (
     BlockScope,
 )
 from rasa.builder.guardrails.store import guardrails_store
-from rasa.builder.guardrails.utils import (
-    check_assistant_chat_for_policy_violations,
-    check_copilot_chat_for_policy_violations,
-)
 from rasa.builder.job_manager import job_manager
 from rasa.builder.jobs import (
     run_prompt_to_bot_job,
@@ -1081,11 +1079,15 @@ async def copilot(request: Request) -> None:
         tracker_context = TrackerContext.from_tracker(
             tracker, max_turns=COPILOT_ASSISTANT_TRACKER_MAX_TURNS
         )
-        if tracker_context is not None:
-            tracker_context = await check_assistant_chat_for_policy_violations(
+        if (
+            tracker_context is not None
+            and llm_service.guardrails_policy_checker is not None
+        ):
+            tracker_context = await llm_service.guardrails_policy_checker.check_assistant_chat_for_policy_violations(  # noqa: E501
                 tracker_context=tracker_context,
                 hello_rasa_user_id=user_id,
                 hello_rasa_project_id=HELLO_RASA_PROJECT_ID,
+                lakera_project_id=LAKERA_ASSISTANT_HISTORY_GUARDRAIL_PROJECT_ID,
             )
 
         # Copilot doesn't need to know about the docs and any file that is not a core
@@ -1103,13 +1105,14 @@ async def copilot(request: Request) -> None:
 
         # 5. Run guardrail policy checks. If any policy violations are detected,
         #    send a response and end the stream.
-        guardrail_response: Optional[
-            GeneratedContent
-        ] = await check_copilot_chat_for_policy_violations(
-            context=context,
-            hello_rasa_user_id=user_id,
-            hello_rasa_project_id=HELLO_RASA_PROJECT_ID,
-        )
+        guardrail_response: Optional[GeneratedContent] = None
+        if llm_service.guardrails_policy_checker is not None:
+            guardrail_response = await llm_service.guardrails_policy_checker.check_copilot_chat_for_policy_violations(  # noqa: E501
+                context=context,
+                hello_rasa_user_id=user_id,
+                hello_rasa_project_id=HELLO_RASA_PROJECT_ID,
+                lakera_project_id=LAKERA_COPILOT_HISTORY_GUARDRAIL_PROJECT_ID,
+            )
         if guardrail_response is not None:
             blocked_or_violation_message = (
                 await _handle_guardrail_violation_and_maybe_block(

rasa/builder/training_service.py

@@ -11,9 +11,10 @@ from rasa.builder.exceptions import AgentLoadError, TrainingError
 from rasa.builder.models import TrainingInput
 from rasa.core.agent import Agent, load_agent
 from rasa.core.channels.studio_chat import StudioChatInput
-from rasa.core.utils import AvailableEndpoints, read_endpoints_from_path
+from rasa.core.config.configuration import Configuration
 from rasa.model import get_latest_model
 from rasa.model_training import TrainingResult, train
+from rasa.shared.constants import DEFAULT_ENDPOINTS_PATH
 from rasa.shared.importers.importer import TrainingDataImporter
 
 structlogger = structlog.get_logger()
@@ -42,14 +43,13 @@ async def train_and_load_agent(input: TrainingInput) -> Agent:
         AgentLoadError: If agent loading fails
     """
     try:
-        # Setup endpoints for training validation
-        await _setup_endpoints(input.endpoints_file)
-
         # Train the model
-        training_result = await _train_model(input.importer)
+        training_result = await _train_model(
+            input.importer, input.endpoints_file, input.config_file
+        )
 
         # Load the agent
-        agent_instance = await _load_agent(training_result.model)
+        agent_instance = await _load_agent(training_result.model, input.endpoints_file)
 
         # Verify agent is ready
         if not agent_instance.is_ready():
@@ -96,7 +96,9 @@ async def try_load_existing_agent(project_folder: str) -> Optional[Agent]:
         )
 
         # Get available endpoints for agent loading
-        available_endpoints = AvailableEndpoints.get_instance()
+        available_endpoints = Configuration.initialise_endpoints(
+            endpoints_path=Path(project_folder) / DEFAULT_ENDPOINTS_PATH
+        ).endpoints
 
         # Load the agent
         agent = await load_agent(
@@ -124,27 +126,17 @@ async def try_load_existing_agent(project_folder: str) -> Optional[Agent]:
         return None
 
 
-async def _setup_endpoints(endpoints_file: Path) -> None:
-    """Setup endpoints configuration for training."""
-    try:
-        # Reset and load endpoints
-        AvailableEndpoints.reset_instance()
-        read_endpoints_from_path(endpoints_file)
-
-        structlogger.debug("training.endpoints_setup", endpoints_file=endpoints_file)
-
-    except Exception as e:
-        raise TrainingError(f"Failed to setup endpoints: {e}")
-
-
-async def _train_model(importer: TrainingDataImporter) -> TrainingResult:
+async def _train_model(
+    importer: TrainingDataImporter, endpoints_file: Path, config_file: Path
+) -> TrainingResult:
     """Train the Rasa model."""
     try:
         structlogger.info("training.started")
 
         training_result = await train(
             domain="",
-            config="",
+            config=str(config_file),
+            endpoints=str(endpoints_file),
             training_files=None,
             file_importer=importer,
         )
@@ -160,12 +152,14 @@ async def _train_model(importer: TrainingDataImporter) -> TrainingResult:
         raise TrainingError(f"Model training failed: {e}")
 
 
-async def _load_agent(model_path: str) -> Agent:
+async def _load_agent(model_path: str, endpoints_file: Path) -> Agent:
     """Load the trained agent."""
     try:
         structlogger.info("training.loading_agent", model_path=model_path)
 
-        available_endpoints = AvailableEndpoints.get_instance()
+        available_endpoints = Configuration.initialise_endpoints(
+            endpoints_path=endpoints_file
+        ).endpoints
         if available_endpoints is None:
             raise AgentLoadError("No endpoints available for agent loading")
 

rasa/builder/validation_service.py

@@ -9,7 +9,7 @@ import structlog
 from rasa.builder import config
 from rasa.builder.exceptions import ValidationError
 from rasa.builder.logging_utils import capture_validation_logs
-from rasa.cli.utils import validate_files
+from rasa.cli.validation.bot_config import validate_files
 from rasa.shared.importers.importer import TrainingDataImporter
 
 structlogger = structlog.get_logger()

rasa/cli/arguments/default_arguments.py

@@ -3,6 +3,7 @@ import logging
 from enum import Enum
 from typing import List, Optional, Text, Union
 
+from rasa.core.constants import DEFAULT_SUB_AGENTS
 from rasa.core.persistor import RemoteStorageType, StorageType, parse_remote_storage
 from rasa.shared.constants import (
     DEFAULT_CONFIG_PATH,
@@ -217,3 +218,14 @@ def add_skip_validation_flag(
         action="append",
         help="Skip YAML validation for selected parts of the training data.",
     )
+
+
+def add_sub_agents_param(
+    parser: Union[argparse.ArgumentParser, argparse._ActionsContainer],
+) -> None:
+    parser.add_argument(
+        "--sub-agents",
+        type=str,
+        default=DEFAULT_SUB_AGENTS,
+        help="Directory that specifies sub-agents to use (default: %(default)s).",
+    )

rasa/cli/arguments/run.py

@@ -7,6 +7,7 @@ from rasa.cli.arguments.default_arguments import (
     add_model_param,
     add_remote_storage_param,
     add_skip_validation_flag,
+    add_sub_agents_param,
 )
 from rasa.core import constants
 from rasa.env import (
@@ -24,6 +25,7 @@ def set_run_arguments(parser: argparse.ArgumentParser) -> None:
     add_server_arguments(parser)
     add_inspect_argument(parser)
     add_skip_validation_flag(parser)
+    add_sub_agents_param(parser)
 
 
 def set_run_action_arguments(parser: argparse.ArgumentParser) -> None:

rasa/cli/arguments/train.py

@@ -10,6 +10,7 @@ from rasa.cli.arguments.default_arguments import (
     add_remote_root_only_param,
     add_remote_storage_param,
     add_stories_param,
+    add_sub_agents_param,
 )
 from rasa.graph_components.providers.training_tracker_provider import (
     TrainingTrackerProvider,
@@ -43,6 +44,7 @@ def set_train_arguments(parser: argparse.ArgumentParser) -> None:
     )
     add_remote_storage_param(parser)
     add_remote_root_only_param(parser)
+    add_sub_agents_param(parser)
 
 
 def set_train_core_arguments(parser: argparse.ArgumentParser) -> None: