rasa-pro 3.10.15__py3-none-any.whl → 3.11.0__py3-none-any.whl

rasa/shared/providers/embedding/litellm_router_embedding_client.py

@@ -0,0 +1,135 @@
+from typing import Any, Dict, List
+import logging
+import structlog
+
+from rasa.shared.exceptions import ProviderClientAPIException
+from rasa.shared.providers._configs.litellm_router_client_config import (
+    LiteLLMRouterClientConfig,
+)
+from rasa.shared.providers.embedding._base_litellm_embedding_client import (
+    _BaseLiteLLMEmbeddingClient,
+)
+from rasa.shared.providers.embedding.embedding_response import EmbeddingResponse
+from rasa.shared.providers.router._base_litellm_router_client import (
+    _BaseLiteLLMRouterClient,
+)
+from rasa.shared.utils.io import suppress_logs
+
+structlogger = structlog.get_logger()
+
+
+class LiteLLMRouterEmbeddingClient(
+    _BaseLiteLLMRouterClient, _BaseLiteLLMEmbeddingClient
+):
+    """A client for interfacing with LiteLLM Router Embedding endpoints.
+
+    Parameters:
+        model_group_id (str): The model group ID.
+        model_configurations (List[Dict[str, Any]]): The list of model configurations.
+        router_settings (Dict[str, Any]): The router settings.
+        kwargs (Optional[Dict[str, Any]]): Additional configuration parameters.
+
+    Raises:
+        ProviderClientValidationError: If validation of the client setup fails.
+    """
+
+    def __init__(
+        self,
+        model_group_id: str,
+        model_configurations: List[Dict[str, Any]],
+        router_settings: Dict[str, Any],
+        **kwargs: Any,
+    ):
+        super().__init__(
+            model_group_id, model_configurations, router_settings, **kwargs
+        )
+
+    @classmethod
+    def from_config(cls, config: Dict[str, Any]) -> "LiteLLMRouterEmbeddingClient":
+        """Instantiates a LiteLLM Router Embedding client from a configuration dict.
+
+        Args:
+            config: (Dict[str, Any]) The configuration dictionary.
+
+        Returns:
+            LiteLLMRouterLLMClient: The instantiated LiteLLM Router LLM client.
+
+        Raises:
+            ValueError: If the configuration is invalid.
+        """
+        try:
+            client_config = LiteLLMRouterClientConfig.from_dict(config)
+        except (KeyError, ValueError) as e:
+            message = "Cannot instantiate a client from the passed configuration."
+            structlogger.error(
+                "litellm_router_llm_client.from_config.error",
+                message=message,
+                config=config,
+                original_error=e,
+            )
+            raise
+
+        return cls(
+            model_group_id=client_config.model_group_id,
+            model_configurations=client_config.litellm_model_list,
+            router_settings=client_config.litellm_router_settings,
+            **client_config.extra_parameters,
+        )
+
+    @suppress_logs(log_level=logging.WARNING)
+    def embed(self, documents: List[str]) -> EmbeddingResponse:
+        """
+        Embeds a list of documents synchronously.
+
+        Args:
+            documents: List of documents to be embedded.
+
+        Returns:
+            List of embedding vectors.
+
+        Raises:
+            ProviderClientAPIException: If API calls raised an error.
+        """
+        self.validate_documents(documents)
+        try:
+            response = self.router_client.embedding(
+                input=documents, **self._embedding_fn_args
+            )
+            return self._format_response(response)
+        except Exception as e:
+            raise ProviderClientAPIException(
+                message="Failed to embed documents", original_exception=e
+            )
+
+    @suppress_logs(log_level=logging.WARNING)
+    async def aembed(self, documents: List[str]) -> EmbeddingResponse:
+        """
+        Embeds a list of documents asynchronously.
+
+        Args:
+            documents: List of documents to be embedded.
+
+        Returns:
+            List of embedding vectors.
+
+        Raises:
+            ProviderClientAPIException: If API calls raised an error.
+        """
+        self.validate_documents(documents)
+        try:
+            response = await self.router_client.aembedding(
+                input=documents, **self._embedding_fn_args
+            )
+            return self._format_response(response)
+        except Exception as e:
+            raise ProviderClientAPIException(
+                message="Failed to embed documents", original_exception=e
+            )
+
+    @property
+    def _embedding_fn_args(self) -> Dict[str, Any]:
+        """Returns the arguments to be passed to the embedding function."""
+        return {
+            **self._litellm_extra_parameters,
+            "model": self._model_group_id,
+        }

rasa/shared/providers/llm/_base_litellm_client.py

@@ -1,6 +1,6 @@
+import logging
 from abc import abstractmethod
 from typing import Dict, List, Any, Union
-import logging
 
 import structlog
 from litellm import (
@@ -9,7 +9,7 @@ from litellm import (
     validate_environment,
 )
 
-from rasa.shared.constants import API_BASE_CONFIG_KEY
+from rasa.shared.constants import API_BASE_CONFIG_KEY, API_KEY
 from rasa.shared.exceptions import (
     ProviderClientAPIException,
     ProviderClientValidationError,
@@ -19,7 +19,7 @@ from rasa.shared.providers._ssl_verification_utils import (
     ensure_ssl_certificates_for_litellm_openai_based_clients,
 )
 from rasa.shared.providers.llm.llm_response import LLMResponse, LLMUsage
-from rasa.shared.utils.io import suppress_logs
+from rasa.shared.utils.io import suppress_logs, resolve_environment_variables
 
 structlogger = structlog.get_logger()
 
@@ -99,12 +99,12 @@ class _BaseLiteLLMClient:
             ProviderClientValidationError if validation fails.
         """
         self._validate_environment_variables()
-        self._validate_api_key_not_in_config()
 
     def _validate_environment_variables(self) -> None:
         """Validate that the required environment variables are set."""
         validation_info = validate_environment(
             self._litellm_model_name,
+            api_key=self._litellm_extra_parameters.get(API_KEY),
             api_base=self._litellm_extra_parameters.get(API_BASE_CONFIG_KEY),
         )
         if missing_environment_variables := validation_info.get(
@@ -121,18 +121,6 @@ class _BaseLiteLLMClient:
             )
             raise ProviderClientValidationError(event_info)
 
-    def _validate_api_key_not_in_config(self) -> None:
-        if "api_key" in self._litellm_extra_parameters:
-            event_info = (
-                "API Key is set through `api_key` extra parameter."
-                "Set API keys through environment variables."
-            )
-            structlogger.error(
-                "base_litellm_client.validate_api_key_not_in_config",
-                event_info=event_info,
-            )
-            raise ProviderClientValidationError(event_info)
-
     @suppress_logs(log_level=logging.WARNING)
     def completion(self, messages: Union[List[str], str]) -> LLMResponse:
         """Synchronously generate completions for given list of messages.
@@ -149,9 +137,8 @@ class _BaseLiteLLMClient:
         """
         try:
             formatted_messages = self._format_messages(messages)
-            response = completion(
-                messages=formatted_messages, **self._completion_fn_args
-            )
+            arguments = resolve_environment_variables(self._completion_fn_args)
+            response = completion(messages=formatted_messages, **arguments)
             return self._format_response(response)
         except Exception as e:
             raise ProviderClientAPIException(e)
@@ -172,9 +159,8 @@ class _BaseLiteLLMClient:
         """
         try:
             formatted_messages = self._format_messages(messages)
-            response = await acompletion(
-                messages=formatted_messages, **self._completion_fn_args
-            )
+            arguments = resolve_environment_variables(self._completion_fn_args)
+            response = await acompletion(messages=formatted_messages, **arguments)
             return self._format_response(response)
         except Exception as e:
             message = ""
@@ -235,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/azure_openai_llm_client.py

@@ -17,6 +17,7 @@ from rasa.shared.constants import (
     OPENAI_API_KEY_ENV_VAR,
     AZURE_API_TYPE_ENV_VAR,
     AZURE_OPENAI_PROVIDER,
+    API_KEY,
 )
 from rasa.shared.exceptions import ProviderClientValidationError
 from rasa.shared.providers._configs.azure_openai_client_config import (
@@ -29,8 +30,7 @@ structlogger = structlog.get_logger()
 
 
 class AzureOpenAILLMClient(_BaseLiteLLMClient):
-    """
-    A client for interfacing with Azure's OpenAI LLM deployments.
+    """A client for interfacing with Azure's OpenAI LLM deployments.
 
     Parameters:
         deployment (str): The deployment name.
@@ -80,11 +80,7 @@ class AzureOpenAILLMClient(_BaseLiteLLMClient):
             or os.getenv(OPENAI_API_VERSION_ENV_VAR)
         )
 
-        # API key can be set through OPENAI_API_KEY too,
-        # because of the backward compatibility
-        self._api_key = os.getenv(AZURE_API_KEY_ENV_VAR) or os.getenv(
-            OPENAI_API_KEY_ENV_VAR
-        )
+        self._api_key_env_var = self._resolve_api_key_env_var()
 
         # Not used by LiteLLM, here for backward compatibility
         self._api_type = (
@@ -117,11 +113,6 @@ class AzureOpenAILLMClient(_BaseLiteLLMClient):
                 "env_var": AZURE_API_VERSION_ENV_VAR,
                 "deprecated_var": OPENAI_API_VERSION_ENV_VAR,
             },
-            "API Key": {
-                "current_value": self._api_key,
-                "env_var": AZURE_API_KEY_ENV_VAR,
-                "deprecated_var": OPENAI_API_KEY_ENV_VAR,
-            },
         }
 
         deprecation_warning_message = (
@@ -154,10 +145,51 @@ class AzureOpenAILLMClient(_BaseLiteLLMClient):
                     )
                 raise_deprecation_warning(message=message)
 
+    def _resolve_api_key_env_var(self) -> str:
+        """Resolves the environment variable to use for the API key.
+
+        Returns:
+            str: The env variable in dollar syntax format to use for the API key.
+        """
+        if API_KEY in self._extra_parameters:
+            # API key is set to an env var in the config itself
+            # in case the model is defined in the endpoints.yml
+            return self._extra_parameters[API_KEY]
+
+        if os.getenv(AZURE_API_KEY_ENV_VAR) is not None:
+            return "${AZURE_API_KEY}"
+
+        if os.getenv(OPENAI_API_KEY_ENV_VAR) is not None:
+            # API key can be set through OPENAI_API_KEY too,
+            # because of the backward compatibility
+            raise_deprecation_warning(
+                message=(
+                    f"Usage of '{OPENAI_API_KEY_ENV_VAR}' environment variable "
+                    "for setting the API key for Azure OpenAI "
+                    "client is deprecated and will be removed "
+                    f"in 4.0.0. Please use '{AZURE_API_KEY_ENV_VAR}' "
+                    "environment variable."
+                )
+            )
+            return "${OPENAI_API_KEY}"
+
+        structlogger.error(
+            "azure_openai_llm_client.api_key_not_set",
+            event_info=(
+                "API key not set, it is required for API calls. "
+                f"Set it either via the environment variable"
+                f"'{AZURE_API_KEY_ENV_VAR}' or directly"
+                f"via the config key '{API_KEY}'."
+            ),
+        )
+        raise ProviderClientValidationError(
+            f"Missing required environment variable/config key '{API_KEY}' for "
+            f"API calls."
+        )
+
     @classmethod
     def from_config(cls, config: Dict[str, Any]) -> "AzureOpenAILLMClient":
-        """
-        Initializes the client from given configuration.
+        """Initializes the client from given configuration.
 
         Args:
             config (Dict[str, Any]): Configuration.
@@ -212,23 +244,17 @@ class AzureOpenAILLMClient(_BaseLiteLLMClient):
 
     @property
     def model(self) -> Optional[str]:
-        """
-        Returns the name of the model deployed on Azure.
-        """
+        """Returns the name of the model deployed on Azure."""
         return self._model
 
     @property
     def api_base(self) -> Optional[str]:
-        """
-        Returns the API base URL for the Azure OpenAI llm client.
-        """
+        """Returns the API base URL for the Azure OpenAI llm client."""
         return self._api_base
 
     @property
     def api_version(self) -> Optional[str]:
-        """
-        Returns the API version for the Azure OpenAI llm client.
-        """
+        """Returns the API version for the Azure OpenAI llm client."""
         return self._api_version
 
     @property
@@ -261,7 +287,7 @@ class AzureOpenAILLMClient(_BaseLiteLLMClient):
             {
                 "api_base": self.api_base,
                 "api_version": self.api_version,
-                "api_key": self._api_key,
+                "api_key": self._api_key_env_var,
             }
         )
         return fn_args
@@ -305,11 +331,6 @@ class AzureOpenAILLMClient(_BaseLiteLLMClient):
                 "env_var": None,
                 "config_key": DEPLOYMENT_CONFIG_KEY,
             },
-            "API Key": {
-                "current_value": self._api_key,
-                "env_var": AZURE_API_KEY_ENV_VAR,
-                "config_key": None,
-            },
         }
 
         missing_settings = [

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

@@ -0,0 +1,182 @@
+from typing import Any, Dict, List, Union
+import logging
+import structlog
+
+from rasa.shared.exceptions import ProviderClientAPIException
+from rasa.shared.providers._configs.litellm_router_client_config import (
+    LiteLLMRouterClientConfig,
+)
+from rasa.shared.providers.llm._base_litellm_client import _BaseLiteLLMClient
+from rasa.shared.providers.llm.llm_response import LLMResponse
+from rasa.shared.providers.router._base_litellm_router_client import (
+    _BaseLiteLLMRouterClient,
+)
+from rasa.shared.utils.io import suppress_logs
+
+structlogger = structlog.get_logger()
+
+
+class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
+    """A client for interfacing with LiteLLM Router LLM endpoints.
+
+    Parameters:
+        model_group_id (str): The model group ID.
+        model_configurations (List[Dict[str, Any]]): The list of model configurations.
+        router_settings (Dict[str, Any]): The router settings.
+        kwargs (Optional[Dict[str, Any]]): Additional configuration parameters.
+
+    Raises:
+        ProviderClientValidationError: If validation of the client setup fails.
+    """
+
+    def __init__(
+        self,
+        model_group_id: str,
+        model_configurations: List[Dict[str, Any]],
+        router_settings: Dict[str, Any],
+        **kwargs: Any,
+    ):
+        super().__init__(
+            model_group_id, model_configurations, router_settings, **kwargs
+        )
+
+    @classmethod
+    def from_config(cls, config: Dict[str, Any]) -> "LiteLLMRouterLLMClient":
+        """Instantiates a LiteLLM Router LLM client from a configuration dict.
+
+        Args:
+            config: (Dict[str, Any]) The configuration dictionary.
+
+        Returns:
+            LiteLLMRouterLLMClient: The instantiated LiteLLM Router LLM client.
+
+        Raises:
+            ValueError: If the configuration is invalid.
+        """
+        try:
+            client_config = LiteLLMRouterClientConfig.from_dict(config)
+        except (KeyError, ValueError) as e:
+            message = "Cannot instantiate a client from the passed configuration."
+            structlogger.error(
+                "litellm_router_llm_client.from_config.error",
+                message=message,
+                config=config,
+                original_error=e,
+            )
+            raise
+
+        return cls(
+            model_group_id=client_config.model_group_id,
+            model_configurations=client_config.litellm_model_list,
+            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.
+        Returns:
+            List of message completions.
+        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(
+                messages=formatted_messages, **self._completion_fn_args
+            )
+            return self._format_response(response)
+        except Exception as e:
+            raise ProviderClientAPIException(e)
+
+    @suppress_logs(log_level=logging.WARNING)
+    async def acompletion(self, messages: Union[List[str], str]) -> LLMResponse:
+        """
+        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.
+        Returns:
+            List of message completions.
+        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(
+                messages=formatted_messages, **self._completion_fn_args
+            )
+            return self._format_response(response)
+        except Exception as e:
+            raise ProviderClientAPIException(e)
+
+    @property
+    def _completion_fn_args(self) -> Dict[str, Any]:
+        """Returns the completion arguments for invoking a call through
+        LiteLLM's completion functions.
+        """
+        return {
+            **self._litellm_extra_parameters,
+            "model": self.model_group_id,
+        }