rasa-pro 3.12.0.dev1__py3-none-any.whl

rasa/shared/providers/llm/_base_litellm_client.py

@@ -0,0 +1,265 @@
+from __future__ import annotations
+
+import logging
+from abc import abstractmethod
+from typing import Any, Dict, List, Union
+
+import structlog
+from litellm import (
+    acompletion,
+    completion,
+    validate_environment,
+)
+
+from rasa.shared.constants import API_BASE_CONFIG_KEY, API_KEY
+from rasa.shared.exceptions import (
+    ProviderClientAPIException,
+    ProviderClientValidationError,
+)
+from rasa.shared.providers._ssl_verification_utils import (
+    ensure_ssl_certificates_for_litellm_non_openai_based_clients,
+    ensure_ssl_certificates_for_litellm_openai_based_clients,
+)
+from rasa.shared.providers.llm.llm_response import LLMResponse, LLMUsage
+from rasa.shared.utils.io import resolve_environment_variables, suppress_logs
+
+structlogger = structlog.get_logger()
+
+_VALIDATE_ENVIRONMENT_MISSING_KEYS_KEY = "missing_keys"
+
+# Suppress LiteLLM info and debug logs - Global level.
+logging.getLogger("LiteLLM").setLevel(logging.WARNING)
+
+
+class _BaseLiteLLMClient:
+    """An abstract base class for LiteLLM clients.
+
+    This class defines the interface and common functionality for all clients
+    based on LiteLLM.
+
+    The class is made private to prevent it from being part of the
+    public-facing interface, as it serves as an internal base class
+    for specific implementations of clients that are currently based on
+    LiteLLM.
+
+    By keeping it private, we ensure that only the derived, concrete
+    implementations are exposed to users, maintaining a cleaner and
+    more controlled API surface.
+    """
+
+    def __init__(self):  # type: ignore
+        self._ensure_certificates()
+
+    @classmethod
+    @abstractmethod
+    def from_config(cls, config: Dict[str, Any]) -> _BaseLiteLLMClient:
+        pass
+
+    @property
+    @abstractmethod
+    def config(self) -> dict:
+        """Returns the configuration for that the llm client
+        in dictionary form.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def _litellm_model_name(self) -> str:
+        """Returns the value of LiteLLM's model parameter to be used in
+        completion/acompletion in LiteLLM format:
+
+        <provider>/<model or deployment name>
+        """
+        pass
+
+    @property
+    def _litellm_extra_parameters(self) -> Dict[str, Any]:
+        """Returns a dictionary of extra parameters which include model
+        parameters as well as LiteLLM specific input parameters.
+
+        By default, this returns an empty dictionary (no extra parameters).
+        """
+        return {}
+
+    @property
+    def _completion_fn_args(self) -> dict:
+        return {
+            **self._litellm_extra_parameters,
+            "model": self._litellm_model_name,
+            # Since all providers covered by LiteLLM use the OpenAI format, but
+            # not all support every OpenAI parameter, raise an exception if
+            # provider/model uses unsupported parameter
+            "drop_params": False,
+        }
+
+    def validate_client_setup(self) -> None:
+        """Perform client validation. By default only environment variables
+        are validated.
+
+        Raises:
+            ProviderClientValidationError if validation fails.
+        """
+        self._validate_environment_variables()
+
+    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(
+            _VALIDATE_ENVIRONMENT_MISSING_KEYS_KEY
+        ):
+            event_info = (
+                f"Environment variables: {missing_environment_variables} "
+                f"not set. Required for API calls."
+            )
+            structlogger.error(
+                "base_litellm_client.validate_environment_variables",
+                event_info=event_info,
+                missing_environment_variables=missing_environment_variables,
+            )
+            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.
+
+        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.
+        """
+        try:
+            formatted_messages = self._format_messages(messages)
+            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)
+
+    @suppress_logs(log_level=logging.WARNING)
+    async def acompletion(self, messages: Union[List[str], str]) -> LLMResponse:
+        """Asynchronously generate completions for given list of messages.
+
+        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.
+        """
+        try:
+            formatted_messages = self._format_messages(messages)
+            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 = ""
+            from rasa.shared.providers.llm.self_hosted_llm_client import (
+                SelfHostedLLMClient,
+            )
+
+            if isinstance(self, SelfHostedLLMClient):
+                message = (
+                    "If you are using 'provider=self-hosted' to call a hosted vllm "
+                    "server make sure your config is correctly setup. You should have "
+                    "the following mandatory keys in your config: "
+                    "provider=self-hosted; "
+                    "model='<your-vllm-model-name>'; "
+                    "api_base='your-hosted-vllm-serv'."
+                    "In case you are getting OpenAI connection errors, such as missing "
+                    "API key, your configuration is incorrect."
+                )
+            raise ProviderClientAPIException(e, message)
+
+    def _format_messages(self, messages: Union[List[str], str]) -> List[Dict[str, str]]:
+        """Formats messages (or a single message) to OpenAI format."""
+        if isinstance(messages, str):
+            messages = [messages]
+        return [{"content": message, "role": "user"} for message in messages]
+
+    def _format_response(self, response: Any) -> LLMResponse:
+        """Parses the LiteLLM response to Rasa format."""
+        formatted_response = LLMResponse(
+            id=response.id,
+            created=response.created,
+            choices=[choice.message.content for choice in response.choices],
+            model=response.model,
+        )
+        if (
+            response.model_extra
+            and (usage := response.model_extra.get("usage")) is not None
+        ):
+            # We use `.get()` for accessing litellm.utils.Usage attributes.
+            # litellm.utils.Usage does not set the attributes if
+            # `prompt_tokens` or `completion_tokens` are absent (None).
+            prompt_tokens = (
+                num_tokens
+                if isinstance(num_tokens := usage.get("prompt_tokens", 0), (int, float))
+                else 0
+            )
+            completion_tokens = (
+                num_tokens
+                if isinstance(
+                    num_tokens := usage.get("completion_tokens", 0), (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
+
+    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
+        client initialization.
+
+        LiteLLM may utilize `openai` clients or other providers that require
+        SSL verification settings through the `SSL_VERIFY` / `SSL_CERTIFICATE`
+        environment variables or the `litellm.ssl_verify` /
+        `litellm.ssl_certificate` global settings.
+
+        This method ensures proper SSL configuration for both cases.
+        """
+        ensure_ssl_certificates_for_litellm_non_openai_based_clients()
+        ensure_ssl_certificates_for_litellm_openai_based_clients()

rasa/shared/providers/llm/azure_openai_llm_client.py

@@ -0,0 +1,415 @@
+from __future__ import annotations
+
+import os
+import re
+from typing import Any, Dict, Optional
+
+import structlog
+
+from rasa.shared.constants import (
+    API_BASE_CONFIG_KEY,
+    API_KEY,
+    API_VERSION_CONFIG_KEY,
+    AZURE_API_BASE_ENV_VAR,
+    AZURE_API_KEY_ENV_VAR,
+    AZURE_API_TYPE_ENV_VAR,
+    AZURE_API_VERSION_ENV_VAR,
+    AZURE_OPENAI_PROVIDER,
+    DEPLOYMENT_CONFIG_KEY,
+    OPENAI_API_BASE_ENV_VAR,
+    OPENAI_API_KEY_ENV_VAR,
+    OPENAI_API_TYPE_ENV_VAR,
+    OPENAI_API_VERSION_ENV_VAR,
+)
+from rasa.shared.exceptions import ProviderClientValidationError
+from rasa.shared.providers._configs.azure_openai_client_config import (
+    AzureOAuthConfig,
+    AzureOpenAIClientConfig,
+)
+from rasa.shared.providers.constants import (
+    DEFAULT_AZURE_API_KEY_NAME,
+    LITE_LLM_API_BASE_FIELD,
+    LITE_LLM_API_KEY_FIELD,
+    LITE_LLM_API_VERSION_FIELD,
+    LITE_LLM_AZURE_AD_TOKEN,
+    LITE_LLM_AZURE_AD_TOKEN_PROVIDER,
+)
+from rasa.shared.providers.llm._base_litellm_client import _BaseLiteLLMClient
+from rasa.shared.utils.io import raise_deprecation_warning
+
+structlogger = structlog.get_logger()
+
+AZURE_CLIENT_ID = "AZURE_CLIENT_ID"
+AZURE_CLIENT_SECRET = "AZURE_CLIENT_SECRET"
+AZURE_TENANT_ID = "AZURE_TENANT_ID"
+CLIENT_SECRET_VARS = (AZURE_CLIENT_ID, AZURE_CLIENT_SECRET, AZURE_TENANT_ID)
+
+AZURE_CLIENT_CERTIFICATE_PATH = "AZURE_CLIENT_CERTIFICATE_PATH"
+AZURE_CLIENT_CERTIFICATE_PASSWORD = "AZURE_CLIENT_CERTIFICATE_PASSWORD"
+AZURE_CLIENT_SEND_CERTIFICATE_CHAIN = "AZURE_CLIENT_SEND_CERTIFICATE_CHAIN"
+CERT_VARS = (AZURE_CLIENT_ID, AZURE_CLIENT_CERTIFICATE_PATH, AZURE_TENANT_ID)
+
+
+class AzureADConfig:
+    def __init__(
+        self, client_id: str, client_secret: str, tenant_id: str, scopes: str
+    ) -> None:
+        self.scopes = scopes
+
+
+class AzureOpenAILLMClient(_BaseLiteLLMClient):
+    """A client for interfacing with Azure's OpenAI LLM deployments.
+
+    Parameters:
+        deployment (str): The deployment name.
+        model (Optional[str]): The name of the deployed model.
+        api_type: (Optional[str]): The api type. If not provided, it will be set via
+            environment variable.
+        api_base (Optional[str]): The base URL for the API endpoints. If not provided,
+            it will be set via environment variables.
+        api_version (Optional[str]): The version of the API to use. If not provided,
+            it will be set via environment variable.
+
+        kwargs (Optional[Dict[str, Any]]): Optional configuration parameters specific
+            to the model deployment.
+
+    Raises:
+        ProviderClientValidationError: If validation of the client setup fails.
+        DeprecationWarning: If deprecated environment variables are used for
+            configuration.
+    """
+
+    def __init__(
+        self,
+        deployment: str,
+        model: Optional[str] = None,
+        api_type: Optional[str] = None,
+        api_base: Optional[str] = None,
+        api_version: Optional[str] = None,
+        oauth: Optional[AzureOAuthConfig] = None,
+        **kwargs: Any,
+    ):
+        super().__init__()  # type: ignore
+        self._deployment = deployment
+        self._model = model
+        self._extra_parameters = kwargs or {}
+
+        # Set api_base with the following priority:
+        # parameter -> Azure Env Var -> (deprecated) OpenAI Env Var
+        self._api_base = (
+            api_base
+            or os.getenv(AZURE_API_BASE_ENV_VAR)
+            or os.getenv(OPENAI_API_BASE_ENV_VAR)
+        )
+
+        # Set api_version with the following priority:
+        # parameter -> Azure Env Var -> (deprecated) OpenAI Env Var
+        self._api_version = (
+            api_version
+            or os.getenv(AZURE_API_VERSION_ENV_VAR)
+            or os.getenv(OPENAI_API_VERSION_ENV_VAR)
+        )
+
+        # Not used by LiteLLM, here for backward compatibility
+        self._api_type = (
+            api_type
+            or os.getenv(AZURE_API_TYPE_ENV_VAR)
+            or os.getenv(OPENAI_API_TYPE_ENV_VAR)
+        )
+
+        os.unsetenv("OPENAI_API_KEY")
+        os.unsetenv("AZURE_API_KEY")
+
+        self._oauth = oauth
+
+        if self._oauth:
+            os.unsetenv(DEFAULT_AZURE_API_KEY_NAME)
+            os.unsetenv(AZURE_API_KEY_ENV_VAR)
+            os.unsetenv(OPENAI_API_KEY_ENV_VAR)
+        self._api_key_env_var = (
+            self._resolve_api_key_env_var() if not self._oauth else None
+        )
+
+        # Run helper function to check and raise deprecation warning if
+        # deprecated environment variables were used for initialization of the
+        # client settings
+        self._raise_evn_var_deprecation_warnings()
+
+        # validate the client settings
+        self.validate_client_setup()
+
+    def _raise_evn_var_deprecation_warnings(self) -> None:
+        """Helper function to check and raise deprecation warning if
+        deprecated environment variables were used for initialization of
+        some client settings.
+        """
+        deprecation_mapping = {
+            "API Base": {
+                "current_value": self.api_base,
+                "env_var": AZURE_API_BASE_ENV_VAR,
+                "deprecated_var": OPENAI_API_BASE_ENV_VAR,
+            },
+            "API Version": {
+                "current_value": self.api_version,
+                "env_var": AZURE_API_VERSION_ENV_VAR,
+                "deprecated_var": OPENAI_API_VERSION_ENV_VAR,
+            },
+        }
+
+        deprecation_warning_message = (
+            "Usage of {deprecated_env_var} environment "
+            "variable for setting the {setting} for Azure "
+            "OpenAI client is deprecated and will be removed "
+            "in 4.0.0. "
+        )
+        deprecation_warning_replacement_message = (
+            "Please use {env_var} environment variable."
+        )
+
+        for setting in deprecation_mapping.keys():
+            current_value = deprecation_mapping[setting]["current_value"]
+            env_var = deprecation_mapping[setting]["env_var"]
+            deprecated_var = deprecation_mapping[setting]["deprecated_var"]
+
+            # Value is set through the non-deprecated env var
+            if current_value == os.getenv(env_var):
+                continue
+
+            # Value is set through the deprecated env var
+            if current_value == os.getenv(deprecated_var):
+                message = deprecation_warning_message.format(
+                    setting=setting, deprecated_env_var=deprecated_var
+                )
+                if env_var is not None:
+                    message += deprecation_warning_replacement_message.format(
+                        env_var=env_var
+                    )
+                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 f"${DEFAULT_AZURE_API_KEY_NAME}"
+
+        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.
+
+        Args:
+            config (Dict[str, Any]): Configuration.
+
+        Raises:
+            ValueError:
+                If any of the required configuration keys are missing.
+                If `api_type` has a value different from `azure`.
+
+        Returns:
+            AzureOpenAILLMClient: Initialized client.
+        """
+        try:
+            azure_openai_config = AzureOpenAIClientConfig.from_dict(config)
+        except ValueError as e:
+            message = "Cannot instantiate a client from the passed configuration."
+            structlogger.error(
+                "azure_openai_llm_client.from_config.error",
+                message=message,
+                config=config,
+                original_error=e,
+            )
+            raise
+
+        return cls(
+            deployment=azure_openai_config.deployment,
+            model=azure_openai_config.model,
+            api_type=azure_openai_config.api_type,
+            api_base=azure_openai_config.api_base,
+            api_version=azure_openai_config.api_version,
+            oauth=azure_openai_config.oauth,
+            **azure_openai_config.extra_parameters,
+        )
+
+    @property
+    def config(self) -> dict:
+        """Returns the configuration for that the llm client
+        in dictionary form.
+        """
+        config = AzureOpenAIClientConfig(
+            deployment=self._deployment,
+            model=self._model,
+            api_base=self._api_base,
+            api_version=self._api_version,
+            api_type=self._api_type,
+            oauth=self._oauth,
+            extra_parameters=self._extra_parameters,
+        )
+        return config.to_dict()
+
+    @property
+    def deployment(self) -> str:
+        return self._deployment
+
+    @property
+    def model(self) -> Optional[str]:
+        """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."""
+        return self._api_base
+
+    @property
+    def api_version(self) -> Optional[str]:
+        """Returns the API version for the Azure OpenAI llm client."""
+        return self._api_version
+
+    @property
+    def api_type(self) -> Optional[str]:
+        return self._api_type
+
+    @property
+    def _litellm_model_name(self) -> str:
+        """Returns the value of LiteLLM's model parameter to be used in
+        completion/acompletion in LiteLLM format:
+
+        <provider>/<model or deployment name>
+        """
+        regex_pattern = rf"^{AZURE_OPENAI_PROVIDER}/"
+        if not re.match(regex_pattern, self._deployment):
+            return f"{AZURE_OPENAI_PROVIDER}/{self._deployment}"
+        return self._deployment
+
+    @property
+    def _litellm_extra_parameters(self) -> Dict[str, Any]:
+        return self._extra_parameters
+
+    @property
+    def _completion_fn_args(self) -> Dict[str, Any]:
+        """Returns the completion arguments for invoking a call through
+        LiteLLM's completion functions.
+        """
+        # Set the API key env var to None if OAuth is used
+        auth_parameter = (
+            {
+                LITE_LLM_AZURE_AD_TOKEN: self._oauth.get_bearer_token(),
+                LITE_LLM_AZURE_AD_TOKEN_PROVIDER: self._oauth.get_bearer_token_provider(),  # noqa: E501
+            }
+            if self._oauth
+            else {LITE_LLM_API_KEY_FIELD: self._api_key_env_var}
+        )
+
+        fn_args = super()._completion_fn_args
+        fn_args.update(
+            {
+                LITE_LLM_API_BASE_FIELD: self.api_base,
+                LITE_LLM_API_VERSION_FIELD: self.api_version,
+                **auth_parameter,
+            }
+        )
+        return fn_args
+
+    def validate_client_setup(self) -> None:
+        """Validates that all required configuration parameters are set."""
+
+        def generate_event_info_for_missing_setting(
+            setting: str,
+            setting_env_var: Optional[str] = None,
+            setting_config_key: Optional[str] = None,
+        ) -> str:
+            """Generate a part of the message with instructions on what to set
+            for the missing client setting.
+            """
+            info = "Set {setting} with {options}. "
+            options = ""
+            if setting_env_var is not None:
+                options += f"environment variable '{setting_env_var}'"
+            if setting_config_key is not None and setting_env_var is not None:
+                options += " or "
+            if setting_config_key is not None:
+                options += f"config key '{setting_config_key}'"
+
+            return info.format(setting=setting, options=options)
+
+        env_var_field = "env_var"
+        config_key_field = "config_key"
+        current_value_field = "current_value"
+        # All required settings for Azure OpenAI client
+        settings: Dict[str, Dict[str, Any]] = {
+            "API Base": {
+                current_value_field: self.api_base,
+                env_var_field: AZURE_API_BASE_ENV_VAR,
+                config_key_field: API_BASE_CONFIG_KEY,
+            },
+            "API Version": {
+                current_value_field: self.api_version,
+                env_var_field: AZURE_API_VERSION_ENV_VAR,
+                config_key_field: API_VERSION_CONFIG_KEY,
+            },
+            "Deployment Name": {
+                current_value_field: self.deployment,
+                env_var_field: None,
+                config_key_field: DEPLOYMENT_CONFIG_KEY,
+            },
+        }
+
+        missing_settings = [
+            setting_name
+            for setting_name, setting_info in settings.items()
+            if setting_info[current_value_field] is None
+        ]
+
+        if missing_settings:
+            event_info = f"Client settings not set: " f"{', '.join(missing_settings)}. "
+
+            for missing_setting in missing_settings:
+                if settings[missing_setting][current_value_field] is not None:
+                    continue
+                event_info += generate_event_info_for_missing_setting(
+                    missing_setting,
+                    settings[missing_setting][env_var_field],
+                    settings[missing_setting][config_key_field],
+                )
+
+            structlogger.error(
+                "azure_openai_llm_client.not_configured",
+                event_info=event_info,
+                missing_settings=missing_settings,
+            )
+            raise ProviderClientValidationError(event_info)