rasa-pro 3.11.3a1.dev7__py3-none-any.whl → 3.12.0.dev2__py3-none-any.whl

rasa/shared/providers/_configs/azure_openai_client_config.py

@@ -1,36 +1,55 @@
+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 rasa.shared.constants import (
-    MODEL_CONFIG_KEY,
-    MODEL_NAME_CONFIG_KEY,
-    OPENAI_API_BASE_CONFIG_KEY,
     API_BASE_CONFIG_KEY,
-    OPENAI_API_TYPE_CONFIG_KEY,
+    API_KEY,
     API_TYPE_CONFIG_KEY,
-    OPENAI_API_VERSION_CONFIG_KEY,
     API_VERSION_CONFIG_KEY,
+    AZURE_API_TYPE,
+    AZURE_OPENAI_PROVIDER,
     DEPLOYMENT_CONFIG_KEY,
     DEPLOYMENT_NAME_CONFIG_KEY,
     ENGINE_CONFIG_KEY,
-    RASA_TYPE_CONFIG_KEY,
     LANGCHAIN_TYPE_CONFIG_KEY,
-    STREAM_CONFIG_KEY,
+    MODEL_CONFIG_KEY,
+    MODEL_NAME_CONFIG_KEY,
     N_REPHRASES_CONFIG_KEY,
+    OPENAI_API_BASE_CONFIG_KEY,
+    OPENAI_API_TYPE_CONFIG_KEY,
+    OPENAI_API_VERSION_CONFIG_KEY,
+    PROVIDER_CONFIG_KEY,
+    RASA_TYPE_CONFIG_KEY,
     REQUEST_TIMEOUT_CONFIG_KEY,
+    STREAM_CONFIG_KEY,
     TIMEOUT_CONFIG_KEY,
-    PROVIDER_CONFIG_KEY,
-    AZURE_OPENAI_PROVIDER,
-    AZURE_API_TYPE,
+)
+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 (
-    resolve_aliases,
     raise_deprecation_warnings,
-    validate_required_keys,
+    resolve_aliases,
     validate_forbidden_keys,
+    validate_required_keys,
 )
+from rasa.shared.utils.common import class_from_module_path
 
 structlogger = structlog.get_logger()
 
@@ -61,6 +80,76 @@ FORBIDDEN_KEYS = [
 ]
 
 
+@dataclass
+class OAuthConfigWrapper(OAuth):
+    """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.
+
+    """
+
+    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_config(cls, oauth_config: Dict[str, Any]) -> OAuth:
+        """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 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():
+            oauth = AzureEntraIDOAuthConfig.from_config(oauth_config)
+        else:
+            module = class_from_module_path(oauth_type)
+
+            if not issubclass(module, OAuth):
+                message = f"Module {oauth_type} does not implement {OAuth.__name__} interface."
+                structlogger.error(
+                    "azure_oauth_config.invalid_oauth_module",
+                    message=message,
+                )
+                raise ValueError(message)
+
+            oauth = module.from_config(oauth_config)
+
+        return cls(oauth=oauth, original_config=original_config)
+
+
 @dataclass
 class AzureOpenAIClientConfig:
     """Parses configuration for Azure OpenAI client, resolves aliases and
@@ -80,11 +169,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 +197,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 +220,26 @@ class AzureOpenAIClientConfig:
         # Validate that the forbidden keys are not present
         validate_forbidden_keys(config, FORBIDDEN_KEYS)
         # Init client config
+
+        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)
+
+        oauth = None
+        if has_oauth_key:
+            oauth = OAuthConfigWrapper.from_config(config.pop(OAUTH_KEY))
+
         this = AzureOpenAIClientConfig(
             # Required parameters
             deployment=config.pop(DEPLOYMENT_CONFIG_KEY),
@@ -142,6 +253,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,6 +267,9 @@ 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

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,24 +1,26 @@
+from __future__ import annotations
+
 from dataclasses import asdict, dataclass, field
 from typing import Any, Dict
 
 import structlog
 
+import rasa.shared.utils.cli
 from rasa.shared.constants import (
     MODEL_CONFIG_KEY,
     MODEL_NAME_CONFIG_KEY,
-    STREAM_CONFIG_KEY,
     N_REPHRASES_CONFIG_KEY,
     PROVIDER_CONFIG_KEY,
-    TIMEOUT_CONFIG_KEY,
     REQUEST_TIMEOUT_CONFIG_KEY,
+    STREAM_CONFIG_KEY,
+    TIMEOUT_CONFIG_KEY,
 )
 from rasa.shared.providers._configs.utils import (
-    validate_required_keys,
-    validate_forbidden_keys,
-    resolve_aliases,
     raise_deprecation_warnings,
+    resolve_aliases,
+    validate_forbidden_keys,
+    validate_required_keys,
 )
-import rasa.shared.utils.cli
 
 structlogger = structlog.get_logger()
 
@@ -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,27 +1,29 @@
+from __future__ import annotations
+
 from dataclasses import asdict, dataclass, field
 from typing import Any, Dict, Optional
 
 import structlog
 
 from rasa.shared.constants import (
-    MODEL_CONFIG_KEY,
-    MODEL_NAME_CONFIG_KEY,
-    RASA_TYPE_CONFIG_KEY,
-    LANGCHAIN_TYPE_CONFIG_KEY,
-    HUGGINGFACE_MULTIPROCESS_CONFIG_KEY,
     HUGGINGFACE_CACHE_FOLDER_CONFIG_KEY,
-    HUGGINGFACE_SHOW_PROGRESS_CONFIG_KEY,
-    HUGGINGFACE_MODEL_KWARGS_CONFIG_KEY,
     HUGGINGFACE_ENCODE_KWARGS_CONFIG_KEY,
     HUGGINGFACE_LOCAL_EMBEDDING_CACHING_FOLDER,
-    PROVIDER_CONFIG_KEY,
     HUGGINGFACE_LOCAL_EMBEDDING_PROVIDER,
-    TIMEOUT_CONFIG_KEY,
+    HUGGINGFACE_MODEL_KWARGS_CONFIG_KEY,
+    HUGGINGFACE_MULTIPROCESS_CONFIG_KEY,
+    HUGGINGFACE_SHOW_PROGRESS_CONFIG_KEY,
+    LANGCHAIN_TYPE_CONFIG_KEY,
+    MODEL_CONFIG_KEY,
+    MODEL_NAME_CONFIG_KEY,
+    PROVIDER_CONFIG_KEY,
+    RASA_TYPE_CONFIG_KEY,
     REQUEST_TIMEOUT_CONFIG_KEY,
+    TIMEOUT_CONFIG_KEY,
 )
 from rasa.shared.providers._configs.utils import (
-    resolve_aliases,
     raise_deprecation_warnings,
+    resolve_aliases,
     validate_required_keys,
 )
 from rasa.shared.utils.io import raise_deprecation_warning
@@ -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,29 +1,31 @@
+from __future__ import annotations
+
 import copy
 from dataclasses import dataclass, field
 from typing import Any, Dict, List
 
 import structlog
+
 from rasa.shared.constants import (
-    ROUTER_CONFIG_KEY,
-    MODELS_CONFIG_KEY,
-    MODEL_GROUP_ID_CONFIG_KEY,
-    MODEL_NAME_CONFIG_KEY,
-    LITELLM_PARAMS_KEY,
-    PROVIDER_CONFIG_KEY,
-    DEPLOYMENT_CONFIG_KEY,
     API_TYPE_CONFIG_KEY,
+    DEPLOYMENT_CONFIG_KEY,
+    LITELLM_PARAMS_KEY,
     MODEL_CONFIG_KEY,
+    MODEL_GROUP_ID_CONFIG_KEY,
     MODEL_LIST_KEY,
+    MODEL_NAME_CONFIG_KEY,
+    MODELS_CONFIG_KEY,
+    PROVIDER_CONFIG_KEY,
+    ROUTER_CONFIG_KEY,
     USE_CHAT_COMPLETIONS_ENDPOINT_CONFIG_KEY,
 )
 from rasa.shared.providers._configs.model_group_config import (
-    ModelGroupConfig,
     ModelConfig,
+    ModelGroupConfig,
 )
 from rasa.shared.providers.mappings import get_prefix_from_provider
 from rasa.shared.utils.llm import DEPLOYMENT_CENTRIC_PROVIDERS
 
-
 structlogger = structlog.get_logger()
 
 _LITELLM_UNSUPPORTED_KEYS = [
@@ -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,19 +1,25 @@
+from __future__ import annotations
+
 from dataclasses import asdict, dataclass, field
 from typing import List, Optional
 
 import structlog
+
 from rasa.shared.constants import (
     API_BASE_CONFIG_KEY,
     API_KEY,
     API_TYPE_CONFIG_KEY,
     API_VERSION_CONFIG_KEY,
     DEPLOYMENT_CONFIG_KEY,
-    PROVIDER_CONFIG_KEY,
+    EXTRA_PARAMETERS_KEY,
     MODEL_CONFIG_KEY,
     MODEL_GROUP_ID_CONFIG_KEY,
-    MODELS_CONFIG_KEY,
     MODEL_GROUPS_CONFIG_KEY,
-    EXTRA_PARAMETERS_KEY,
+    MODELS_CONFIG_KEY,
+    PROVIDER_CONFIG_KEY,
+)
+from rasa.shared.providers._configs._lite_llm_config.lite_llm_config_adapter import (
+    to_lite_llm_config,
 )
 from rasa.shared.providers.mappings import get_client_config_class_from_provider
 
@@ -40,7 +46,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.
@@ -130,7 +136,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, 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_config(cls: OAuthType, config: Dict[str, Any]) -> OAuthType:
+        """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,32 +1,34 @@
+from __future__ import annotations
+
 from dataclasses import asdict, dataclass, field
 from typing import Any, Dict, Optional
 
 import structlog
 
 from rasa.shared.constants import (
+    API_BASE_CONFIG_KEY,
+    API_TYPE_CONFIG_KEY,
+    API_VERSION_CONFIG_KEY,
+    LANGCHAIN_TYPE_CONFIG_KEY,
     MODEL_CONFIG_KEY,
     MODEL_NAME_CONFIG_KEY,
+    N_REPHRASES_CONFIG_KEY,
     OPENAI_API_BASE_CONFIG_KEY,
-    API_BASE_CONFIG_KEY,
+    OPENAI_API_TYPE,
     OPENAI_API_TYPE_CONFIG_KEY,
-    API_TYPE_CONFIG_KEY,
     OPENAI_API_VERSION_CONFIG_KEY,
-    API_VERSION_CONFIG_KEY,
+    OPENAI_PROVIDER,
+    PROVIDER_CONFIG_KEY,
     RASA_TYPE_CONFIG_KEY,
-    LANGCHAIN_TYPE_CONFIG_KEY,
-    STREAM_CONFIG_KEY,
-    N_REPHRASES_CONFIG_KEY,
     REQUEST_TIMEOUT_CONFIG_KEY,
+    STREAM_CONFIG_KEY,
     TIMEOUT_CONFIG_KEY,
-    PROVIDER_CONFIG_KEY,
-    OPENAI_API_TYPE,
-    OPENAI_PROVIDER,
 )
 from rasa.shared.providers._configs.utils import (
-    resolve_aliases,
-    validate_required_keys,
     raise_deprecation_warnings,
+    resolve_aliases,
     validate_forbidden_keys,
+    validate_required_keys,
 )
 
 structlogger = structlog.get_logger()
@@ -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,13 +1,15 @@
+from __future__ import annotations
+
 from dataclasses import asdict, dataclass, field
 from typing import Optional
 
 import structlog
 
 from rasa.shared.constants import (
+    API_BASE_CONFIG_KEY,
     MODEL_CONFIG_KEY,
-    RASA_PROVIDER,
     PROVIDER_CONFIG_KEY,
-    API_BASE_CONFIG_KEY,
+    RASA_PROVIDER,
 )
 from rasa.shared.providers._configs.utils import (
     validate_required_keys,
@@ -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,29 +1,30 @@
+from __future__ import annotations
+
 from dataclasses import asdict, dataclass, field
 from typing import Any, Dict, Optional
 
 import structlog
 
 from rasa.shared.constants import (
+    API_BASE_CONFIG_KEY,
+    API_TYPE_CONFIG_KEY,
+    API_VERSION_CONFIG_KEY,
+    LANGCHAIN_TYPE_CONFIG_KEY,
     MODEL_CONFIG_KEY,
     MODEL_NAME_CONFIG_KEY,
+    N_REPHRASES_CONFIG_KEY,
     OPENAI_API_BASE_CONFIG_KEY,
-    API_BASE_CONFIG_KEY,
     OPENAI_API_TYPE_CONFIG_KEY,
-    API_TYPE_CONFIG_KEY,
     OPENAI_API_VERSION_CONFIG_KEY,
-    API_VERSION_CONFIG_KEY,
+    OPENAI_PROVIDER,
+    PROVIDER_CONFIG_KEY,
     RASA_TYPE_CONFIG_KEY,
-    LANGCHAIN_TYPE_CONFIG_KEY,
-    STREAM_CONFIG_KEY,
-    N_REPHRASES_CONFIG_KEY,
     REQUEST_TIMEOUT_CONFIG_KEY,
-    TIMEOUT_CONFIG_KEY,
-    PROVIDER_CONFIG_KEY,
-    OPENAI_PROVIDER,
     SELF_HOSTED_PROVIDER,
+    STREAM_CONFIG_KEY,
+    TIMEOUT_CONFIG_KEY,
     USE_CHAT_COMPLETIONS_ENDPOINT_CONFIG_KEY,
 )
-
 from rasa.shared.providers._configs.utils import (
     raise_deprecation_warnings,
     resolve_aliases,
@@ -114,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"