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

rasa/shared/providers/llm/azure_openai_llm_client.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 import re
 from typing import Any, Dict, Optional
@@ -21,13 +23,38 @@ 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_BASE_FIELD,
+    LITE_LLM_API_KEY_FIELD,
+    LITE_LLM_API_VERSION_FIELD,
+    LITE_LLM_AZURE_AD_TOKEN,
+)
 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.
@@ -41,6 +68,7 @@ class AzureOpenAILLMClient(_BaseLiteLLMClient):
             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.
 
@@ -57,6 +85,7 @@ class AzureOpenAILLMClient(_BaseLiteLLMClient):
         api_type: Optional[str] = None,
         api_base: Optional[str] = None,
         api_version: Optional[str] = None,
+        oauth: Optional[AzureEntraIDOAuthConfig] = None,
         **kwargs: Any,
     ):
         super().__init__()  # type: ignore
@@ -80,8 +109,6 @@ class AzureOpenAILLMClient(_BaseLiteLLMClient):
             or os.getenv(OPENAI_API_VERSION_ENV_VAR)
         )
 
-        self._api_key_env_var = self._resolve_api_key_env_var()
-
         # Not used by LiteLLM, here for backward compatibility
         self._api_type = (
             api_type
@@ -89,6 +116,19 @@ class AzureOpenAILLMClient(_BaseLiteLLMClient):
             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
@@ -157,7 +197,7 @@ class AzureOpenAILLMClient(_BaseLiteLLMClient):
             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,
@@ -188,7 +228,7 @@ class AzureOpenAILLMClient(_BaseLiteLLMClient):
         )
 
     @classmethod
-    def from_config(cls, config: Dict[str, Any]) -> "AzureOpenAILLMClient":
+    def from_config(cls, config: Dict[str, Any]) -> AzureOpenAILLMClient:
         """Initializes the client from given configuration.
 
         Args:
@@ -215,11 +255,12 @@ class AzureOpenAILLMClient(_BaseLiteLLMClient):
             raise
 
         return cls(
-            azure_openai_config.deployment,
-            azure_openai_config.model,
-            azure_openai_config.api_type,
-            azure_openai_config.api_base,
-            azure_openai_config.api_version,
+            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,
         )
 
@@ -234,6 +275,7 @@ class AzureOpenAILLMClient(_BaseLiteLLMClient):
             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()
@@ -282,12 +324,23 @@ class AzureOpenAILLMClient(_BaseLiteLLMClient):
         """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: 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}
+
         fn_args = super()._completion_fn_args
         fn_args.update(
             {
-                "api_base": self.api_base,
-                "api_version": self.api_version,
-                "api_key": self._api_key_env_var,
+                LITE_LLM_API_BASE_FIELD: self.api_base,
+                LITE_LLM_API_VERSION_FIELD: self.api_version,
+                **auth_parameter,
             }
         )
         return fn_args
@@ -314,41 +367,44 @@ class AzureOpenAILLMClient(_BaseLiteLLMClient):
 
             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": self.api_base,
-                "env_var": AZURE_API_BASE_ENV_VAR,
-                "config_key": API_BASE_CONFIG_KEY,
+                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": self.api_version,
-                "env_var": AZURE_API_VERSION_ENV_VAR,
-                "config_key": API_VERSION_CONFIG_KEY,
+                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": self.deployment,
-                "env_var": None,
-                "config_key": DEPLOYMENT_CONFIG_KEY,
+                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"] is None
+            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"] is not None:
+                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"],
-                    settings[missing_setting]["config_key"],
+                    settings[missing_setting][env_var_field],
+                    settings[missing_setting][config_key_field],
                 )
 
             structlogger.error(

rasa/shared/providers/llm/default_litellm_llm_client.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Any, Dict
 
 from rasa.shared.constants import (
@@ -35,7 +37,7 @@ class DefaultLiteLLMClient(_BaseLiteLLMClient):
         self.validate_client_setup()
 
     @classmethod
-    def from_config(cls, config: Dict[str, Any]) -> "DefaultLiteLLMClient":
+    def from_config(cls, config: Dict[str, Any]) -> DefaultLiteLLMClient:
         default_config = DefaultLiteLLMClientConfig.from_dict(config)
         return cls(
             model=default_config.model,

rasa/shared/providers/llm/litellm_router_llm_client.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from typing import Any, Dict, List, Union
 
@@ -7,6 +9,7 @@ from rasa.shared.exceptions import ProviderClientAPIException
 from rasa.shared.providers._configs.litellm_router_client_config import (
     LiteLLMRouterClientConfig,
 )
+from rasa.shared.providers.constants import LITE_LLM_MODEL_FIELD
 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 (
@@ -42,7 +45,7 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
         )
 
     @classmethod
-    def from_config(cls, config: Dict[str, Any]) -> "LiteLLMRouterLLMClient":
+    def from_config(cls, config: Dict[str, Any]) -> LiteLLMRouterLLMClient:
         """Instantiates a LiteLLM Router LLM client from a configuration dict.
 
         Args:
@@ -87,6 +90,10 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
             ProviderClientAPIException: If the API request fails.
         """
         try:
+            structlogger.info(
+                "litellm_router_llm_client.text_completion",
+                _completion_fn_args=self._completion_fn_args,
+            )
             response = self.router_client.text_completion(
                 prompt=prompt, **self._completion_fn_args
             )
@@ -115,7 +122,7 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
             raise ProviderClientAPIException(e)
 
     @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.
 
@@ -125,8 +132,14 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
         text_completion method is called.
 
         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.
         Raises:
@@ -144,7 +157,9 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _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.
 
@@ -154,8 +169,14 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
         text_completion method is called.
 
         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.
         Raises:
@@ -179,5 +200,5 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
         """
         return {
             **self._litellm_extra_parameters,
-            "model": self.model_group_id,
+            LITE_LLM_MODEL_FIELD: self.model_group_id,
         }

rasa/shared/providers/llm/llm_client.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Dict, List, Protocol, Union, runtime_checkable
 
 from rasa.shared.providers.llm.llm_response import LLMResponse
@@ -11,7 +13,7 @@ class LLMClient(Protocol):
     """
 
     @classmethod
-    def from_config(cls, config: dict) -> "LLMClient":
+    def from_config(cls, config: dict) -> LLMClient:
         """
         Initializes the llm client with the given configuration.
 
@@ -30,7 +32,7 @@ class LLMClient(Protocol):
         """
         ...
 
-    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.
 
@@ -38,14 +40,22 @@ class LLMClient(Protocol):
         strings) and return a list of completions (as strings).
 
         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:
             LLMResponse
         """
         ...
 
-    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.
 
@@ -53,8 +63,14 @@ class LLMClient(Protocol):
         strings) and return a list of completions (as strings).
 
         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:
             LLMResponse
         """

rasa/shared/providers/llm/openai_llm_client.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 import re
 from typing import Any, Dict, Optional
@@ -11,6 +13,10 @@ from rasa.shared.constants import (
     OPENAI_PROVIDER,
 )
 from rasa.shared.providers._configs.openai_client_config import OpenAIClientConfig
+from rasa.shared.providers.constants import (
+    LITE_LLM_API_KEY_FIELD,
+    LITE_LLM_API_VERSION_FIELD,
+)
 from rasa.shared.providers.llm._base_litellm_client import _BaseLiteLLMClient
 
 structlogger = structlog.get_logger()
@@ -57,7 +63,7 @@ class OpenAILLMClient(_BaseLiteLLMClient):
         self.validate_client_setup()
 
     @classmethod
-    def from_config(cls, config: Dict[str, Any]) -> "OpenAILLMClient":
+    def from_config(cls, config: Dict[str, Any]) -> OpenAILLMClient:
         """
         Initializes the client from given configuration.
 
@@ -148,8 +154,8 @@ class OpenAILLMClient(_BaseLiteLLMClient):
         fn_args = super()._completion_fn_args
         fn_args.update(
             {
-                "api_base": self.api_base,
-                "api_version": self.api_version,
+                LITE_LLM_API_KEY_FIELD: self.api_base,
+                LITE_LLM_API_VERSION_FIELD: self.api_version,
             }
         )
         return fn_args

rasa/shared/providers/llm/rasa_llm_client.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Any, Dict, Optional
 
 import structlog
@@ -9,6 +11,10 @@ from rasa.shared.constants import (
 from rasa.shared.providers._configs.rasa_llm_client_config import (
     RasaLLMClientConfig,
 )
+from rasa.shared.providers.constants import (
+    LITE_LLM_API_BASE_FIELD,
+    LITE_LLM_API_KEY_FIELD,
+)
 from rasa.shared.providers.llm._base_litellm_client import _BaseLiteLLMClient
 from rasa.utils.licensing import retrieve_license_from_env
 
@@ -82,12 +88,15 @@ class RasaLLMClient(_BaseLiteLLMClient):
         """Returns the completion arguments for invoking a call using completions."""
         fn_args = super()._completion_fn_args
         fn_args.update(
-            {"api_base": self.api_base, "api_key": retrieve_license_from_env()}
+            {
+                LITE_LLM_API_BASE_FIELD: self.api_base,
+                LITE_LLM_API_KEY_FIELD: retrieve_license_from_env(),
+            }
         )
         return fn_args
 
     @classmethod
-    def from_config(cls, config: Dict[str, Any]) -> "RasaLLMClient":
+    def from_config(cls, config: Dict[str, Any]) -> RasaLLMClient:
         try:
             client_config = RasaLLMClientConfig.from_dict(config)
         except ValueError as e:

rasa/shared/providers/llm/self_hosted_llm_client.py

@@ -1,12 +1,11 @@
+from __future__ import annotations
+
 import logging
 import os
 from typing import Any, Dict, List, Optional, Union
 
 import structlog
-from litellm import (
-    atext_completion,
-    text_completion,
-)
+from litellm import atext_completion, text_completion
 
 from rasa.shared.constants import (
     API_KEY,
@@ -17,6 +16,10 @@ from rasa.shared.exceptions import ProviderClientAPIException
 from rasa.shared.providers._configs.self_hosted_llm_client_config import (
     SelfHostedLLMClientConfig,
 )
+from rasa.shared.providers.constants import (
+    LITE_LLM_API_BASE_FIELD,
+    LITE_LLM_API_VERSION_FIELD,
+)
 from rasa.shared.providers.llm._base_litellm_client import _BaseLiteLLMClient
 from rasa.shared.providers.llm.llm_response import LLMResponse
 from rasa.shared.utils.io import suppress_logs
@@ -67,7 +70,7 @@ class SelfHostedLLMClient(_BaseLiteLLMClient):
             self._apply_dummy_api_key_if_missing()
 
     @classmethod
-    def from_config(cls, config: Dict[str, Any]) -> "SelfHostedLLMClient":
+    def from_config(cls, config: Dict[str, Any]) -> SelfHostedLLMClient:
         try:
             client_config = SelfHostedLLMClientConfig.from_dict(config)
         except ValueError as e:
@@ -184,8 +187,8 @@ class SelfHostedLLMClient(_BaseLiteLLMClient):
         fn_args = super()._completion_fn_args
         fn_args.update(
             {
-                "api_base": self.api_base,
-                "api_version": self.api_version,
+                LITE_LLM_API_BASE_FIELD: self.api_base,
+                LITE_LLM_API_VERSION_FIELD: self.api_version,
             }
         )
         return fn_args
@@ -214,7 +217,14 @@ class SelfHostedLLMClient(_BaseLiteLLMClient):
         Asynchronously generate completions for given prompt.
 
         Args:
-            prompt: Prompt 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.
         Raises:
@@ -226,7 +236,9 @@ class SelfHostedLLMClient(_BaseLiteLLMClient):
         except Exception as e:
             raise ProviderClientAPIException(e)
 
-    async def acompletion(self, messages: Union[List[str], str]) -> LLMResponse:
+    async def acompletion(
+        self, messages: Union[List[dict], List[str], str]
+    ) -> LLMResponse:
         """Asynchronous completion of the model with the given messages.
 
         Method overrides the base class method to call the appropriate
@@ -235,7 +247,14 @@ class SelfHostedLLMClient(_BaseLiteLLMClient):
         atext_completion method is called.
 
         Args:
-            messages: The messages to be used for completion.
+            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:
             The completion response.
@@ -244,7 +263,7 @@ class SelfHostedLLMClient(_BaseLiteLLMClient):
             return await super().acompletion(messages)
         return await self._atext_completion(messages)
 
-    def completion(self, messages: Union[List[str], str]) -> LLMResponse:
+    def completion(self, messages: Union[List[dict], List[str], str]) -> LLMResponse:
         """Completion of the model with the given messages.
 
         Method overrides the base class method to call the appropriate

rasa/shared/providers/router/_base_litellm_router_client.py

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

rasa/shared/providers/router/router_client.py

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

rasa/shared/utils/constants.py

@@ -2,3 +2,6 @@ DEFAULT_ENCODING = "utf-8"
 
 READ_YAML_FILE_CACHE_MAXSIZE_ENV_VAR = "READ_YAML_FILE_CACHE_MAXSIZE"
 DEFAULT_READ_YAML_FILE_CACHE_MAXSIZE = 256
+RASA_PRO_BETA_PREDICATES_IN_RESPONSE_CONDITIONS_ENV_VAR_NAME = (
+    "RASA_PRO_BETA_PREDICATES_IN_RESPONSE_CONDITIONS"
+)