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

rasa/shared/providers/embedding/azure_openai_embedding_client.py

@@ -4,23 +4,29 @@ from typing import Any, Dict, List, 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,
     OPENAI_API_BASE_ENV_VAR,
     OPENAI_API_KEY_ENV_VAR,
     OPENAI_API_TYPE_ENV_VAR,
     OPENAI_API_VERSION_ENV_VAR,
-    API_BASE_CONFIG_KEY,
-    API_KEY,
-    API_VERSION_CONFIG_KEY,
-    AZURE_OPENAI_PROVIDER,
 )
 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,21 @@ class AzureOpenAIEmbeddingClient(_BaseLiteLLMEmbeddingClient):
 
     @property
     def _embedding_fn_args(self) -> dict:
+        auth_parameter = (
+            {
+                LITE_LLM_AZURE_AD_TOKEN: self._oauth.get_bearer_token(),
+            }
+            if self._oauth
+            else {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,5 +1,8 @@
-from typing import Any, Dict, List
+from __future__ import annotations
+
 import logging
+from typing import Any, Dict, List
+
 import structlog
 
 from rasa.shared.exceptions import ProviderClientAPIException
@@ -45,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,11 +1,13 @@
+from __future__ import annotations
+
 import logging
 from abc import abstractmethod
-from typing import Dict, List, Any, Union
+from typing import Any, Dict, List, Union
 
 import structlog
 from litellm import (
-    completion,
     acompletion,
+    completion,
     validate_environment,
 )
 
@@ -19,7 +21,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, resolve_environment_variables
+from rasa.shared.utils.io import resolve_environment_variables, suppress_logs
 
 structlogger = structlog.get_logger()
 
@@ -50,7 +52,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

rasa/shared/providers/llm/azure_openai_llm_client.py

@@ -1,33 +1,60 @@
+from __future__ import annotations
+
 import os
 import re
-from typing import Dict, Any, Optional
+from typing import Any, Dict, Optional
 
 import structlog
 
 from rasa.shared.constants import (
-    OPENAI_API_BASE_ENV_VAR,
-    OPENAI_API_VERSION_ENV_VAR,
-    AZURE_API_BASE_ENV_VAR,
-    AZURE_API_VERSION_ENV_VAR,
     API_BASE_CONFIG_KEY,
+    API_KEY,
     API_VERSION_CONFIG_KEY,
-    DEPLOYMENT_CONFIG_KEY,
+    AZURE_API_BASE_ENV_VAR,
     AZURE_API_KEY_ENV_VAR,
-    OPENAI_API_TYPE_ENV_VAR,
-    OPENAI_API_KEY_ENV_VAR,
     AZURE_API_TYPE_ENV_VAR,
+    AZURE_API_VERSION_ENV_VAR,
     AZURE_OPENAI_PROVIDER,
-    API_KEY,
+    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 (
+    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,21 @@ 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 = (
+            {
+                LITE_LLM_AZURE_AD_TOKEN: self._oauth.get_bearer_token(),
+            }
+            if self._oauth
+            else {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 +365,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,4 +1,6 @@
-from typing import Dict, Any
+from __future__ import annotations
+
+from typing import Any, Dict
 
 from rasa.shared.constants import (
     AWS_BEDROCK_PROVIDER,
@@ -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,11 +1,15 @@
-from typing import Any, Dict, List, Union
+from __future__ import annotations
+
 import logging
+from typing import Any, Dict, List, Union
+
 import structlog
 
 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 (
@@ -41,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:
@@ -86,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
             )
@@ -106,6 +114,10 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
             ProviderClientAPIException: If the API request fails.
         """
         try:
+            structlogger.info(
+                "litellm_router_llm_client.atext_completion",
+                _completion_fn_args=self._completion_fn_args,
+            )
             response = await self.router_client.atext_completion(
                 prompt=prompt, **self._completion_fn_args
             )
@@ -135,6 +147,10 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
             return self._text_completion(messages)
         try:
             formatted_messages = self._format_messages(messages)
+            structlogger.info(
+                "litellm_router_llm_client.completion",
+                _completion_fn_args=self._completion_fn_args,
+            )
             response = self.router_client.completion(
                 messages=formatted_messages, **self._completion_fn_args
             )
@@ -164,6 +180,10 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
             return await self._atext_completion(messages)
         try:
             formatted_messages = self._format_messages(messages)
+            structlogger.info(
+                "litellm_router_llm_client.acompletion",
+                _completion_fn_args=self._completion_fn_args,
+            )
             response = await self.router_client.acompletion(
                 messages=formatted_messages, **self._completion_fn_args
             )
@@ -178,5 +198,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,4 +1,6 @@
-from typing import Protocol, Dict, List, runtime_checkable, Union
+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.
 

rasa/shared/providers/llm/llm_response.py

@@ -1,8 +1,5 @@
 from dataclasses import dataclass, field, asdict
-from typing import Dict, List, Optional, Text, Any, Union
-import structlog
-
-structlogger = structlog.get_logger()
+from typing import Dict, List, Optional
 
 
 @dataclass
@@ -19,18 +16,6 @@ class LLMUsage:
     def __post_init__(self) -> None:
         self.total_tokens = self.prompt_tokens + self.completion_tokens
 
-    @classmethod
-    def from_dict(cls, data: Dict[Text, Any]) -> "LLMUsage":
-        """
-        Creates an LLMUsage object from a dictionary.
-        If any keys are missing, they will default to zero
-        or whatever default you prefer.
-        """
-        return cls(
-            prompt_tokens=data.get("prompt_tokens"),
-            completion_tokens=data.get("completion_tokens"),
-        )
-
     def to_dict(self) -> dict:
         """Converts the LLMUsage dataclass instance into a dictionary."""
         return asdict(self)
@@ -57,32 +42,6 @@ class LLMResponse:
     """Optional dictionary for storing additional information related to the
     completion that may not be covered by other fields."""
 
-    @classmethod
-    def from_dict(cls, data: Dict[Text, Any]) -> "LLMResponse":
-        """
-        Creates an LLMResponse from a dictionary.
-        """
-        usage_data = data.get("usage", {})
-        usage_obj = LLMUsage.from_dict(usage_data) if usage_data else None
-
-        return cls(
-            id=data["id"],
-            choices=data["choices"],
-            created=data["created"],
-            model=data.get("model"),
-            usage=usage_obj,
-            additional_info=data.get("additional_info"),
-        )
-
-    @classmethod
-    def ensure_llm_response(cls, response: Union[str, "LLMResponse"]) -> "LLMResponse":
-        if isinstance(response, LLMResponse):
-            return response
-
-        structlogger.warn("llm_response.deprecated_response_type", response=response)
-        data = {"id": None, "choices": [response], "created": None}
-        return LLMResponse.from_dict(data)
-
     def to_dict(self) -> dict:
         """Converts the LLMResponse dataclass instance into a dictionary."""
         result = asdict(self)

rasa/shared/providers/llm/openai_llm_client.py

@@ -1,16 +1,22 @@
+from __future__ import annotations
+
 import os
 import re
-from typing import Dict, Any, Optional
+from typing import Any, Dict, Optional
 
 import structlog
 
 from rasa.shared.constants import (
     OPENAI_API_BASE_ENV_VAR,
-    OPENAI_API_VERSION_ENV_VAR,
     OPENAI_API_TYPE_ENV_VAR,
+    OPENAI_API_VERSION_ENV_VAR,
     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,17 +1,22 @@
+from __future__ import annotations
+
 from typing import Any, Dict, Optional
 
 import structlog
 
 from rasa.shared.constants import (
-    RASA_PROVIDER,
     OPENAI_PROVIDER,
+    RASA_PROVIDER,
 )
 from rasa.shared.providers._configs.rasa_llm_client_config import (
     RasaLLMClientConfig,
 )
-from rasa.utils.licensing import retrieve_license_from_env
+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
 
 structlogger = structlog.get_logger()
 
@@ -88,12 +93,15 @@ class RasaLLMClient(_BaseLiteLLMClient):
         """
         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: