rasa-pro 3.12.0.dev13__py3-none-any.whl → 3.12.0rc1__py3-none-any.whl

rasa/shared/providers/_configs/azure_openai_client_config.py

@@ -1,10 +1,20 @@
+from __future__ import annotations
+
+from copy import deepcopy
 from dataclasses import asdict, dataclass, field
-from typing import Any, Dict, Optional
+from typing import (
+    Any,
+    Dict,
+    Optional,
+    Set,
+)
 
 import structlog
+from pydantic import BaseModel
 
 from rasa.shared.constants import (
     API_BASE_CONFIG_KEY,
+    API_KEY,
     API_TYPE_CONFIG_KEY,
     API_VERSION_CONFIG_KEY,
     AZURE_API_TYPE,
@@ -25,12 +35,22 @@ from rasa.shared.constants import (
     STREAM_CONFIG_KEY,
     TIMEOUT_CONFIG_KEY,
 )
+from rasa.shared.providers._configs.azure_entra_id_config import (
+    AzureEntraIDOAuthConfig,
+    AzureEntraIDOAuthType,
+)
+from rasa.shared.providers._configs.oauth_config import (
+    OAUTH_KEY,
+    OAUTH_TYPE_FIELD,
+    OAuth,
+)
 from rasa.shared.providers._configs.utils import (
     raise_deprecation_warnings,
     resolve_aliases,
     validate_forbidden_keys,
     validate_required_keys,
 )
+from rasa.shared.utils.common import class_from_module_path
 
 structlogger = structlog.get_logger()
 
@@ -61,6 +81,86 @@ FORBIDDEN_KEYS = [
 ]
 
 
+class OAuthConfigWrapper(OAuth, BaseModel):
+    """Wrapper for OAuth configuration.
+
+    It's main purpose is to provide to_dict method which is used to serialize
+    the oauth configuration to the original format.
+
+    """
+
+    # Pydantic configuration to allow arbitrary user defined types
+    class Config:
+        arbitrary_types_allowed = True
+
+    oauth: OAuth
+    original_config: Dict[str, Any]
+
+    def get_bearer_token(self) -> str:
+        """Returns a bearer token."""
+        return self.oauth.get_bearer_token()
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Converts the OAuth configuration to the original format."""
+        return self.original_config
+
+    @staticmethod
+    def _valid_type_values() -> Set[str]:
+        """Returns the valid built-in values for the `type` field in the `oauth`."""
+        return AzureEntraIDOAuthType.valid_string_values()
+
+    @classmethod
+    def from_dict(cls, oauth_config: Dict[str, Any]) -> OAuthConfigWrapper:
+        """Initializes a dataclass from the passed config.
+
+        Args:
+            oauth_config: (dict) The config from which to initialize.
+
+        Returns:
+            AzureOAuthConfig
+        """
+        original_config = deepcopy(oauth_config)
+
+        oauth_type: Optional[str] = oauth_config.get(OAUTH_TYPE_FIELD, None)
+
+        if oauth_type is None:
+            message = (
+                "Oauth configuration must contain "
+                f"'{OAUTH_TYPE_FIELD}' field and it must be set to one of the "
+                f"following values: {OAuthConfigWrapper._valid_type_values()}, "
+                f"or to the path of module which is "
+                f"implementing {OAuth.__name__} protocol."
+            )
+            structlogger.error(
+                "azure_oauth_config.missing_oauth_type",
+                message=message,
+            )
+            raise ValueError(message)
+
+        if oauth_type in AzureEntraIDOAuthType.valid_string_values():
+            return cls(
+                oauth=AzureEntraIDOAuthConfig.from_dict(oauth_config),
+                original_config=original_config,
+            )
+
+        module = class_from_module_path(oauth_type)
+
+        if not issubclass(module, OAuth):
+            message = (
+                f"Module {oauth_type} does not implement "
+                f"{OAuth.__name__} interface."
+            )
+            structlogger.error(
+                "azure_oauth_config.invalid_oauth_module",
+                message=message,
+            )
+            raise ValueError(message)
+
+        return cls(
+            oauth=module.from_dict(oauth_config), original_config=original_config
+        )
+
+
 @dataclass
 class AzureOpenAIClientConfig:
     """Parses configuration for Azure OpenAI client, resolves aliases and
@@ -80,11 +180,13 @@ class AzureOpenAIClientConfig:
     # API Type is not used by LiteLLM backend, but we define
     # it here for backward compatibility.
     api_type: Optional[str] = AZURE_API_TYPE
-
     # Provider is not used by LiteLLM backend, but we define it here since it's
     # used as switch between different clients.
     provider: str = AZURE_OPENAI_PROVIDER
 
+    # OAuth related parameters
+    oauth: Optional[OAuthConfigWrapper] = None
+
     extra_parameters: dict = field(default_factory=dict)
 
     def __post_init__(self) -> None:
@@ -106,7 +208,7 @@ class AzureOpenAIClientConfig:
             raise ValueError(message)
 
     @classmethod
-    def from_dict(cls, config: dict) -> "AzureOpenAIClientConfig":
+    def from_dict(cls, config: dict) -> AzureOpenAIClientConfig:
         """Initializes a dataclass from the passed config.
 
         Args:
@@ -129,6 +231,16 @@ class AzureOpenAIClientConfig:
         # Validate that the forbidden keys are not present
         validate_forbidden_keys(config, FORBIDDEN_KEYS)
         # Init client config
+
+        cls._validate_authentication_configuration(config)
+
+        has_oauth_key = config.get(OAUTH_KEY, None) is not None
+        oauth = (
+            OAuthConfigWrapper.from_dict(config.pop(OAUTH_KEY))
+            if has_oauth_key
+            else None
+        )
+
         this = AzureOpenAIClientConfig(
             # Required parameters
             deployment=config.pop(DEPLOYMENT_CONFIG_KEY),
@@ -142,6 +254,8 @@ class AzureOpenAIClientConfig:
             # in clients.
             api_base=config.pop(API_BASE_CONFIG_KEY, None),
             api_version=config.pop(API_VERSION_CONFIG_KEY, None),
+            # OAuth related parameters, set only if auth_type is set to 'entra_id'
+            oauth=oauth,
             # The rest of parameters (e.g. model parameters) are considered
             # as extra parameters (this also includes timeout).
             extra_parameters=config,
@@ -154,12 +268,33 @@ class AzureOpenAIClientConfig:
         # Extra parameters should also be on the top level
         d.pop("extra_parameters", None)
         d.update(self.extra_parameters)
+
+        d.pop("oauth", None)
+        d.update({"oauth": self.oauth.to_dict()} if self.oauth else {})
         return d
 
     @staticmethod
     def resolve_config_aliases(config: Dict[str, Any]) -> Dict[str, Any]:
         return resolve_aliases(config, DEPRECATED_ALIASES_TO_STANDARD_KEY_MAPPING)
 
+    @staticmethod
+    def _validate_authentication_configuration(config: Dict[str, Any]) -> None:
+        """Validates the authentication configuration."""
+        has_api_key = config.get(API_KEY, None) is not None
+        has_oauth_key = config.get(OAUTH_KEY, None) is not None
+
+        if has_api_key and has_oauth_key:
+            message = (
+                "Azure OpenAI client configuration cannot contain "
+                f"both '{API_KEY}' and '{OAUTH_KEY}' fields. Please provide either "
+                f"'{API_KEY}' or '{OAUTH_KEY}' fields."
+            )
+            structlogger.error(
+                "azure_openai_client_config.multiple_auth_types_specified",
+                message=message,
+            )
+            raise ValueError(message)
+
 
 def is_azure_openai_config(config: dict) -> bool:
     """Check whether the configuration is meant to configure

rasa/shared/providers/_configs/client_config.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Protocol, runtime_checkable
 
 
@@ -9,7 +11,7 @@ class ClientConfig(Protocol):
     """
 
     @classmethod
-    def from_dict(cls, config: dict) -> "ClientConfig":
+    def from_dict(cls, config: dict) -> ClientConfig:
         """
         Initializes the client config with the given configuration.
 

rasa/shared/providers/_configs/default_litellm_client_config.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from dataclasses import asdict, dataclass, field
 from typing import Any, Dict
 
@@ -69,7 +71,7 @@ class DefaultLiteLLMClientConfig:
             raise ValueError(message)
 
     @classmethod
-    def from_dict(cls, config: dict) -> "DefaultLiteLLMClientConfig":
+    def from_dict(cls, config: dict) -> DefaultLiteLLMClientConfig:
         """
         Initializes a dataclass from the passed config.
 

rasa/shared/providers/_configs/huggingface_local_embedding_client_config.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from dataclasses import asdict, dataclass, field
 from typing import Any, Dict, Optional
 
@@ -90,7 +92,7 @@ class HuggingFaceLocalEmbeddingClientConfig:
             raise ValueError(message)
 
     @classmethod
-    def from_dict(cls, config: dict) -> "HuggingFaceLocalEmbeddingClientConfig":
+    def from_dict(cls, config: dict) -> HuggingFaceLocalEmbeddingClientConfig:
         """
         Initializes a dataclass from the passed config.
 

rasa/shared/providers/_configs/litellm_router_client_config.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import copy
 from dataclasses import dataclass, field
 from typing import Any, Dict, List
@@ -120,7 +122,7 @@ class LiteLLMRouterClientConfig:
             raise ValueError(message)
 
     @classmethod
-    def from_dict(cls, config: dict) -> "LiteLLMRouterClientConfig":
+    def from_dict(cls, config: dict) -> LiteLLMRouterClientConfig:
         """Initializes a dataclass from the passed config.
 
         Args:

rasa/shared/providers/_configs/model_group_config.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from dataclasses import asdict, dataclass, field
 from typing import List, Optional
 
@@ -41,7 +43,7 @@ class ModelConfig:
     api_type: Optional[str] = None
 
     @classmethod
-    def from_dict(cls, config: dict) -> "ModelConfig":
+    def from_dict(cls, config: dict) -> ModelConfig:
         """Initializes a dataclass from the passed config. The provider config param is
         used to determine the client config class to use. The client config class takes
         care of resolving config aliases and throwing deprecation warnings.
@@ -131,7 +133,7 @@ class ModelGroupConfig:
             raise ValueError(message)
 
     @classmethod
-    def from_dict(cls, config: dict) -> "ModelGroupConfig":
+    def from_dict(cls, config: dict) -> ModelGroupConfig:
         """Initializes a dataclass from the passed config.
 
         Args:

rasa/shared/providers/_configs/oauth_config.py

@@ -0,0 +1,33 @@
+import abc
+from typing import Any, Dict, Type, TypeVar
+
+OAUTH_TYPE_FIELD = "type"
+OAUTH_KEY = "oauth"
+
+OAuthType = TypeVar("OAuthType", bound="OAuth")
+
+
+class OAuth(abc.ABC):
+    """Interface for OAuth configuration."""
+
+    @classmethod
+    @abc.abstractmethod
+    def from_dict(
+        cls: Type[OAuthType], config: Dict[str, Any]
+    ) -> OAuthType:  # ignore[type]
+        """Initializes a dataclass from the passed config.
+
+        Args:
+            config: (dict) The config from which to initialize.
+
+        Returns:
+            OAuth
+        """
+
+    @abc.abstractmethod
+    def get_bearer_token(self) -> str:
+        """Returns a bearer token.
+
+        Bear token is used to authenticate requests to the Azure
+        Oopen AI instance's API protected by the Gateway.
+        """

rasa/shared/providers/_configs/openai_client_config.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from dataclasses import asdict, dataclass, field
 from typing import Any, Dict, Optional
 
@@ -111,7 +113,7 @@ class OpenAIClientConfig:
             raise ValueError(message)
 
     @classmethod
-    def from_dict(cls, config: dict) -> "OpenAIClientConfig":
+    def from_dict(cls, config: dict) -> OpenAIClientConfig:
         """
         Initializes a dataclass from the passed config.
 

rasa/shared/providers/_configs/rasa_llm_client_config.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from dataclasses import asdict, dataclass, field
 from typing import Optional
 
@@ -37,7 +39,7 @@ class RasaLLMClientConfig:
     extra_parameters: dict = field(default_factory=dict)
 
     @classmethod
-    def from_dict(cls, config: dict) -> "RasaLLMClientConfig":
+    def from_dict(cls, config: dict) -> RasaLLMClientConfig:
         """
         Initializes a dataclass from the passed config.
 

rasa/shared/providers/_configs/self_hosted_llm_client_config.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from dataclasses import asdict, dataclass, field
 from typing import Any, Dict, Optional
 
@@ -113,7 +115,7 @@ class SelfHostedLLMClientConfig:
             raise ValueError(message)
 
     @classmethod
-    def from_dict(cls, config: dict) -> "SelfHostedLLMClientConfig":
+    def from_dict(cls, config: dict) -> SelfHostedLLMClientConfig:
         """
         Initializes a dataclass from the passed config.
 

rasa/shared/providers/constants.py

@@ -0,0 +1,6 @@
+DEFAULT_AZURE_API_KEY_NAME = "AZURE_API_KEY"
+LITE_LLM_API_BASE_FIELD = "api_base"
+LITE_LLM_API_KEY_FIELD = "api_key"
+LITE_LLM_API_VERSION_FIELD = "api_version"
+LITE_LLM_MODEL_FIELD = "model"
+LITE_LLM_AZURE_AD_TOKEN = "azure_ad_token"

rasa/shared/providers/embedding/azure_openai_embedding_client.py

@@ -19,8 +19,14 @@ from rasa.shared.constants import (
 )
 from rasa.shared.exceptions import ProviderClientValidationError
 from rasa.shared.providers._configs.azure_openai_client_config import (
+    AzureEntraIDOAuthConfig,
     AzureOpenAIClientConfig,
 )
+from rasa.shared.providers.constants import (
+    DEFAULT_AZURE_API_KEY_NAME,
+    LITE_LLM_API_KEY_FIELD,
+    LITE_LLM_AZURE_AD_TOKEN,
+)
 from rasa.shared.providers.embedding._base_litellm_embedding_client import (
     _BaseLiteLLMEmbeddingClient,
 )
@@ -41,6 +47,8 @@ class AzureOpenAIEmbeddingClient(_BaseLiteLLMEmbeddingClient):
             If not provided, it will be set via environment variable.
         api_version (Optional[str]): The version of the API to use.
             If not provided, it will be set via environment variable.
+        oauth (Optional[AzureEntraIDOAuthConfig]): Optional OAuth configuration.
+            If provided, the client will use OAuth for authentication.
         kwargs (Optional[Dict[str, Any]]): Optional configuration parameters specific
             to the embedding model deployment.
 
@@ -57,6 +65,7 @@ class AzureOpenAIEmbeddingClient(_BaseLiteLLMEmbeddingClient):
         api_base: Optional[str] = None,
         api_type: Optional[str] = None,
         api_version: Optional[str] = None,
+        oauth: Optional[AzureEntraIDOAuthConfig] = None,
         **kwargs: Any,
     ):
         super().__init__()  # type: ignore
@@ -84,7 +93,11 @@ class AzureOpenAIEmbeddingClient(_BaseLiteLLMEmbeddingClient):
         # Litellm does not support use of OPENAI_API_KEY, so we need to map it
         # because of backward compatibility. However, we're first looking at
         # AZURE_API_KEY.
-        self._api_key_env_var = self._resolve_api_key_env_var()
+
+        self._oauth = oauth
+        self._api_key_env_var = (
+            self._resolve_api_key_env_var() if not self._oauth else None
+        )
 
         self.validate_client_setup()
 
@@ -100,7 +113,7 @@ class AzureOpenAIEmbeddingClient(_BaseLiteLLMEmbeddingClient):
             return self._extra_parameters[API_KEY]
 
         if os.getenv(AZURE_API_KEY_ENV_VAR) is not None:
-            return "${AZURE_API_KEY}"
+            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,
@@ -163,6 +176,7 @@ class AzureOpenAIEmbeddingClient(_BaseLiteLLMEmbeddingClient):
             api_base=azure_openai_config.api_base,
             api_type=azure_openai_config.api_type,
             api_version=azure_openai_config.api_version,
+            oauth=azure_openai_config.oauth,
             **azure_openai_config.extra_parameters,
         )
 
@@ -177,6 +191,7 @@ class AzureOpenAIEmbeddingClient(_BaseLiteLLMEmbeddingClient):
             api_base=self.api_base,
             api_type=self.api_type,
             api_version=self.api_version,
+            oauth=self._oauth,
             extra_parameters=self._extra_parameters,
         )
         return config.to_dict()
@@ -219,13 +234,23 @@ class AzureOpenAIEmbeddingClient(_BaseLiteLLMEmbeddingClient):
 
     @property
     def _embedding_fn_args(self) -> dict:
+        auth_parameter: Dict[str, str] = {}
+
+        if self._oauth:
+            auth_parameter = {
+                **auth_parameter,
+                LITE_LLM_AZURE_AD_TOKEN: self._oauth.get_bearer_token(),
+            }
+        elif self._api_key_env_var:
+            auth_parameter = {LITE_LLM_API_KEY_FIELD: self._api_key_env_var}
+
         return {
             **self._litellm_extra_parameters,
             "model": self._litellm_model_name,
             "api_base": self.api_base,
             "api_type": self.api_type,
             "api_version": self.api_version,
-            "api_key": self._api_key_env_var,
+            **auth_parameter,
         }
 
     @property

rasa/shared/providers/embedding/litellm_router_embedding_client.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from typing import Any, Dict, List
 
@@ -46,7 +48,7 @@ class LiteLLMRouterEmbeddingClient(
         )
 
     @classmethod
-    def from_config(cls, config: Dict[str, Any]) -> "LiteLLMRouterEmbeddingClient":
+    def from_config(cls, config: Dict[str, Any]) -> LiteLLMRouterEmbeddingClient:
         """Instantiates a LiteLLM Router Embedding client from a configuration dict.
 
         Args:

rasa/shared/providers/llm/_base_litellm_client.py

@@ -1,15 +1,13 @@
+from __future__ import annotations
+
 import logging
 from abc import abstractmethod
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List, Union, cast
 
 import structlog
-from litellm import (
-    acompletion,
-    completion,
-    validate_environment,
-)
+from litellm import acompletion, completion, validate_environment
 
-from rasa.shared.constants import API_BASE_CONFIG_KEY, API_KEY
+from rasa.shared.constants import API_BASE_CONFIG_KEY, API_KEY, ROLE_USER
 from rasa.shared.exceptions import (
     ProviderClientAPIException,
     ProviderClientValidationError,
@@ -50,7 +48,7 @@ class _BaseLiteLLMClient:
 
     @classmethod
     @abstractmethod
-    def from_config(cls, config: Dict[str, Any]) -> "_BaseLiteLLMClient":
+    def from_config(cls, config: Dict[str, Any]) -> _BaseLiteLLMClient:
         pass
 
     @property
@@ -122,12 +120,18 @@ class _BaseLiteLLMClient:
             raise ProviderClientValidationError(event_info)
 
     @suppress_logs(log_level=logging.WARNING)
-    def completion(self, messages: Union[List[str], str]) -> LLMResponse:
+    def completion(self, messages: Union[List[dict], 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.
+            messages: The message can be,
+                - a list of preformatted messages. Each message should be a dictionary
+                    with the following keys:
+                    - content: The message content.
+                    - role: The role of the message (e.g. user or system).
+                - a list of messages. Each message is a string and will be formatted
+                    as a user message.
+                - a single message as a string which will be formatted as user message.
 
         Returns:
             List of message completions.
@@ -136,7 +140,7 @@ class _BaseLiteLLMClient:
             ProviderClientAPIException: If the API request fails.
         """
         try:
-            formatted_messages = self._format_messages(messages)
+            formatted_messages = self._get_formatted_messages(messages)
             arguments = resolve_environment_variables(self._completion_fn_args)
             response = completion(messages=formatted_messages, **arguments)
             return self._format_response(response)
@@ -144,12 +148,20 @@ class _BaseLiteLLMClient:
             raise ProviderClientAPIException(e)
 
     @suppress_logs(log_level=logging.WARNING)
-    async def acompletion(self, messages: Union[List[str], str]) -> LLMResponse:
+    async def acompletion(
+        self, messages: Union[List[dict], 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.
+            messages: The message can be,
+                - a list of preformatted messages. Each message should be a dictionary
+                    with the following keys:
+                    - content: The message content.
+                    - role: The role of the message (e.g. user or system).
+                - a list of messages. Each message is a string and will be formatted
+                    as a user message.
+                - a single message as a string which will be formatted as user message.
 
         Returns:
             List of message completions.
@@ -158,7 +170,7 @@ class _BaseLiteLLMClient:
             ProviderClientAPIException: If the API request fails.
         """
         try:
-            formatted_messages = self._format_messages(messages)
+            formatted_messages = self._get_formatted_messages(messages)
             arguments = resolve_environment_variables(self._completion_fn_args)
             response = await acompletion(messages=formatted_messages, **arguments)
             return self._format_response(response)
@@ -181,11 +193,24 @@ class _BaseLiteLLMClient:
                 )
             raise ProviderClientAPIException(e, message)
 
+    def _get_formatted_messages(
+        self, messages: Union[List[dict], List[str], str]
+    ) -> List[Dict[str, str]]:
+        """Returns a list of formatted messages."""
+        if (
+            isinstance(messages, list)
+            and len(messages) > 0
+            and isinstance(messages[0], dict)
+        ):
+            # Check if the messages are already formatted. If so, return them as is.
+            return cast(List[Dict[str, str]], messages)
+        return self._format_messages(messages)
+
     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]
+        return [{"content": message, "role": ROLE_USER} for message in messages]
 
     def _format_response(self, response: Any) -> LLMResponse:
         """Parses the LiteLLM response to Rasa format."""