rasa-pro 3.14.0.dev7__py3-none-any.whl → 3.14.0.dev9__py3-none-any.whl

rasa/shared/exceptions.py

@@ -163,6 +163,10 @@ class ProviderClientAPIException(RasaException):
         return s
 
 
+class LLMToolResponseDecodeError(ProviderClientAPIException):
+    """Raised when a JSON decoding error occurs in LLM tool response."""
+
+
 class ProviderClientValidationError(RasaException):
     """Raised for errors that occur during validation of the API client."""
 
@@ -173,3 +177,11 @@ class FinetuningDataPreparationException(RasaException):
 
 class AgentInitializationException(RasaException):
     """Raised when there is an error during the initialization of an agent."""
+
+
+class AgentAuthInitializationException(RasaException):
+    """Raised when there is an error during the initialization of agent auth client."""
+
+
+class AuthenticationError(RasaException):
+    """Raised when there is an authentication error."""

rasa/shared/importers/importer.py

@@ -207,6 +207,12 @@ class TrainingDataImporter(ABC):
                 )
             ]
 
+        return TrainingDataImporter.wrap_in_builtins(importers)
+
+    @staticmethod
+    def wrap_in_builtins(
+        importers: List["TrainingDataImporter"],
+    ) -> "TrainingDataImporter":
         return LanguageImporter(
             E2EImporter(
                 FlowSyncImporter(ResponsesSyncImporter(CombinedDataImporter(importers)))

rasa/shared/importers/utils.py

@@ -1,9 +1,37 @@
-from typing import Iterable, List, Optional, Text
+from typing import Any, Dict, Iterable, List, Optional, Text
+
+from pydantic import BaseModel, Field
 
 from rasa.shared.core.domain import Domain
 from rasa.shared.core.flows import FlowsList
+from rasa.shared.core.flows.yaml_flows_io import KEY_FLOWS, get_flows_as_json
 from rasa.shared.core.training_data.structures import StoryGraph
+from rasa.shared.importers.importer import TrainingDataImporter
+from rasa.shared.nlu.training_data.formats.rasa_yaml import RasaYAMLWriter
 from rasa.shared.nlu.training_data.training_data import TrainingData
+from rasa.utils.json_utils import extract_values
+
+
+class CALMUserData(BaseModel):
+    """All pieces that will be uploaded to Rasa Studio."""
+
+    flows: Dict[str, Any] = Field(default_factory=dict)
+    domain: Dict[str, Any] = Field(default_factory=dict)
+    config: Dict[str, Any] = Field(default_factory=dict)
+    endpoints: Dict[str, Any] = Field(default_factory=dict)
+    nlu: Dict[str, Any] = Field(default_factory=dict)
+
+
+DOMAIN_KEYS = [
+    "version",
+    "actions",
+    "responses",
+    "slots",
+    "intents",
+    "entities",
+    "forms",
+    "session_config",
+]
 
 
 def training_data_from_paths(paths: Iterable[Text], language: Text) -> TrainingData:
@@ -40,3 +68,51 @@ def flows_from_paths(files: List[Text], domain: Optional[Domain] = None) -> Flow
         )
     flows.validate(domain)
     return flows
+
+
+def extract_calm_import_parts_from_importer(
+    importer: TrainingDataImporter,
+    config: Optional[Dict[str, Any]] = None,
+    endpoints: Optional[Dict[str, Any]] = None,
+) -> CALMUserData:
+    """Extracts CALMUserData from a TrainingDataImporter.
+
+    Args:
+        importer: The training data importer
+        data_paths: The path(s) to the training data for flows
+        config: Optional config dict, if not provided will use importer.get_config()
+        endpoints: Optional endpoints dict, defaults to empty dict
+
+    Returns:
+        CALMUserData containing flows, domain, config, endpoints, and nlu data
+    """
+    # Extract config
+    if config is None:
+        config = importer.get_config()
+
+    # Extract domain
+    domain_from_files = importer.get_user_domain().as_dict()
+    domain = extract_values(domain_from_files, DOMAIN_KEYS)
+
+    # Extract flows
+    flows = importer.get_user_flows()
+    flows_dict = {KEY_FLOWS: get_flows_as_json(flows)}
+
+    # Extract NLU data
+    nlu_data = importer.get_nlu_data()
+    nlu_examples = nlu_data.filter_training_examples(
+        lambda ex: ex.get("intent") in nlu_data.intents
+    )
+    nlu_dict = RasaYAMLWriter().training_data_to_dict(nlu_examples)
+
+    # Use provided endpoints or default to empty dict
+    if endpoints is None:
+        endpoints = {}
+
+    return CALMUserData(
+        flows=flows_dict or {},
+        domain=domain or {},
+        config=config or {},
+        endpoints=endpoints or {},
+        nlu=nlu_dict or {},
+    )

rasa/shared/nlu/training_data/schemas/responses.yml

@@ -72,6 +72,9 @@ schema;responses:
               allowempty: True
           channel:
             type: "str"
+          allow_interruptions:
+            type: "bool"
+            required: False
           metadata:
             type: "any"
           condition:

rasa/shared/providers/_utils.py

@@ -1,87 +1,103 @@
 from typing import Any, Dict, Optional
 
+import boto3
 import structlog
-from litellm import validate_environment
+from botocore.exceptions import BotoCoreError, ClientError
 
 from rasa.shared.constants import (
     API_BASE_CONFIG_KEY,
     API_VERSION_CONFIG_KEY,
     AWS_ACCESS_KEY_ID_CONFIG_KEY,
-    AWS_ACCESS_KEY_ID_ENV_VAR,
+    AWS_BEDROCK_PROVIDER,
     AWS_REGION_NAME_CONFIG_KEY,
-    AWS_REGION_NAME_ENV_VAR,
+    AWS_SAGEMAKER_CHAT_PROVIDER,
+    AWS_SAGEMAKER_PROVIDER,
     AWS_SECRET_ACCESS_KEY_CONFIG_KEY,
-    AWS_SECRET_ACCESS_KEY_ENV_VAR,
     AWS_SESSION_TOKEN_CONFIG_KEY,
-    AWS_SESSION_TOKEN_ENV_VAR,
     AZURE_API_BASE_ENV_VAR,
     AZURE_API_VERSION_ENV_VAR,
     DEPLOYMENT_CONFIG_KEY,
 )
 from rasa.shared.exceptions import ProviderClientValidationError
-from rasa.shared.providers.embedding._base_litellm_embedding_client import (
-    _VALIDATE_ENVIRONMENT_MISSING_KEYS_KEY,
-)
+from rasa.shared.utils.io import resolve_environment_variables
 
 structlogger = structlog.get_logger()
 
 
 def validate_aws_setup_for_litellm_clients(
-    litellm_model_name: str, litellm_call_kwargs: dict, source_log: str
+    litellm_model_name: str, litellm_call_kwargs: Dict, source_log: str, provider: str
 ) -> None:
-    """Validates the AWS setup for LiteLLM clients to ensure all required
-    environment variables or corresponding call kwargs are set.
+    """Validates the AWS setup for LiteLLM clients to ensure credentials are set.
 
     Args:
         litellm_model_name (str): The name of the LiteLLM model being validated.
         litellm_call_kwargs (dict): Additional keyword arguments passed to the client,
             which may include configuration values for AWS credentials.
         source_log (str): The source log identifier for structured logging.
+        provider (str): The provider for which the validation is being performed.
 
     Raises:
         ProviderClientValidationError: If any required AWS environment variable
             or corresponding configuration key is missing.
     """
-
-    # Mapping of environment variable names to their corresponding config keys
-    envs_to_args = {
-        AWS_ACCESS_KEY_ID_ENV_VAR: AWS_ACCESS_KEY_ID_CONFIG_KEY,
-        AWS_SECRET_ACCESS_KEY_ENV_VAR: AWS_SECRET_ACCESS_KEY_CONFIG_KEY,
-        AWS_REGION_NAME_ENV_VAR: AWS_REGION_NAME_CONFIG_KEY,
-        AWS_SESSION_TOKEN_ENV_VAR: AWS_SESSION_TOKEN_CONFIG_KEY,
-    }
-
-    # Validate the environment setup for the model
-    validation_info = validate_environment(litellm_model_name)
-    missing_environment_variables = validation_info.get(
-        _VALIDATE_ENVIRONMENT_MISSING_KEYS_KEY, []
+    # expand environment variables if referenced in the config
+    resolved_litellm_call_kwargs: Dict = resolve_environment_variables(
+        litellm_call_kwargs
+    )  # type: ignore[assignment]
+
+    # boto3 only accepts bedrock and sagemaker as valid clients
+    # therefore we need to convert the provider name if it is defined
+    # as sagemaker_chat
+    provider = (
+        AWS_SAGEMAKER_PROVIDER if provider == AWS_SAGEMAKER_CHAT_PROVIDER else provider
     )
-    # Filter out missing environment variables that have been set trough arguments
-    # in extra parameters
-    missing_environment_variables = [
-        missing_env_var
-        for missing_env_var in missing_environment_variables
-        if litellm_call_kwargs.get(envs_to_args.get(missing_env_var)) is None
-    ]
 
-    if missing_environment_variables:
-        missing_environment_details = [
-            (
-                f"'{missing_env_var}' environment variable or "
-                f"'{envs_to_args.get(missing_env_var)}' config key"
-            )
-            for missing_env_var in missing_environment_variables
+    # if the AWS credentials are defined in the endpoints yaml model config,
+    # either as referenced secret env vars or direct values, we need to pass them
+    # to the boto3 client to ensure that the client can connect to the AWS service.
+    additional_kwargs: Dict[str, Any] = {}
+    if AWS_ACCESS_KEY_ID_CONFIG_KEY in resolved_litellm_call_kwargs:
+        additional_kwargs[AWS_ACCESS_KEY_ID_CONFIG_KEY] = resolved_litellm_call_kwargs[
+            AWS_ACCESS_KEY_ID_CONFIG_KEY
         ]
+    if AWS_SECRET_ACCESS_KEY_CONFIG_KEY in resolved_litellm_call_kwargs:
+        additional_kwargs[AWS_SECRET_ACCESS_KEY_CONFIG_KEY] = (
+            resolved_litellm_call_kwargs[AWS_SECRET_ACCESS_KEY_CONFIG_KEY]
+        )
+    if AWS_SESSION_TOKEN_CONFIG_KEY in resolved_litellm_call_kwargs:
+        additional_kwargs[AWS_SESSION_TOKEN_CONFIG_KEY] = resolved_litellm_call_kwargs[
+            AWS_SESSION_TOKEN_CONFIG_KEY
+        ]
+    if AWS_REGION_NAME_CONFIG_KEY in resolved_litellm_call_kwargs:
+        additional_kwargs["region_name"] = resolved_litellm_call_kwargs[
+            AWS_REGION_NAME_CONFIG_KEY
+        ]
+
+    try:
+        # We are using the boto3 client because it can discover the AWS credentials
+        # from the environment variables, credentials file, or IAM roles.
+        # This is necessary to ensure that the client can connect to the AWS service.
+        aws_client = boto3.client(provider, **additional_kwargs)
+
+        # Using different method calls available to different AWS clients
+        # to test the connection
+        if provider == AWS_SAGEMAKER_PROVIDER:
+            aws_client.list_models()
+        elif provider == AWS_BEDROCK_PROVIDER:
+            aws_client.get_model_invocation_logging_configuration()
+
+    except (ClientError, BotoCoreError) as exc:
         event_info = (
-            f"The following environment variables or configuration keys are "
-            f"missing: "
-            f"{', '.join(missing_environment_details)}. "
-            f"These settings are required for API calls."
+            f"Failed to validate AWS setup for LiteLLM clients: {exc}. "
+            f"Ensure that you are using one of the available authentication methods:"
+            f"credentials file, environment variables, or IAM roles. "
+            f"Also, ensure that the AWS region is set correctly. "
         )
         structlogger.error(
-            f"{source_log}.validate_aws_environment_variables",
+            f"{source_log}.validate_aws_credentials_for_litellm_clients",
             event_info=event_info,
-            missing_environment_variables=missing_environment_variables,
+            exception=str(exc),
+            model_name=litellm_model_name,
         )
         raise ProviderClientValidationError(event_info)
 

rasa/shared/providers/embedding/default_litellm_embedding_client.py

@@ -37,6 +37,7 @@ class DefaultLiteLLMEmbeddingClient(_BaseLiteLLMEmbeddingClient):
 
     @classmethod
     def from_config(cls, config: Dict[str, Any]) -> "DefaultLiteLLMEmbeddingClient":
+        """Creates a DefaultLiteLLMEmbeddingClient instance from a config dict."""
         default_config = DefaultLiteLLMClientConfig.from_dict(config)
         return cls(
             model=default_config.model,
@@ -121,6 +122,7 @@ class DefaultLiteLLMEmbeddingClient(_BaseLiteLLMEmbeddingClient):
                 self._litellm_model_name,
                 self._litellm_extra_parameters,
                 "default_litellm_embedding_client",
+                provider=self.provider,
             )
         else:
             super().validate_client_setup()

rasa/shared/providers/llm/_base_litellm_client.py

@@ -158,7 +158,7 @@ class _BaseLiteLLMClient:
             )
             return self._format_response(response)
         except Exception as e:
-            raise ProviderClientAPIException(e)
+            raise ProviderClientAPIException(e) from e
 
     @suppress_logs(log_level=logging.WARNING)
     async def acompletion(
@@ -209,7 +209,7 @@ class _BaseLiteLLMClient:
                     "In case you are getting OpenAI connection errors, such as missing "
                     "API key, your configuration is incorrect."
                 )
-            raise ProviderClientAPIException(e, message)
+            raise ProviderClientAPIException(e, message) from e
 
     def _get_formatted_messages(
         self, messages: Union[List[dict], List[str], str]

rasa/shared/providers/llm/default_litellm_llm_client.py

@@ -39,6 +39,7 @@ class DefaultLiteLLMClient(_BaseLiteLLMClient):
 
     @classmethod
     def from_config(cls, config: Dict[str, Any]) -> DefaultLiteLLMClient:
+        """Creates a DefaultLiteLLMClient instance from a configuration dictionary."""
         default_config = DefaultLiteLLMClientConfig.from_dict(config)
         return cls(
             model=default_config.model,
@@ -110,6 +111,7 @@ class DefaultLiteLLMClient(_BaseLiteLLMClient):
                 self._litellm_model_name,
                 self._litellm_extra_parameters,
                 "default_litellm_llm_client",
+                provider=self.provider,
             )
         else:
             super().validate_client_setup()

rasa/shared/providers/llm/llm_response.py

@@ -9,7 +9,7 @@ from litellm.utils import ChatCompletionMessageToolCall
 from pydantic import BaseModel
 
 from rasa.shared.constants import KEY_TOOL_CALLS
-from rasa.shared.exceptions import ProviderClientAPIException
+from rasa.shared.exceptions import LLMToolResponseDecodeError
 
 structlogger = structlog.get_logger()
 
@@ -75,13 +75,13 @@ class LLMToolCall(BaseModel):
                 tool_name=data.function.name,
                 tool_call=data.function.arguments,
             )
-            raise ProviderClientAPIException(
+            raise LLMToolResponseDecodeError(
                 original_exception=e,
                 message=(
-                    f"Invalid arguments for tool call `{data.function.name}`: "
+                    f"Invalid arguments for tool call - `{data.function.name}`: "
                     f"`{data.function.arguments}`"
                 ),
-            )
+            ) from e
 
         return cls(
             id=data.id,

rasa/shared/utils/common.py

@@ -17,6 +17,7 @@ from typing import (
     Optional,
     Sequence,
     Text,
+    Tuple,
     Type,
 )
 
@@ -102,9 +103,12 @@ def sort_list_of_dicts_by_first_key(dicts: List[Dict]) -> List[Dict]:
 
 def cached_method(f: Callable[..., Any]) -> Callable[..., Any]:
     """Caches method calls based on the call's `args` and `kwargs`.
+
     Works for `async` and `sync` methods. Don't apply this to functions.
+
     Args:
         f: The decorated method whose return value should be cached.
+
     Returns:
         The return value which the method gives for the first call with the given
         arguments.
@@ -358,6 +362,7 @@ def validate_environment(
     component_name: str,
 ) -> None:
     """Make sure all needed requirements for a component are met.
+
     Args:
          required_env_vars: List of environment variables that should be set
          required_packages: List of packages that should be installed
@@ -389,3 +394,22 @@ Sign up at: https://feedback.rasa.com
 {separator}
 """
     print_success(message)
+
+
+def conditional_import(module_name: str, class_name: str) -> Tuple[Any, bool]:
+    """Conditionally import a class, returning (class, is_available) tuple.
+
+    Args:
+        module_name: The module path to import from
+        class_name: The class name to import
+
+    Returns:
+        A tuple of (class, is_available) where class is the imported class
+        or None if import failed, and is_available is a boolean indicating
+        whether the import was successful.
+    """
+    try:
+        module = __import__(module_name, fromlist=[class_name])
+        return getattr(module, class_name), True
+    except ImportError:
+        return None, False

rasa/shared/utils/llm.py

@@ -762,7 +762,7 @@ def get_prompt_template(
                 log_source_method=log_source_method,
             )
             return prompt_template
-    except (FileIOException, FileNotFoundException):
+    except (FileIOException, FileNotFoundException) as e:
         structlogger.warning(
             "utils.llm.get_prompt_template" ".failed_to_read_custom_prompt_template",
             event_info=(
@@ -770,6 +770,7 @@ def get_prompt_template(
             ),
             log_source_component=log_source_component,
             log_source_method=log_source_method,
+            error=str(e),
         )
     return default_prompt_template
 

rasa/shared/utils/mcp/server_connection.py

@@ -2,13 +2,18 @@
 
 import asyncio
 import warnings
-from contextlib import AsyncExitStack
-from typing import Any, Dict, Optional
+from contextlib import AsyncExitStack, asynccontextmanager
+from typing import Any, AsyncIterator, ClassVar, Dict, List, Optional
 
 import structlog
+from httpx import HTTPStatusError
 from mcp import ClientSession
 from mcp.client.streamable_http import streamablehttp_client
 
+from rasa.shared.agents.auth.agent_auth_manager import AgentAuthManager
+from rasa.shared.agents.auth.auth_strategy import AgentAuthStrategy
+from rasa.shared.exceptions import AuthenticationError
+
 structlogger = structlog.get_logger()
 
 
@@ -22,16 +27,26 @@ warnings.filterwarnings(
 
 class MCPServerConnection:
     """
-    Manages connection to an MCP server.
+    Manages connection to an MCP server, with optional authentication.
 
-    This class handles the lifecycle of connections to MCP servers,
-    including connection establishment, session management, and cleanup.
+    This class handles:
+    - Connection establishment and cleanup
+    - Session lifecycle management
+    - Authentication via AgentAuthManager (API Key, OAuth2, mTLS, etc.)
     """
 
     # Timeout for ping operations in seconds
     PING_TIMEOUT_SECONDS = 3.0
 
-    def __init__(self, server_name: str, server_url: str, server_type: str):
+    _SUPPORTED_SERVER_TYPES: ClassVar[List[str]] = ["http", "https"]
+
+    def __init__(
+        self,
+        server_name: str,
+        server_url: str,
+        server_type: str,
+        auth_manager: Optional[AgentAuthManager] = None,
+    ):
         """
         Initialize the MCP server connection.
 
@@ -39,43 +54,101 @@ class MCPServerConnection:
             server_name: Server name to identify the server
             server_url: Server URL
             server_type: Server type (currently only 'http' is supported)
+            auth_manager: Optional AgentAuthManager instance for this connection
         """
         self.server_name = server_name
         self.server_url = server_url
         self.server_type = server_type
+        self._auth_manager = auth_manager
         self.session: Optional[ClientSession] = None
         self.exit_stack: Optional[AsyncExitStack] = None
 
     @classmethod
     def from_config(cls, server_config: Dict[str, Any]) -> "MCPServerConnection":
         """Initialize the MCP server connection from a configuration dictionary."""
+        auth_config = server_config.get("additional_params")
+        _auth_manager = AgentAuthManager.load_auth(auth_config) if auth_config else None
         return cls(
             server_config["name"],
             server_config["url"],
             server_config.get("type", "http"),
+            _auth_manager,
         )
 
+    @staticmethod
+    @asynccontextmanager
+    async def open_mcp_session(
+        url: str, auth_strategy: Optional[AgentAuthStrategy] = None
+    ) -> AsyncIterator[ClientSession]:
+        """
+        Open a streamable MCP session, ensuring that initialization
+        completes before yielding.
+        """
+        async with streamablehttp_client(url, auth=auth_strategy) as (
+            read_stream,
+            write_stream,
+            _,
+        ):
+            async with ClientSession(read_stream, write_stream) as session:
+                await session.initialize()  # handshake done here
+                yield session
+
     async def connect(self) -> None:
         """Establish connection to the MCP server.
 
         Raises:
             ValueError: If the server type is not supported.
             ConnectionError: If connection fails.
+            AuthenticationError: If authentication fails.
         """
-        if self.server_type != "http":
+        if self.server_type not in self._SUPPORTED_SERVER_TYPES:
             raise ValueError(f"Unsupported server type: {self.server_type}")
 
         # Create a new exit stack for this connection to avoid task boundary issues
         self.exit_stack = AsyncExitStack()
 
         try:
-            read_stream, write_stream, _ = await self.exit_stack.enter_async_context(
-                streamablehttp_client(self.server_url)
+            # Get authentication strategy that adheres to httpx.Auth.
+            auth_strategy = (
+                self._auth_manager.get_auth() if self._auth_manager else None
             )
+
+            # Register the wrapped context manager in the stack
             self.session = await self.exit_stack.enter_async_context(
-                ClientSession(read_stream, write_stream)
+                self.open_mcp_session(self.server_url, auth_strategy)
             )
-            await self.session.initialize()
+
+        except Exception as eg:
+            for exc in getattr(eg, "exceptions", [eg]):
+                event_info = (
+                    f"Failed to connect to MCP server `{self.server_name}`: {exc!s}"
+                )
+                if isinstance(exc, HTTPStatusError):
+                    status_code = exc.response.status_code
+                    structlogger.error(
+                        "mcp_server_connection.connect.http_status_error",
+                        event_info=event_info,
+                        server_name=self.server_name,
+                        server_url=self.server_url,
+                        status_code=status_code,
+                        response_text=exc.response.reason_phrase,
+                    )
+                    await self._cleanup()
+                    if status_code in [400, 401, 403]:
+                        raise AuthenticationError(eg) from eg
+                    else:
+                        raise ConnectionError(eg) from eg
+                else:
+                    structlogger.error(
+                        "mcp_server_connection.connect.other_exception",
+                        event_info=event_info,
+                        server_name=self.server_name,
+                        server_url=self.server_url,
+                        error=str(exc),
+                    )
+            await self._cleanup()
+            raise ConnectionError(eg) from eg
+
         except asyncio.CancelledError as e:
             event_info = f"Connection to MCP server `{self.server_name}` was cancelled."
             structlogger.error(
@@ -88,18 +161,6 @@ class MCPServerConnection:
             await self._cleanup()
             raise ConnectionError(e) from e
 
-        except Exception as e:
-            event_info = f"Failed to connect to MCP server `{self.server_name}`: {e}"
-            structlogger.error(
-                "mcp_server_connection.connect.connection_failed",
-                event_info=event_info,
-                server_name=self.server_name,
-                server_url=self.server_url,
-            )
-            # Clean up on error
-            await self._cleanup()
-            raise ConnectionError(e) from e
-
     async def ensure_active_session(self) -> ClientSession:
         """
         Ensure an active session is available.

rasa/shared/utils/mcp/utils.py

@@ -0,0 +1,20 @@
+"""MCP utilities."""
+
+
+def mcp_server_exists(mcp_server: str) -> bool:
+    """Check if an MCP server exists in the configured endpoints.
+
+    Args:
+        mcp_server: The name of the MCP server to check.
+
+    Returns:
+        True if the MCP server exists, False otherwise.
+    """
+    from rasa.core.available_endpoints import AvailableEndpoints
+
+    mcp_server_list = AvailableEndpoints.get_instance().mcp_servers
+    if mcp_server_list is None:
+        return False
+
+    mcp_server_names = [server.name for server in mcp_server_list]
+    return mcp_server in mcp_server_names