rasa-pro 3.14.0.dev6__tar.gz → 3.14.0.dev7__tar.gz

{rasa_pro-3.14.0.dev6 → rasa_pro-3.14.0.dev7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: rasa-pro
-Version: 3.14.0.dev6
+Version: 3.14.0.dev7
 Summary: State-of-the-art open-core Conversational AI framework for Enterprises that natively leverages generative AI for effortless assistant development.
 Keywords: nlp,machine-learning,machine-learning-library,bot,bots,botkit,rasa conversational-agents,conversational-ai,chatbot,chatbot-framework,bot-framework
 Author: Rasa Technologies GmbH

{rasa_pro-3.14.0.dev6 → rasa_pro-3.14.0.dev7}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "rasa-pro"
-version = "3.14.0.dev6"
+version = "3.14.0.dev7"
 description = "State-of-the-art open-core Conversational AI framework for Enterprises that natively leverages generative AI for effortless assistant development."
 authors = [ "Rasa Technologies GmbH <hi@rasa.com>",]
 maintainers = [ "Tom Bocklisch <tom@rasa.com>",]
@@ -361,6 +361,7 @@ convention = "google"
 
 
 
+
 [tool.poetry.group.dev.dependencies]
 ruff = ">=0.3.5,<0.4.0"
 docker = "~7.1.0"

rasa_pro-3.14.0.dev7/rasa/agents/exceptions.py

@@ -0,0 +1,38 @@
+from rasa.exceptions import RasaException, ValidationError
+
+
+class AgentNotFoundException(RasaException):
+    """Raised when an agent is not found."""
+
+    def __init__(self, agent_name: str):
+        super().__init__(f"The agent {agent_name} is not defined.")
+
+
+class DuplicatedAgentNameException(ValidationError):
+    """Raised when agent names are duplicated."""
+
+    def __init__(self, duplicated_names: list[str]) -> None:
+        """Initialize the exception."""
+        event_info = f"Agent names are duplicated: {', '.join(duplicated_names)}"
+
+        super().__init__(
+            code="agent.duplicated_name",
+            event_info=event_info,
+            duplicated_names=duplicated_names,
+        )
+
+
+class AgentNameFlowConflictException(ValidationError):
+    """Raised when agent names conflict with flow names."""
+
+    def __init__(self, conflicting_names: list[str]) -> None:
+        """Initialize the exception."""
+        event_info = (
+            f"Agent names conflict with flow names: {', '.join(conflicting_names)}"
+        )
+
+        super().__init__(
+            code="agent.flow_name_conflict",
+            event_info=event_info,
+            conflicting_names=conflicting_names,
+        )

{rasa_pro-3.14.0.dev6 → rasa_pro-3.14.0.dev7}/rasa/agents/protocol/a2a/a2a_agent.py

@@ -239,7 +239,7 @@ class A2AAgent(AgentProtocol):
                 error_message=f"Streaming error: {exception!s}",
             )
 
-        # In case we didn't receive any response or the response ended unexpectedly
+        # In case we didn't receive any response or the stream ended unexpectedly
         structlogger.error(
             "a2a_agent.run_streaming_agent.unexpected_end",
             event_info="Unexpected end of streaming response from A2A server",

{rasa_pro-3.14.0.dev6 → rasa_pro-3.14.0.dev7}/rasa/agents/protocol/mcp/mcp_base_agent.py

@@ -253,7 +253,8 @@ class MCPBaseAgent(AgentProtocol):
             except ConnectionError as ce:
                 structlogger.warning(
                     "mcp_agent.connect.connection_error",
-                    event_info=f"Connection attempt {attempt} failed: {ce}",
+                    event_info=f"Connection attempt {attempt} failed.",
+                    error=str(ce),
                     attempt=attempt,
                     max_retries=self._max_retries,
                 )
@@ -296,7 +297,7 @@ class MCPBaseAgent(AgentProtocol):
         except Exception as e:
             event_info = (
                 f"Agent `{self._name}` failed to connect to MCP server - "
-                f"{server_name} @ {server_config.url}: {e}"
+                f"`{server_name}` @ `{server_config.url}`: {e!s}"
             )
             structlogger.error(
                 "mcp_agent.connect.failed_to_connect",
@@ -333,7 +334,7 @@ class MCPBaseAgent(AgentProtocol):
         except Exception as e:
             structlogger.error(
                 "mcp_agent.disconnect_server.error",
-                event_info=f"Failed to disconnect from server `{server_name}`: {e}",
+                event_info=f"Failed to disconnect from server `{server_name}`: {e!s}",
                 server_name=server_name,
                 agent_name=self._name,
                 agent_id=str(make_agent_identifier(self._name, self.protocol_type)),
@@ -566,8 +567,8 @@ class MCPBaseAgent(AgentProtocol):
 
         server_id = self._tool_to_server_mapper[tool_name]
         connection = self._server_connections[server_id]
-        session = await connection.ensure_active_session()
         try:
+            session = await connection.ensure_active_session()
             result = await session.call_tool(tool_name, arguments)
             return AgentToolResult.from_mcp_tool_result(tool_name, result)
         except Exception as e:
@@ -577,7 +578,7 @@ class MCPBaseAgent(AgentProtocol):
                 is_error=True,
                 error_message=(
                     f"Failed to execute tool `{tool_name}` via MCP server `{server_id}`"
-                    f"@ {connection.server_url}: {e}"
+                    f" @ `{connection.server_url}`: {e!s}"
                 ),
             )
 
@@ -606,7 +607,7 @@ class MCPBaseAgent(AgentProtocol):
                 tool_name=tool_name,
                 result=None,
                 is_error=True,
-                error_message=f"Failed to execute built-in tool `{tool_name}`: {e}",
+                error_message=f"Failed to execute built-in tool `{tool_name}`: {e!s}",
             )
         return await self._execute_mcp_tool(tool_name, arguments)
 

{rasa_pro-3.14.0.dev6 → rasa_pro-3.14.0.dev7}/rasa/agents/protocol/mcp/mcp_task_agent.py

@@ -98,7 +98,7 @@ class MCPTaskAgent(MCPBaseAgent):
             for name in re.findall(r"\bslots\.(\w+)", condition)
         }
 
-        slot_names = agent_input.slot_names
+        slot_names = [slot.name for slot in agent_input.slots]
 
         # Keep only slots that actually exist in agent_input.slots
         valid_slot_names = [

{rasa_pro-3.14.0.dev6 → rasa_pro-3.14.0.dev7}/rasa/agents/utils.py

@@ -2,6 +2,7 @@ from typing import List, Optional
 
 from rasa.agents.agent_manager import AgentManager
 from rasa.agents.exceptions import AgentNotFoundException
+from rasa.agents.validation import validate_agent_names_not_conflicting_with_flows
 from rasa.core.available_agents import (
     AgentConfig,
     AgentMCPServerConfig,
@@ -46,6 +47,10 @@ async def initialize_agents(
     agent_manager: AgentManager = AgentManager()
     endpoints = AvailableEndpoints.get_instance()
 
+    # Validate agent names don't conflict with flow names
+    flow_names = {flow.id for flow in flows.underlying_flows}
+    validate_agent_names_not_conflicting_with_flows(sub_agents.agents, flow_names)
+
     for flow in flows.underlying_flows:
         for step in flow.steps:
             if isinstance(step, CallFlowStep):

rasa_pro-3.14.0.dev7/rasa/agents/validation.py

@@ -0,0 +1,484 @@
+"""Validation functions for agent configurations with reduced redundancies."""
+
+import os
+import urllib.parse
+from collections import Counter
+from typing import Any, Dict, List, NoReturn, Set
+
+from pydantic import ValidationError as PydanticValidationError
+from ruamel import yaml
+
+from rasa.agents.exceptions import (
+    AgentNameFlowConflictException,
+    DuplicatedAgentNameException,
+)
+from rasa.core.available_agents import (
+    DEFAULT_AGENTS_CONFIG_FOLDER,
+    AgentConfig,
+    AgentConfiguration,
+    AgentConnections,
+    AgentInfo,
+    ProtocolConfig,
+)
+from rasa.core.available_endpoints import AvailableEndpoints
+from rasa.exceptions import ValidationError
+
+# Centralized allowed keys configuration to eliminate duplication
+ALLOWED_KEYS = {
+    "agent": {"name", "protocol", "description"},
+    "configuration": {
+        "llm",
+        "prompt_template",
+        "module",
+        "timeout",
+        "max_retries",
+        "agent_card",
+    },
+    "connections": {"mcp_servers"},
+}
+
+
+def validate_agent_names_unique(agents: List[AgentConfig]) -> None:
+    """Validate that agent names are unique across all loaded agents.
+
+    Args:
+        agents: List of agent configurations.
+
+    Raises:
+        DuplicatedAgentNameException: If agent names are not unique.
+    """
+    agent_names = [agent_config.agent.name for agent_config in agents]
+    name_counts = Counter(agent_names)
+
+    duplicated_names = [name for name, count in name_counts.items() if count > 1]
+
+    if duplicated_names:
+        raise DuplicatedAgentNameException(duplicated_names)
+
+
+def validate_agent_names_not_conflicting_with_flows(
+    agents: Dict[str, AgentConfig], flow_names: Set[str]
+) -> None:
+    """Validate that agent names do not conflict with flow names.
+
+    Args:
+        agents: Dictionary of agent configurations.
+        flow_names: Set of flow names to check against.
+
+    Raises:
+        AgentNameFlowConflictException: If agent names conflict with flow names.
+    """
+    conflicting_names = [
+        agent_name for agent_name in agents.keys() if agent_name in flow_names
+    ]
+
+    if conflicting_names:
+        raise AgentNameFlowConflictException(conflicting_names)
+
+
+def _validate_mcp_config(agent_config: AgentConfig) -> None:
+    """Validate MCP-specific configuration requirements."""
+    agent_name = agent_config.agent.name
+
+    # Check connections.mcp_servers exists
+    if agent_config.connections is None or agent_config.connections.mcp_servers is None:
+        raise ValidationError(
+            code="agent.validation.mcp.missing_connections",
+            event_info=f"For protocol 'MCP', agent '{agent_name}' must have "
+            "'connections.mcp_servers' configured.",
+        )
+
+    # Check mcp_servers list is not empty
+    if not agent_config.connections.mcp_servers:
+        raise ValidationError(
+            code="agent.validation.mcp.empty_servers_list",
+            event_info=f"For protocol 'MCP', agent '{agent_name}' must have "
+            "at least one MCP server configured in 'connections.mcp_servers'.",
+        )
+
+    # Validate each MCP server configuration
+    for i, server in enumerate(agent_config.connections.mcp_servers):
+        if not server.name:
+            raise ValidationError(
+                code="agent.validation.mcp.server_missing_name",
+                event_info=f"For protocol 'MCP', agent '{agent_name}' MCP server "
+                f"at index {i} must have a 'name' field.",
+            )
+
+
+def _validate_a2a_config(agent_config: AgentConfig) -> None:
+    """Validate A2A-specific configuration requirements."""
+    agent_name = agent_config.agent.name
+
+    # Check configuration.agent_card exists
+    if (
+        agent_config.configuration is None
+        or agent_config.configuration.agent_card is None
+    ):
+        raise ValidationError(
+            code="agent.validation.a2a.missing_agent_card",
+            event_info=f"For protocol 'A2A', agent '{agent_name}' must have "
+            "'configuration.agent_card' specified.",
+        )
+
+    # Validate agent_card path or URL
+    agent_card = agent_config.configuration.agent_card
+    if not agent_card:
+        raise ValidationError(
+            code="agent.validation.a2a.empty_agent_card",
+            event_info=f"Agent '{agent_name}' has empty 'agent_card' value",
+        )
+
+    # Check if it's a URL
+    parsed_url = urllib.parse.urlparse(agent_card)
+    if parsed_url.scheme and parsed_url.netloc:
+        # It's a URL, validate format
+        if parsed_url.scheme not in ["http", "https"]:
+            raise ValidationError(
+                code="agent.validation.a2a.invalid_agent_card_url",
+                event_info=f"Agent '{agent_name}' has invalid URL scheme "
+                f"'{parsed_url.scheme}' for 'agent_card'",
+            )
+    else:
+        # It's a file path, check if it exists
+        if not os.path.exists(agent_card):
+            raise ValidationError(
+                code="agent.validation.a2a.agent_card_file_not_found",
+                event_info=f"Agent '{agent_name}' has 'agent_card' file that doesn't "
+                f"exist: {agent_card}",
+            )
+
+
+def _validate_optional_keys(agent_config: AgentConfig) -> None:
+    """Validate optional keys in agent configuration."""
+    agent_name = agent_config.agent.name
+
+    # Validate prompt_template if present
+    if agent_config.configuration and agent_config.configuration.prompt_template:
+        prompt_path = agent_config.configuration.prompt_template
+        if not os.path.exists(prompt_path):
+            raise ValidationError(
+                code="agent.validation.optional.prompt_template_not_found",
+                event_info=f"Agent '{agent_name}' has prompt template that "
+                f"does not exist: {prompt_path}",
+            )
+
+    # Validate module if present
+    if agent_config.configuration and agent_config.configuration.module:
+        import importlib
+
+        module_name = agent_config.configuration.module
+        try:
+            module_path, class_name = module_name.rsplit(".", 1)
+            getattr(importlib.import_module(module_path), class_name)
+        except (ImportError, AttributeError) as e:
+            raise ValidationError(
+                code="agent.validation.optional.module_not_found",
+                event_info=f"Agent '{agent_name}' has module '{module_name}' "
+                f"that could not be imported: {e}",
+            )
+
+
+def _validate_llm_references(
+    llm_config: Dict[str, Any],
+    endpoints: "AvailableEndpoints",
+    agent_name: str,
+) -> None:
+    """Validate LLM configuration references against endpoints."""
+    if "model_group" in llm_config:
+        from rasa.engine.validation import (
+            _validate_component_model_client_config_has_references_to_endpoints,
+        )
+
+        component_config = {"llm": llm_config}
+        _validate_component_model_client_config_has_references_to_endpoints(
+            component_config, "llm", component_name=agent_name
+        )
+
+
+def _validate_mcp_server_references(
+    mcp_servers: list, endpoints: "AvailableEndpoints", agent_name: str
+) -> None:
+    """Validate MCP server references against endpoints."""
+    if not endpoints.mcp_servers:
+        raise ValidationError(
+            code="agent.validation.endpoints.no_mcp_servers",
+            event_info=(
+                f"Agent '{agent_name}' references MCP servers but no MCP "
+                "servers are defined in endpoints.yml."
+            ),
+        )
+
+    available_mcp_server_names = [server.name for server in endpoints.mcp_servers]
+
+    for i, mcp_server in enumerate(mcp_servers):
+        server_name = mcp_server.name
+        if server_name not in available_mcp_server_names:
+            raise ValidationError(
+                code="agent.validation.endpoints.invalid_mcp_server",
+                event_info=(
+                    f"MCP server '{server_name}' at index {i} for Agent "
+                    f"'{agent_name}' does not exist in endpoints.yml. Available MCP "
+                    f"servers: {', '.join(available_mcp_server_names)}"
+                ),
+            )
+
+
+def _handle_pydantic_validation_error(
+    error: PydanticValidationError, agent_name: str
+) -> NoReturn:
+    """Handle specific Pydantic validation errors that are actually possible.
+
+    Args:
+        error: The Pydantic validation error to handle
+        agent_name: Name of the agent for error messages
+    """
+    missing_fields = []
+    invalid_protocol = False
+    type_error = None
+
+    for pydantic_error in error.errors():
+        error_type = pydantic_error["type"]
+        field_path = ".".join(str(loc) for loc in pydantic_error["loc"])
+
+        if error_type == "missing":
+            for field in ["name", "protocol", "description"]:
+                if field in field_path:
+                    missing_fields.append(field)
+        elif error_type == "enum" and "protocol" in field_path:
+            invalid_protocol = True
+        elif error_type in ["string_type", "int_parsing"]:
+            type_error = (field_path, pydantic_error["msg"])
+
+    # Handle missing required fields
+    if missing_fields:
+        raise ValidationError(
+            code="agent.validation.mandatory.fields_missing",
+            event_info=(
+                f"Agent '{agent_name}' is missing required fields "
+                f"in agent section: {', '.join(missing_fields)}"
+            ),
+        )
+
+    # Handle invalid protocol
+    elif invalid_protocol:
+        raise ValidationError(
+            code="agent.validation.pydantic.invalid_protocol",
+            event_info=(
+                f"Agent '{agent_name}' has invalid protocol value. "
+                "Supported protocols: MCP, A2A"
+            ),
+        )
+
+    # Handle type errors
+    elif type_error:
+        field, msg = type_error
+        raise ValidationError(
+            code="agent.validation.pydantic.type_error",
+            event_info=(
+                f"Agent '{agent_name}' has invalid type for field " f"'{field}': {msg}"
+            ),
+        )
+
+    # Handle other Pydantic validation errors
+    else:
+        raise ValidationError(
+            code="agent.validation.pydantic.failed",
+            event_info=f"Agent '{agent_name}' validation failed: {error}",
+        )
+
+
+def _validate_endpoint_references(agent_config: AgentConfig) -> None:
+    """Validate that LLM and MCP server references in agent config are valid."""
+    agent_name = agent_config.agent.name
+
+    # Get available endpoints
+    endpoints = AvailableEndpoints.get_instance()
+
+    # Validate LLM configuration references
+    if agent_config.configuration and agent_config.configuration.llm:
+        _validate_llm_references(agent_config.configuration.llm, endpoints, agent_name)
+
+    # Validate MCP server references
+    if agent_config.connections and agent_config.connections.mcp_servers:
+        _validate_mcp_server_references(
+            agent_config.connections.mcp_servers, endpoints, agent_name
+        )
+
+
+def _validate_section_keys(
+    data: Dict[str, Any], section: str, allowed_keys: set
+) -> None:
+    """Generic function to validate keys in a specific section."""
+    if section not in data:
+        return
+
+    section_data = data[section]
+    if not isinstance(section_data, dict):
+        return
+
+    additional_keys = set(section_data.keys()) - allowed_keys
+    if additional_keys:
+        agent_name = data.get("agent", {}).get("name", "unknown")
+        raise ValidationError(
+            code=f"agent.validation.structure.additional_{section}_keys",
+            event_info=(
+                f"Agent '{agent_name}' contains additional keys in "
+                f"'{section}' section: {', '.join(sorted(additional_keys))}"
+            ),
+        )
+
+
+def _validate_mandatory_fields(data: Dict[str, Any], agent_name: str) -> None:
+    """Validate that all mandatory fields are present in the agent section."""
+    if "agent" not in data:
+        raise ValidationError(
+            code="agent.validation.mandatory.agent_section_missing",
+            event_info=f"Agent '{agent_name}' is missing 'agent' section",
+        )
+
+    agent_data = data["agent"]
+    if not isinstance(agent_data, dict):
+        raise ValidationError(
+            code="agent.validation.mandatory.agent_section_invalid",
+            event_info=(
+                f"Agent '{agent_name}' has invalid 'agent' section - "
+                "must be a dictionary"
+            ),
+        )
+
+    # Check for required fields
+    missing_fields = []
+    for field in ["name", "protocol", "description"]:
+        if field not in agent_data or not agent_data[field]:
+            missing_fields.append(field)
+
+    if missing_fields:
+        raise ValidationError(
+            code="agent.validation.mandatory.fields_missing",
+            event_info=(
+                f"Agent '{agent_name}' is missing required fields in agent section: "
+                f"{', '.join(missing_fields)}"
+            ),
+        )
+
+
+def _validate_no_additional_keys_raw_data(data: Dict[str, Any]) -> None:
+    """Validate that no additional, unexpected keys are present in the raw data."""
+    # Use the generic validation function for each section
+    _validate_section_keys(data, "agent", ALLOWED_KEYS["agent"])
+    _validate_section_keys(data, "configuration", ALLOWED_KEYS["configuration"])
+    _validate_section_keys(data, "connections", ALLOWED_KEYS["connections"])
+
+
+def _validate_folder_structure(agent_folder: str) -> None:
+    """Validate agent folder structure."""
+    if not os.path.exists(agent_folder):
+        raise ValidationError(
+            code="agent.validation.folder.not_found",
+            event_info=f"Agent folder does not exist: {agent_folder}",
+        )
+
+    if not os.path.isdir(agent_folder):
+        raise ValidationError(
+            code="agent.validation.folder.not_directory",
+            event_info=f"Agent folder is not a directory: {agent_folder}",
+        )
+
+
+def _validate_config_file_exists(config_path: str, agent_name: str) -> None:
+    """Validate config.yml file exists."""
+    if not os.path.isfile(config_path):
+        raise ValidationError(
+            code="agent.validation.folder.missing_config",
+            event_info=f"Agent '{agent_name}' is missing 'config.yml' file",
+        )
+
+
+def validate_agent_config(agent_config: AgentConfig) -> None:
+    """Validate an agent configuration using all applicable validators."""
+    protocol = agent_config.agent.protocol
+
+    # Run protocol-specific validation
+    if protocol == ProtocolConfig.MCP:
+        _validate_mcp_config(agent_config)
+    elif protocol == ProtocolConfig.A2A:
+        _validate_a2a_config(agent_config)
+
+    # Run optional keys validation
+    _validate_optional_keys(agent_config)
+
+    # Run endpoint references validation
+    _validate_endpoint_references(agent_config)
+
+
+def validate_agent_folder(agent_folder: str = DEFAULT_AGENTS_CONFIG_FOLDER) -> None:
+    """Validate all agent configurations in a folder."""
+    # Validate folder structure
+    _validate_folder_structure(agent_folder)
+
+    # Scan for agent folders
+    for agent_folder_name in os.listdir(agent_folder):
+        agent_path = os.path.join(agent_folder, agent_folder_name)
+
+        if not os.path.isdir(agent_path):
+            continue
+
+        config_path = os.path.join(agent_path, "config.yml")
+
+        # Validate config file exists
+        _validate_config_file_exists(config_path, agent_folder_name)
+
+        # Read and validate the config content
+        try:
+            # First read the raw YAML data to validate structure
+            with open(config_path, "r") as f:
+                data = yaml.safe_load(f)
+
+            # Validate no additional keys
+            _validate_no_additional_keys_raw_data(data)
+
+            # Validate mandatory fields before creating Pydantic models
+            _validate_mandatory_fields(data, agent_folder_name)
+
+            # Create the agent config using AvailableAgents
+            from rasa.core.available_agents import AvailableAgents
+
+            agent_config = AvailableAgents._read_agent_config(config_path)
+
+            # Validate the agent config (protocol-specific and endpoint references)
+            validate_agent_config(agent_config)
+        except PydanticValidationError as e:
+            _handle_pydantic_validation_error(e, agent_folder_name)
+        except Exception as e:
+            # Handle non-Pydantic exceptions
+            raise ValidationError(
+                code="agent.validation.folder.config_validation_failed",
+                event_info=f"Agent '{agent_folder_name}' validation failed: {e}",
+            )
+
+
+def validate_agent_config_data(data: Dict[str, Any]) -> "AgentConfig":
+    """Validate agent configuration data."""
+    try:
+        # Validate no additional keys before creating Pydantic models
+        _validate_no_additional_keys_raw_data(data)
+
+        # Create the agent config (this will trigger Pydantic validation)
+        agent_config = AgentConfig(
+            agent=AgentInfo(**data.get("agent", {})),
+            configuration=AgentConfiguration(**data.get("configuration", {}))
+            if data.get("configuration")
+            else None,
+            connections=AgentConnections(**data.get("connections", {}))
+            if data.get("connections")
+            else None,
+        )
+
+        return agent_config
+
+    except PydanticValidationError as e:
+        _handle_pydantic_validation_error(e, "Agent configuration")
+    except Exception:
+        # Re-raise other exceptions
+        raise

{rasa_pro-3.14.0.dev6 → rasa_pro-3.14.0.dev7}/rasa/api.py

@@ -71,22 +71,23 @@ def run(
 
 
 def train(
-    domain: Text,
-    config: Text,
-    training_files: "Union[Text, List[Text]]",
-    output: Text = rasa.shared.constants.DEFAULT_MODELS_PATH,
+    domain: str,
+    config: str,
+    training_files: "Union[str, List[str]]",
+    output: str = rasa.shared.constants.DEFAULT_MODELS_PATH,
     dry_run: bool = False,
     force_training: bool = False,
-    fixed_model_name: Optional[Text] = None,
+    fixed_model_name: Optional[str] = None,
     persist_nlu_training_data: bool = False,
     core_additional_arguments: Optional[Dict] = None,
     nlu_additional_arguments: Optional[Dict] = None,
-    model_to_finetune: Optional[Text] = None,
+    model_to_finetune: Optional[str] = None,
     finetuning_epoch_fraction: float = 1.0,
     remote_storage: Optional[StorageType] = None,
     file_importer: Optional["TrainingDataImporter"] = None,
     keep_local_model_copy: bool = False,
     remote_root_only: bool = False,
+    sub_agents: Optional[str] = None,
 ) -> "TrainingResult":
     """Runs Rasa Core and NLU training in `async` loop.
 
@@ -116,6 +117,7 @@ def train(
             remote storage is configured.
         remote_root_only: If `True`, the model will be stored in the root of the
             remote model storage.
+        sub_agents: Path to sub-agents directory.
 
     Returns:
         An instance of `TrainingResult`.
@@ -140,6 +142,7 @@ def train(
             file_importer=file_importer,
             keep_local_model_copy=keep_local_model_copy,
             remote_root_only=remote_root_only,
+            sub_agents=sub_agents,
         )
     )