rasa-pro 3.11.0rc2__py3-none-any.whl → 3.11.1__py3-none-any.whl

rasa/shared/providers/_utils.py

@@ -0,0 +1,79 @@
+import structlog
+
+from rasa.shared.constants import (
+    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_CONFIG_KEY,
+    AWS_SESSION_TOKEN_ENV_VAR,
+)
+from rasa.shared.exceptions import ProviderClientValidationError
+from litellm import validate_environment
+from rasa.shared.providers.embedding._base_litellm_embedding_client import (
+    _VALIDATE_ENVIRONMENT_MISSING_KEYS_KEY,
+)
+
+structlogger = structlog.get_logger()
+
+
+def validate_aws_setup_for_litellm_clients(
+    litellm_model_name: str, litellm_call_kwargs: dict, source_log: str
+) -> None:
+    """Validates the AWS setup for LiteLLM clients to ensure all required
+    environment variables or corresponding call kwargs 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.
+
+    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, []
+    )
+    # 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
+        ]
+        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."
+        )
+        structlogger.error(
+            f"{source_log}.validate_aws_environment_variables",
+            event_info=event_info,
+            missing_environment_variables=missing_environment_variables,
+        )
+        raise ProviderClientValidationError(event_info)

rasa/shared/providers/embedding/default_litellm_embedding_client.py

@@ -1,8 +1,13 @@
 from typing import Any, Dict
 
+from rasa.shared.constants import (
+    AWS_BEDROCK_PROVIDER,
+    AWS_SAGEMAKER_PROVIDER,
+)
 from rasa.shared.providers._configs.default_litellm_client_config import (
     DefaultLiteLLMClientConfig,
 )
+from rasa.shared.providers._utils import validate_aws_setup_for_litellm_clients
 from rasa.shared.providers.embedding._base_litellm_embedding_client import (
     _BaseLiteLLMEmbeddingClient,
 )
@@ -100,3 +105,22 @@ class DefaultLiteLLMEmbeddingClient(_BaseLiteLLMEmbeddingClient):
             "model": self._litellm_model_name,
             **self._litellm_extra_parameters,
         }
+
+    def validate_client_setup(self) -> None:
+        # TODO: Temporarily disable environment variable validation for AWS setup
+        #       (Bedrock and SageMaker) until resolved by either:
+        #       1. An update from the LiteLLM package addressing the issue.
+        #       2. The implementation of a Bedrock client on our end.
+        #       ---
+        #       This fix ensures a consistent user experience for Bedrock (and
+        #       SageMaker) in Rasa by allowing AWS secrets to be provided as extra
+        #       parameters without triggering validation errors due to missing AWS
+        #       environment variables.
+        if self.provider.lower() in [AWS_BEDROCK_PROVIDER, AWS_SAGEMAKER_PROVIDER]:
+            validate_aws_setup_for_litellm_clients(
+                self._litellm_model_name,
+                self._litellm_extra_parameters,
+                "default_litellm_embedding_client",
+            )
+        else:
+            super().validate_client_setup()

rasa/shared/providers/embedding/litellm_router_embedding_client.py

@@ -72,7 +72,7 @@ class LiteLLMRouterEmbeddingClient(
         return cls(
             model_group_id=client_config.model_group_id,
             model_configurations=client_config.litellm_model_list,
-            router_settings=client_config.router,
+            router_settings=client_config.litellm_router_settings,
             **client_config.extra_parameters,
         )
 

rasa/shared/providers/llm/_base_litellm_client.py

@@ -221,6 +221,32 @@ class _BaseLiteLLMClient:
         )
         return formatted_response
 
+    def _format_text_completion_response(self, response: Any) -> LLMResponse:
+        """Parses the LiteLLM text completion response to Rasa format."""
+        formatted_response = LLMResponse(
+            id=response.id,
+            created=response.created,
+            choices=[choice.text for choice in response.choices],
+            model=response.model,
+        )
+        if (usage := response.usage) is not None:
+            prompt_tokens = (
+                num_tokens
+                if isinstance(num_tokens := usage.prompt_tokens, (int, float))
+                else 0
+            )
+            completion_tokens = (
+                num_tokens
+                if isinstance(num_tokens := usage.completion_tokens, (int, float))
+                else 0
+            )
+            formatted_response.usage = LLMUsage(prompt_tokens, completion_tokens)
+        structlogger.debug(
+            "base_litellm_client.formatted_response",
+            formatted_response=formatted_response.to_dict(),
+        )
+        return formatted_response
+
     @staticmethod
     def _ensure_certificates() -> None:
         """Configures SSL certificates for LiteLLM. This method is invoked during

rasa/shared/providers/llm/default_litellm_llm_client.py

@@ -1,8 +1,13 @@
 from typing import Dict, Any
 
+from rasa.shared.constants import (
+    AWS_BEDROCK_PROVIDER,
+    AWS_SAGEMAKER_PROVIDER,
+)
 from rasa.shared.providers._configs.default_litellm_client_config import (
     DefaultLiteLLMClientConfig,
 )
+from rasa.shared.providers._utils import validate_aws_setup_for_litellm_clients
 from rasa.shared.providers.llm._base_litellm_client import _BaseLiteLLMClient
 
 
@@ -82,3 +87,22 @@ class DefaultLiteLLMClient(_BaseLiteLLMClient):
         to the client provider and deployed model.
         """
         return self._extra_parameters
+
+    def validate_client_setup(self) -> None:
+        # TODO: Temporarily change the environment variable validation for AWS setup
+        #       (Bedrock and SageMaker) until resolved by either:
+        #       1. An update from the LiteLLM package addressing the issue.
+        #       2. The implementation of a Bedrock client on our end.
+        #       ---
+        #       This fix ensures a consistent user experience for Bedrock (and
+        #       SageMaker) in Rasa by allowing AWS secrets to be provided as extra
+        #       parameters without triggering validation errors due to missing AWS
+        #       environment variables.
+        if self.provider.lower() in [AWS_BEDROCK_PROVIDER, AWS_SAGEMAKER_PROVIDER]:
+            validate_aws_setup_for_litellm_clients(
+                self._litellm_model_name,
+                self._litellm_extra_parameters,
+                "default_litellm_llm_client",
+            )
+        else:
+            super().validate_client_setup()

rasa/shared/providers/llm/litellm_router_llm_client.py

@@ -68,15 +68,61 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
         return cls(
             model_group_id=client_config.model_group_id,
             model_configurations=client_config.litellm_model_list,
-            router_settings=client_config.router,
+            router_settings=client_config.litellm_router_settings,
+            use_chat_completions_endpoint=client_config.use_chat_completions_endpoint,
             **client_config.extra_parameters,
         )
 
+    @suppress_logs(log_level=logging.WARNING)
+    def _text_completion(self, prompt: Union[List[str], str]) -> LLMResponse:
+        """
+        Synchronously generate completions for given prompt.
+
+        Args:
+            prompt: Prompt to generate the completion for.
+        Returns:
+            List of message completions.
+        Raises:
+            ProviderClientAPIException: If the API request fails.
+        """
+        try:
+            response = self.router_client.text_completion(
+                prompt=prompt, **self._completion_fn_args
+            )
+            return self._format_text_completion_response(response)
+        except Exception as e:
+            raise ProviderClientAPIException(e)
+
+    @suppress_logs(log_level=logging.WARNING)
+    async def _atext_completion(self, prompt: Union[List[str], str]) -> LLMResponse:
+        """
+        Asynchronously generate completions for given prompt.
+
+        Args:
+            prompt: Prompt to generate the completion for.
+        Returns:
+            List of message completions.
+        Raises:
+            ProviderClientAPIException: If the API request fails.
+        """
+        try:
+            response = await self.router_client.atext_completion(
+                prompt=prompt, **self._completion_fn_args
+            )
+            return self._format_text_completion_response(response)
+        except Exception as e:
+            raise ProviderClientAPIException(e)
+
     @suppress_logs(log_level=logging.WARNING)
     def completion(self, messages: Union[List[str], str]) -> LLMResponse:
         """
         Synchronously generate completions for given list of messages.
 
+        Method overrides the base class method to call the appropriate
+        completion method based on the configuration. If the chat completions
+        endpoint is enabled, the completion method is called. Otherwise, the
+        text_completion method is called.
+
         Args:
             messages: List of messages or a single message to generate the
                 completion for.
@@ -85,6 +131,8 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
         Raises:
             ProviderClientAPIException: If the API request fails.
         """
+        if not self._use_chat_completions_endpoint:
+            return self._text_completion(messages)
         try:
             formatted_messages = self._format_messages(messages)
             response = self.router_client.completion(
@@ -99,6 +147,11 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
         """
         Asynchronously generate completions for given list of messages.
 
+        Method overrides the base class method to call the appropriate
+        completion method based on the configuration. If the chat completions
+        endpoint is enabled, the completion method is called. Otherwise, the
+        text_completion method is called.
+
         Args:
             messages: List of messages or a single message to generate the
                 completion for.
@@ -107,6 +160,8 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
         Raises:
             ProviderClientAPIException: If the API request fails.
         """
+        if not self._use_chat_completions_endpoint:
+            return await self._atext_completion(messages)
         try:
             formatted_messages = self._format_messages(messages)
             response = await self.router_client.acompletion(

rasa/shared/providers/llm/self_hosted_llm_client.py

@@ -10,13 +10,14 @@ import structlog
 from rasa.shared.constants import (
     SELF_HOSTED_VLLM_PREFIX,
     SELF_HOSTED_VLLM_API_KEY_ENV_VAR,
+    API_KEY,
 )
 from rasa.shared.providers._configs.self_hosted_llm_client_config import (
     SelfHostedLLMClientConfig,
 )
 from rasa.shared.exceptions import ProviderClientAPIException
 from rasa.shared.providers.llm._base_litellm_client import _BaseLiteLLMClient
-from rasa.shared.providers.llm.llm_response import LLMResponse, LLMUsage
+from rasa.shared.providers.llm.llm_response import LLMResponse
 from rasa.shared.utils.io import suppress_logs
 
 structlogger = structlog.get_logger()
@@ -61,7 +62,8 @@ class SelfHostedLLMClient(_BaseLiteLLMClient):
         self._api_version = api_version
         self._use_chat_completions_endpoint = use_chat_completions_endpoint
         self._extra_parameters = kwargs or {}
-        self._apply_dummy_api_key_if_missing()
+        if self._extra_parameters.get(API_KEY) is None:
+            self._apply_dummy_api_key_if_missing()
 
     @classmethod
     def from_config(cls, config: Dict[str, Any]) -> "SelfHostedLLMClient":
@@ -259,32 +261,6 @@ class SelfHostedLLMClient(_BaseLiteLLMClient):
             return super().completion(messages)
         return self._text_completion(messages)
 
-    def _format_text_completion_response(self, response: Any) -> LLMResponse:
-        """Parses the LiteLLM text completion response to Rasa format."""
-        formatted_response = LLMResponse(
-            id=response.id,
-            created=response.created,
-            choices=[choice.text for choice in response.choices],
-            model=response.model,
-        )
-        if (usage := response.usage) is not None:
-            prompt_tokens = (
-                num_tokens
-                if isinstance(num_tokens := usage.prompt_tokens, (int, float))
-                else 0
-            )
-            completion_tokens = (
-                num_tokens
-                if isinstance(num_tokens := usage.completion_tokens, (int, float))
-                else 0
-            )
-            formatted_response.usage = LLMUsage(prompt_tokens, completion_tokens)
-        structlogger.debug(
-            "base_litellm_client.formatted_response",
-            formatted_response=formatted_response.to_dict(),
-        )
-        return formatted_response
-
     @staticmethod
     def _apply_dummy_api_key_if_missing() -> None:
         if not os.getenv(SELF_HOSTED_VLLM_API_KEY_ENV_VAR):

rasa/shared/providers/router/_base_litellm_router_client.py

@@ -1,4 +1,5 @@
 from typing import Any, Dict, List
+import os
 import structlog
 
 from litellm import Router
@@ -7,6 +8,12 @@ from rasa.shared.constants import (
     MODEL_LIST_KEY,
     MODEL_GROUP_ID_CONFIG_KEY,
     ROUTER_CONFIG_KEY,
+    SELF_HOSTED_VLLM_PREFIX,
+    SELF_HOSTED_VLLM_API_KEY_ENV_VAR,
+    LITELLM_PARAMS_KEY,
+    API_KEY,
+    MODEL_CONFIG_KEY,
+    USE_CHAT_COMPLETIONS_ENDPOINT_CONFIG_KEY,
 )
 from rasa.shared.exceptions import ProviderClientValidationError
 from rasa.shared.providers._configs.litellm_router_client_config import (
@@ -42,12 +49,15 @@ class _BaseLiteLLMRouterClient:
         model_group_id: str,
         model_configurations: List[Dict[str, Any]],
         router_settings: Dict[str, Any],
+        use_chat_completions_endpoint: bool = True,
         **kwargs: Any,
     ):
         self._model_group_id = model_group_id
         self._model_configurations = model_configurations
         self._router_settings = router_settings
+        self._use_chat_completions_endpoint = use_chat_completions_endpoint
         self._extra_parameters = kwargs or {}
+        self.additional_client_setup()
         try:
             resolved_model_configurations = (
                 self._resolve_env_vars_in_model_configurations()
@@ -67,6 +77,21 @@ class _BaseLiteLLMRouterClient:
             )
             raise ProviderClientValidationError(f"{event_info} Original error: {e}")
 
+    def additional_client_setup(self) -> None:
+        """Additional setup for the LiteLLM Router client."""
+        # If the model configuration is self-hosted VLLM, set a dummy API key if not
+        # provided. A bug in the LiteLLM library requires an API key to be set even if
+        # it is not required.
+        for model_configuration in self.model_configurations:
+            if (
+                f"{SELF_HOSTED_VLLM_PREFIX}/"
+                in model_configuration[LITELLM_PARAMS_KEY][MODEL_CONFIG_KEY]
+                and API_KEY not in model_configuration[LITELLM_PARAMS_KEY]
+                and not os.getenv(SELF_HOSTED_VLLM_API_KEY_ENV_VAR)
+            ):
+                os.environ[SELF_HOSTED_VLLM_API_KEY_ENV_VAR] = "dummy api key"
+                return
+
     @classmethod
     def from_config(cls, config: Dict[str, Any]) -> "_BaseLiteLLMRouterClient":
         """Instantiates a LiteLLM Router Embedding client from a configuration dict.
@@ -95,7 +120,8 @@ class _BaseLiteLLMRouterClient:
         return cls(
             model_group_id=client_config.model_group_id,
             model_configurations=client_config.litellm_model_list,
-            router_settings=client_config.router,
+            router_settings=client_config.litellm_router_settings,
+            use_chat_completions_endpoint=client_config.use_chat_completions_endpoint,
             **client_config.extra_parameters,
         )
 
@@ -119,6 +145,11 @@ class _BaseLiteLLMRouterClient:
         """Returns the instantiated LiteLLM Router client."""
         return self._router_client
 
+    @property
+    def use_chat_completions_endpoint(self) -> bool:
+        """Returns whether to use the chat completions endpoint."""
+        return self._use_chat_completions_endpoint
+
     @property
     def _litellm_extra_parameters(self) -> Dict[str, Any]:
         """
@@ -136,6 +167,9 @@ class _BaseLiteLLMRouterClient:
             MODEL_GROUP_ID_CONFIG_KEY: self.model_group_id,
             MODEL_LIST_KEY: self.model_configurations,
             ROUTER_CONFIG_KEY: self.router_settings,
+            USE_CHAT_COMPLETIONS_ENDPOINT_CONFIG_KEY: (
+                self.use_chat_completions_endpoint
+            ),
             **self._litellm_extra_parameters,
         }
 

rasa/shared/utils/common.py

@@ -3,14 +3,16 @@ import functools
 import importlib
 import inspect
 import logging
+import os
 import pkgutil
 import sys
 from types import ModuleType
-from typing import Text, Dict, Optional, Any, List, Callable, Collection, Type
+from typing import Sequence, Text, Dict, Optional, Any, List, Callable, Collection, Type
 
 import rasa.shared.utils.io
+from rasa.exceptions import MissingDependencyException
 from rasa.shared.constants import DOCS_URL_MIGRATION_GUIDE
-from rasa.shared.exceptions import RasaException
+from rasa.shared.exceptions import ProviderClientValidationError, RasaException
 
 logger = logging.getLogger(__name__)
 
@@ -193,7 +195,7 @@ def mark_as_experimental_feature(feature_name: Text) -> None:
 def mark_as_beta_feature(feature_name: Text) -> None:
     """Warns users that they are using a beta feature."""
     logger.warning(
-        f"🔬 Beta Feature: {feature_name} is in beta. It may have unexpected"
+        f"🔬 Beta Feature: {feature_name} is in beta. It may have unexpected "
         "behaviour and might be changed in the future."
     )
 
@@ -295,3 +297,28 @@ def warn_and_exit_if_module_path_contains_rasa_plus(
             docs=DOCS_URL_MIGRATION_GUIDE,
         )
         sys.exit(1)
+
+
+def validate_environment(
+    required_env_vars: Sequence[str],
+    required_packages: Sequence[str],
+    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
+         component_name: component name that needs the requirements
+    """
+    for e in required_env_vars:
+        if not os.environ.get(e):
+            raise ProviderClientValidationError(
+                f"Missing environment variable for {component_name}: {e}"
+            )
+    for p in required_packages:
+        try:
+            importlib.import_module(p)
+        except ImportError:
+            raise MissingDependencyException(
+                f"Missing package for {component_name}: {p}"
+            )

rasa/shared/utils/health_check/health_check.py

@@ -1,4 +1,5 @@
 import os
+import sys
 from typing import Optional, Dict, Any
 
 from rasa.shared.constants import (
@@ -9,7 +10,6 @@ from rasa.shared.constants import (
 from rasa.shared.exceptions import ProviderClientValidationError
 from rasa.shared.providers.embedding.embedding_client import EmbeddingClient
 from rasa.shared.providers.llm.llm_client import LLMClient
-from rasa.shared.utils.cli import print_error_and_exit
 from rasa.shared.utils.llm import llm_factory, structlogger, embedder_factory
 
 
@@ -25,15 +25,15 @@ def try_instantiate_llm_client(
     except (ProviderClientValidationError, ValueError) as e:
         structlogger.error(
             f"{log_source_function}.llm_instantiation_failed",
-            message="Unable to instantiate LLM client.",
+            event_info=(
+                f"Unable to create the LLM client for component - "
+                f"{log_source_component}. "
+                f"Please make sure you specified the required environment variables "
+                f"and configuration keys. "
+            ),
             error=e,
         )
-        print_error_and_exit(
-            f"Unable to create the LLM client for component - {log_source_component}. "
-            f"Please make sure you specified the required environment variables "
-            f"and configuration keys. "
-            f"Error: {e}"
-        )
+        sys.exit(1)
 
 
 def try_instantiate_embedder(
@@ -48,14 +48,14 @@ def try_instantiate_embedder(
     except (ProviderClientValidationError, ValueError) as e:
         structlogger.error(
             f"{log_source_function}.embedder_instantiation_failed",
-            message="Unable to instantiate Embedding client.",
+            event_info=(
+                f"Unable to create the Embedding client for component - "
+                f"{log_source_component}. Please make sure you specified the required "
+                f"environment variables and configuration keys."
+            ),
             error=e,
         )
-        print_error_and_exit(
-            f"Unable to create the Embedding client for component - "
-            f"{log_source_component}. Please make sure you specified the required "
-            f"environment variables and configuration keys. Error: {e}"
-        )
+        sys.exit(1)
 
 
 def perform_llm_health_check(
@@ -202,13 +202,14 @@ def send_test_llm_api_request(
     except Exception as e:
         structlogger.error(
             f"{log_source_function}.send_test_llm_api_request_failed",
-            event_info="Test call to the LLM API failed.",
+            event_info=(
+                f"Test call to the LLM API failed for component - "
+                f"{log_source_component}.",
+            ),
+            config=llm_client.config,
             error=e,
         )
-        print_error_and_exit(
-            f"Test call to the LLM API failed for component - {log_source_component}. "
-            f"Error: {e}"
-        )
+        sys.exit(1)
 
 
 def send_test_embeddings_api_request(
@@ -232,13 +233,14 @@ def send_test_embeddings_api_request(
     except Exception as e:
         structlogger.error(
             f"{log_source_function}.send_test_llm_api_request_failed",
-            event_info="Test call to the Embeddings API failed.",
+            event_info=(
+                f"Test call to the Embeddings API failed for component - "
+                f"{log_source_component}."
+            ),
+            config=embedder.config,
             error=e,
         )
-        print_error_and_exit(
-            f"Test call to the Embeddings API failed for component - "
-            f"{log_source_component}. Error: {e}"
-        )
+        sys.exit(1)
 
 
 def is_api_health_check_enabled() -> bool: