rasa-pro 3.11.4__py3-none-any.whl → 3.12.0__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,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/embedding_client.py

@@ -1,4 +1,4 @@
-from typing import Protocol, Dict, List, TYPE_CHECKING, runtime_checkable
+from typing import TYPE_CHECKING, Dict, List, Protocol, runtime_checkable
 
 if TYPE_CHECKING:
     from rasa.shared.providers.embedding.embedding_response import EmbeddingResponse

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,15 +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, cast
 
 import structlog
-from litellm import (
-    completion,
-    acompletion,
-    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,
@@ -19,7 +17,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 +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."""

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,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,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
             )
@@ -114,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.
 
@@ -124,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:
@@ -143,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.
 
@@ -153,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:
@@ -176,7 +198,8 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
         """Returns the completion arguments for invoking a call through
         LiteLLM's completion functions.
         """
+
         return {
             **self._litellm_extra_parameters,
-            "model": self.model_group_id,
+            LITE_LLM_MODEL_FIELD: self.model_group_id,
         }