rasa-pro 3.11.0rc1__py3-none-any.whl → 3.11.0rc3__py3-none-any.whl

rasa/shared/providers/embedding/_base_litellm_embedding_client.py

@@ -6,6 +6,7 @@ import litellm
 import structlog
 from litellm import aembedding, embedding, validate_environment
 
+from rasa.shared.constants import API_BASE_CONFIG_KEY, API_KEY
 from rasa.shared.exceptions import (
     ProviderClientAPIException,
     ProviderClientValidationError,
@@ -81,11 +82,14 @@ class _BaseLiteLLMEmbeddingClient:
             ProviderClientValidationError if validation fails.
         """
         self._validate_environment_variables()
-        self._validate_api_key_not_in_config()
 
     def _validate_environment_variables(self) -> None:
         """Validate that the required environment variables are set."""
-        validation_info = validate_environment(self._litellm_model_name)
+        validation_info = validate_environment(
+            self._litellm_model_name,
+            api_key=self._litellm_extra_parameters.get(API_KEY),
+            api_base=self._litellm_extra_parameters.get(API_BASE_CONFIG_KEY),
+        )
         if missing_environment_variables := validation_info.get(
             _VALIDATE_ENVIRONMENT_MISSING_KEYS_KEY
         ):
@@ -100,18 +104,6 @@ class _BaseLiteLLMEmbeddingClient:
             )
             raise ProviderClientValidationError(event_info)
 
-    def _validate_api_key_not_in_config(self) -> None:
-        if "api_key" in self._litellm_extra_parameters:
-            event_info = (
-                "API Key is set through `api_key` extra parameter."
-                "Set API keys through environment variables."
-            )
-            structlogger.error(
-                "base_litellm_client.validate_api_key_not_in_config",
-                event_info=event_info,
-            )
-            raise ProviderClientValidationError(event_info)
-
     def validate_documents(self, documents: List[str]) -> None:
         """Validates a list of documents to ensure they are suitable for embedding.
 

rasa/shared/providers/embedding/litellm_router_embedding_client.py

@@ -72,7 +72,7 @@ class LiteLLMRouterEmbeddingClient(
         return cls(
             model_group_id=client_config.model_group_id,
             model_configurations=client_config.litellm_model_list,
-            router_settings=client_config.router,
+            router_settings=client_config.litellm_router_settings,
             **client_config.extra_parameters,
         )
 

rasa/shared/providers/llm/_base_litellm_client.py

@@ -9,6 +9,7 @@ from litellm import (
     validate_environment,
 )
 
+from rasa.shared.constants import API_BASE_CONFIG_KEY, API_KEY
 from rasa.shared.exceptions import (
     ProviderClientAPIException,
     ProviderClientValidationError,
@@ -101,7 +102,11 @@ class _BaseLiteLLMClient:
 
     def _validate_environment_variables(self) -> None:
         """Validate that the required environment variables are set."""
-        validation_info = validate_environment(self._litellm_model_name)
+        validation_info = validate_environment(
+            self._litellm_model_name,
+            api_key=self._litellm_extra_parameters.get(API_KEY),
+            api_base=self._litellm_extra_parameters.get(API_BASE_CONFIG_KEY),
+        )
         if missing_environment_variables := validation_info.get(
             _VALIDATE_ENVIRONMENT_MISSING_KEYS_KEY
         ):
@@ -216,6 +221,32 @@ class _BaseLiteLLMClient:
         )
         return formatted_response
 
+    def _format_text_completion_response(self, response: Any) -> LLMResponse:
+        """Parses the LiteLLM text completion response to Rasa format."""
+        formatted_response = LLMResponse(
+            id=response.id,
+            created=response.created,
+            choices=[choice.text for choice in response.choices],
+            model=response.model,
+        )
+        if (usage := response.usage) is not None:
+            prompt_tokens = (
+                num_tokens
+                if isinstance(num_tokens := usage.prompt_tokens, (int, float))
+                else 0
+            )
+            completion_tokens = (
+                num_tokens
+                if isinstance(num_tokens := usage.completion_tokens, (int, float))
+                else 0
+            )
+            formatted_response.usage = LLMUsage(prompt_tokens, completion_tokens)
+        structlogger.debug(
+            "base_litellm_client.formatted_response",
+            formatted_response=formatted_response.to_dict(),
+        )
+        return formatted_response
+
     @staticmethod
     def _ensure_certificates() -> None:
         """Configures SSL certificates for LiteLLM. This method is invoked during

rasa/shared/providers/llm/litellm_router_llm_client.py

@@ -68,15 +68,61 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
         return cls(
             model_group_id=client_config.model_group_id,
             model_configurations=client_config.litellm_model_list,
-            router_settings=client_config.router,
+            router_settings=client_config.litellm_router_settings,
+            use_chat_completions_endpoint=client_config.use_chat_completions_endpoint,
             **client_config.extra_parameters,
         )
 
+    @suppress_logs(log_level=logging.WARNING)
+    def _text_completion(self, prompt: Union[List[str], str]) -> LLMResponse:
+        """
+        Synchronously generate completions for given prompt.
+
+        Args:
+            prompt: Prompt to generate the completion for.
+        Returns:
+            List of message completions.
+        Raises:
+            ProviderClientAPIException: If the API request fails.
+        """
+        try:
+            response = self.router_client.text_completion(
+                prompt=prompt, **self._completion_fn_args
+            )
+            return self._format_text_completion_response(response)
+        except Exception as e:
+            raise ProviderClientAPIException(e)
+
+    @suppress_logs(log_level=logging.WARNING)
+    async def _atext_completion(self, prompt: Union[List[str], str]) -> LLMResponse:
+        """
+        Asynchronously generate completions for given prompt.
+
+        Args:
+            prompt: Prompt to generate the completion for.
+        Returns:
+            List of message completions.
+        Raises:
+            ProviderClientAPIException: If the API request fails.
+        """
+        try:
+            response = await self.router_client.atext_completion(
+                prompt=prompt, **self._completion_fn_args
+            )
+            return self._format_text_completion_response(response)
+        except Exception as e:
+            raise ProviderClientAPIException(e)
+
     @suppress_logs(log_level=logging.WARNING)
     def completion(self, messages: Union[List[str], str]) -> LLMResponse:
         """
         Synchronously generate completions for given list of messages.
 
+        Method overrides the base class method to call the appropriate
+        completion method based on the configuration. If the chat completions
+        endpoint is enabled, the completion method is called. Otherwise, the
+        text_completion method is called.
+
         Args:
             messages: List of messages or a single message to generate the
                 completion for.
@@ -85,6 +131,8 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
         Raises:
             ProviderClientAPIException: If the API request fails.
         """
+        if not self._use_chat_completions_endpoint:
+            return self._text_completion(messages)
         try:
             formatted_messages = self._format_messages(messages)
             response = self.router_client.completion(
@@ -99,6 +147,11 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
         """
         Asynchronously generate completions for given list of messages.
 
+        Method overrides the base class method to call the appropriate
+        completion method based on the configuration. If the chat completions
+        endpoint is enabled, the completion method is called. Otherwise, the
+        text_completion method is called.
+
         Args:
             messages: List of messages or a single message to generate the
                 completion for.
@@ -107,6 +160,8 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
         Raises:
             ProviderClientAPIException: If the API request fails.
         """
+        if not self._use_chat_completions_endpoint:
+            return await self._atext_completion(messages)
         try:
             formatted_messages = self._format_messages(messages)
             response = await self.router_client.acompletion(

rasa/shared/providers/llm/self_hosted_llm_client.py

@@ -10,13 +10,14 @@ import structlog
 from rasa.shared.constants import (
     SELF_HOSTED_VLLM_PREFIX,
     SELF_HOSTED_VLLM_API_KEY_ENV_VAR,
+    API_KEY,
 )
 from rasa.shared.providers._configs.self_hosted_llm_client_config import (
     SelfHostedLLMClientConfig,
 )
 from rasa.shared.exceptions import ProviderClientAPIException
 from rasa.shared.providers.llm._base_litellm_client import _BaseLiteLLMClient
-from rasa.shared.providers.llm.llm_response import LLMResponse, LLMUsage
+from rasa.shared.providers.llm.llm_response import LLMResponse
 from rasa.shared.utils.io import suppress_logs
 
 structlogger = structlog.get_logger()
@@ -61,7 +62,8 @@ class SelfHostedLLMClient(_BaseLiteLLMClient):
         self._api_version = api_version
         self._use_chat_completions_endpoint = use_chat_completions_endpoint
         self._extra_parameters = kwargs or {}
-        self._apply_dummy_api_key_if_missing()
+        if self._extra_parameters.get(API_KEY) is None:
+            self._apply_dummy_api_key_if_missing()
 
     @classmethod
     def from_config(cls, config: Dict[str, Any]) -> "SelfHostedLLMClient":
@@ -259,32 +261,6 @@ class SelfHostedLLMClient(_BaseLiteLLMClient):
             return super().completion(messages)
         return self._text_completion(messages)
 
-    def _format_text_completion_response(self, response: Any) -> LLMResponse:
-        """Parses the LiteLLM text completion response to Rasa format."""
-        formatted_response = LLMResponse(
-            id=response.id,
-            created=response.created,
-            choices=[choice.text for choice in response.choices],
-            model=response.model,
-        )
-        if (usage := response.usage) is not None:
-            prompt_tokens = (
-                num_tokens
-                if isinstance(num_tokens := usage.prompt_tokens, (int, float))
-                else 0
-            )
-            completion_tokens = (
-                num_tokens
-                if isinstance(num_tokens := usage.completion_tokens, (int, float))
-                else 0
-            )
-            formatted_response.usage = LLMUsage(prompt_tokens, completion_tokens)
-        structlogger.debug(
-            "base_litellm_client.formatted_response",
-            formatted_response=formatted_response.to_dict(),
-        )
-        return formatted_response
-
     @staticmethod
     def _apply_dummy_api_key_if_missing() -> None:
         if not os.getenv(SELF_HOSTED_VLLM_API_KEY_ENV_VAR):

rasa/shared/providers/router/_base_litellm_router_client.py

@@ -1,4 +1,5 @@
 from typing import Any, Dict, List
+import os
 import structlog
 
 from litellm import Router
@@ -7,6 +8,12 @@ from rasa.shared.constants import (
     MODEL_LIST_KEY,
     MODEL_GROUP_ID_CONFIG_KEY,
     ROUTER_CONFIG_KEY,
+    SELF_HOSTED_VLLM_PREFIX,
+    SELF_HOSTED_VLLM_API_KEY_ENV_VAR,
+    LITELLM_PARAMS_KEY,
+    API_KEY,
+    MODEL_CONFIG_KEY,
+    USE_CHAT_COMPLETIONS_ENDPOINT_CONFIG_KEY,
 )
 from rasa.shared.exceptions import ProviderClientValidationError
 from rasa.shared.providers._configs.litellm_router_client_config import (
@@ -42,12 +49,15 @@ class _BaseLiteLLMRouterClient:
         model_group_id: str,
         model_configurations: List[Dict[str, Any]],
         router_settings: Dict[str, Any],
+        use_chat_completions_endpoint: bool = True,
         **kwargs: Any,
     ):
         self._model_group_id = model_group_id
         self._model_configurations = model_configurations
         self._router_settings = router_settings
+        self._use_chat_completions_endpoint = use_chat_completions_endpoint
         self._extra_parameters = kwargs or {}
+        self.additional_client_setup()
         try:
             resolved_model_configurations = (
                 self._resolve_env_vars_in_model_configurations()
@@ -67,6 +77,21 @@ class _BaseLiteLLMRouterClient:
             )
             raise ProviderClientValidationError(f"{event_info} Original error: {e}")
 
+    def additional_client_setup(self) -> None:
+        """Additional setup for the LiteLLM Router client."""
+        # If the model configuration is self-hosted VLLM, set a dummy API key if not
+        # provided. A bug in the LiteLLM library requires an API key to be set even if
+        # it is not required.
+        for model_configuration in self.model_configurations:
+            if (
+                f"{SELF_HOSTED_VLLM_PREFIX}/"
+                in model_configuration[LITELLM_PARAMS_KEY][MODEL_CONFIG_KEY]
+                and API_KEY not in model_configuration[LITELLM_PARAMS_KEY]
+                and not os.getenv(SELF_HOSTED_VLLM_API_KEY_ENV_VAR)
+            ):
+                os.environ[SELF_HOSTED_VLLM_API_KEY_ENV_VAR] = "dummy api key"
+                return
+
     @classmethod
     def from_config(cls, config: Dict[str, Any]) -> "_BaseLiteLLMRouterClient":
         """Instantiates a LiteLLM Router Embedding client from a configuration dict.
@@ -95,7 +120,8 @@ class _BaseLiteLLMRouterClient:
         return cls(
             model_group_id=client_config.model_group_id,
             model_configurations=client_config.litellm_model_list,
-            router_settings=client_config.router,
+            router_settings=client_config.litellm_router_settings,
+            use_chat_completions_endpoint=client_config.use_chat_completions_endpoint,
             **client_config.extra_parameters,
         )
 
@@ -119,6 +145,11 @@ class _BaseLiteLLMRouterClient:
         """Returns the instantiated LiteLLM Router client."""
         return self._router_client
 
+    @property
+    def use_chat_completions_endpoint(self) -> bool:
+        """Returns whether to use the chat completions endpoint."""
+        return self._use_chat_completions_endpoint
+
     @property
     def _litellm_extra_parameters(self) -> Dict[str, Any]:
         """
@@ -136,6 +167,9 @@ class _BaseLiteLLMRouterClient:
             MODEL_GROUP_ID_CONFIG_KEY: self.model_group_id,
             MODEL_LIST_KEY: self.model_configurations,
             ROUTER_CONFIG_KEY: self.router_settings,
+            USE_CHAT_COMPLETIONS_ENDPOINT_CONFIG_KEY: (
+                self.use_chat_completions_endpoint
+            ),
             **self._litellm_extra_parameters,
         }
 

rasa/shared/utils/common.py

@@ -193,7 +193,7 @@ def mark_as_experimental_feature(feature_name: Text) -> None:
 def mark_as_beta_feature(feature_name: Text) -> None:
     """Warns users that they are using a beta feature."""
     logger.warning(
-        f"🔬 Beta Feature: {feature_name} is in beta. It may have unexpected"
+        f"🔬 Beta Feature: {feature_name} is in beta. It may have unexpected "
         "behaviour and might be changed in the future."
     )
 

rasa/shared/utils/health_check/embeddings_health_check_mixin.py

@@ -0,0 +1,31 @@
+from typing import Optional, Dict, Any
+
+
+class EmbeddingsHealthCheckMixin:
+    """Mixin class that provides methods for performing embeddings health checks during
+    training and inference within components.
+
+    This mixin offers static methods that wrap the following health check functions:
+    - `perform_embeddings_health_check`
+    """
+
+    @staticmethod
+    def perform_embeddings_health_check(
+        custom_embeddings_config: Optional[Dict[str, Any]],
+        default_embeddings_config: Dict[str, Any],
+        log_source_method: str,
+        log_source_component: str,
+    ) -> None:
+        """Wraps the `perform_embeddings_health_check` function to enable
+        tracing and instrumentation.
+        """
+        from rasa.shared.utils.health_check.health_check import (
+            perform_embeddings_health_check,
+        )
+
+        perform_embeddings_health_check(
+            custom_embeddings_config,
+            default_embeddings_config,
+            log_source_method,
+            log_source_component,
+        )

rasa/shared/utils/health_check/health_check.py

@@ -0,0 +1,256 @@
+import os
+from typing import Optional, Dict, Any
+
+from rasa.shared.constants import (
+    LLM_API_HEALTH_CHECK_ENV_VAR,
+    MODELS_CONFIG_KEY,
+    LLM_API_HEALTH_CHECK_DEFAULT_VALUE,
+)
+from rasa.shared.exceptions import ProviderClientValidationError
+from rasa.shared.providers.embedding.embedding_client import EmbeddingClient
+from rasa.shared.providers.llm.llm_client import LLMClient
+from rasa.shared.utils.cli import print_error_and_exit
+from rasa.shared.utils.llm import llm_factory, structlogger, embedder_factory
+
+
+def try_instantiate_llm_client(
+    custom_llm_config: Optional[Dict],
+    default_llm_config: Optional[Dict],
+    log_source_function: str,
+    log_source_component: str,
+) -> LLMClient:
+    """Validate llm configuration."""
+    try:
+        return llm_factory(custom_llm_config, default_llm_config)
+    except (ProviderClientValidationError, ValueError) as e:
+        structlogger.error(
+            f"{log_source_function}.llm_instantiation_failed",
+            message="Unable to instantiate LLM client.",
+            error=e,
+        )
+        print_error_and_exit(
+            f"Unable to create the LLM client for component - {log_source_component}. "
+            f"Please make sure you specified the required environment variables "
+            f"and configuration keys. "
+            f"Error: {e}"
+        )
+
+
+def try_instantiate_embedder(
+    custom_embeddings_config: Optional[Dict],
+    default_embeddings_config: Optional[Dict],
+    log_source_function: str,
+    log_source_component: str,
+) -> EmbeddingClient:
+    """Validate embeddings configuration."""
+    try:
+        return embedder_factory(custom_embeddings_config, default_embeddings_config)
+    except (ProviderClientValidationError, ValueError) as e:
+        structlogger.error(
+            f"{log_source_function}.embedder_instantiation_failed",
+            message="Unable to instantiate Embedding client.",
+            error=e,
+        )
+        print_error_and_exit(
+            f"Unable to create the Embedding client for component - "
+            f"{log_source_component}. Please make sure you specified the required "
+            f"environment variables and configuration keys. Error: {e}"
+        )
+
+
+def perform_llm_health_check(
+    custom_config: Optional[Dict[str, Any]],
+    default_config: Dict[str, Any],
+    log_source_function: str,
+    log_source_component: str,
+) -> None:
+    """Try to instantiate the LLM Client to validate the provided config.
+    If the LLM_API_HEALTH_CHECK environment variable is true, perform a test call
+    to the LLM API. If config contains multiple models, perform a test call for each
+    model in the model group.
+
+    This method supports both single model configurations and model group configurations
+    (configs that have the `models` key).
+    """
+    # Instantiate the LLM client or Router LLM client to validate the provided config.
+    llm_client = try_instantiate_llm_client(
+        custom_config, default_config, log_source_function, log_source_component
+    )
+
+    if is_api_health_check_enabled():
+        if (
+            custom_config
+            and MODELS_CONFIG_KEY in custom_config
+            and len(custom_config[MODELS_CONFIG_KEY]) > 1
+        ):
+            # If the config uses a router, instantiate the LLM client for each model
+            # in the model group. This is required to perform a test api call for each
+            # model in the group.
+            # Note: The Router LLM client is not used here as we need to perform a test
+            # api call and not load balance the requests.
+            for model_config in custom_config[MODELS_CONFIG_KEY]:
+                llm_client = try_instantiate_llm_client(
+                    model_config,
+                    default_config,
+                    log_source_function,
+                    log_source_component,
+                )
+                send_test_llm_api_request(
+                    llm_client, log_source_function, log_source_component
+                )
+        else:
+            # Make a test api call to perform a health check for the LLM client.
+            # LLM config from config file and model group config from endpoint config
+            # without router are handled here.
+            send_test_llm_api_request(
+                llm_client,
+                log_source_function,
+                log_source_component,
+            )
+    else:
+        structlogger.warning(
+            f"{log_source_function}.perform_llm_health_check.disabled",
+            event_info=(
+                f"The {LLM_API_HEALTH_CHECK_ENV_VAR} environment variable is set "
+                f"to false, which will disable LLM health check. "
+                f"It is recommended to set this variable to true in production "
+                f"environments."
+            ),
+        )
+        return None
+
+
+def perform_embeddings_health_check(
+    custom_config: Optional[Dict[str, Any]],
+    default_config: Dict[str, Any],
+    log_source_function: str,
+    log_source_component: str,
+) -> None:
+    """Try to instantiate the Embedder to validate the provided config.
+    If the LLM_API_HEALTH_CHECK environment variable is true, perform a test call
+    to the Embeddings API. If config contains multiple models, perform a test call for
+    each model in the model group.
+
+    This method supports both single model configurations and model group configurations
+    (configs that have the `models` key).
+    """
+    # Instantiate the Embedder client or the Embedder Router client to validate the
+    # provided config. Deprecation warnings and errors are logged here.
+    embedder = try_instantiate_embedder(
+        custom_config, default_config, log_source_function, log_source_component
+    )
+
+    if is_api_health_check_enabled():
+        if (
+            custom_config
+            and MODELS_CONFIG_KEY in custom_config
+            and len(custom_config[MODELS_CONFIG_KEY]) > 1
+        ):
+            # If the config uses a router, instantiate the Embedder client for each
+            # model in the model group. This is required to perform a test api call
+            # for every model in the group.
+            # Note: The Router Embedding client is not used here as we need to perform
+            # a test API call and not load balance the requests.
+            for model_config in custom_config[MODELS_CONFIG_KEY]:
+                embedder = try_instantiate_embedder(
+                    model_config,
+                    default_config,
+                    log_source_function,
+                    log_source_component,
+                )
+                send_test_embeddings_api_request(
+                    embedder, log_source_function, log_source_component
+                )
+        else:
+            # Make a test api call to perform a health check for the Embedding client.
+            # Embeddings config from config file and model group config from endpoint
+            # config without router are handled here.
+            send_test_embeddings_api_request(
+                embedder, log_source_function, log_source_component
+            )
+    else:
+        structlogger.warning(
+            f"{log_source_function}" f".perform_embeddings_health_check.disabled",
+            event_info=(
+                f"The {LLM_API_HEALTH_CHECK_ENV_VAR} environment variable is set "
+                f"to false, which will disable embeddings API health check. "
+                f"It is recommended to set this variable to true in production "
+                f"environments."
+            ),
+        )
+        return None
+
+
+def send_test_llm_api_request(
+    llm_client: LLMClient, log_source_function: str, log_source_component: str
+) -> None:
+    """Sends a test request to the LLM API to perform a health check.
+
+    Raises:
+        Exception: If the API call fails.
+    """
+    structlogger.info(
+        f"{log_source_function}.send_test_llm_api_request",
+        event_info=(
+            f"Sending a test LLM API request for the component - "
+            f"{log_source_component}."
+        ),
+        config=llm_client.config,
+    )
+    try:
+        llm_client.completion("hello")
+    except Exception as e:
+        structlogger.error(
+            f"{log_source_function}.send_test_llm_api_request_failed",
+            event_info="Test call to the LLM API failed.",
+            error=e,
+        )
+        print_error_and_exit(
+            f"Test call to the LLM API failed for component - {log_source_component}. "
+            f"Error: {e}"
+        )
+
+
+def send_test_embeddings_api_request(
+    embedder: EmbeddingClient, log_source_function: str, log_source_component: str
+) -> None:
+    """Sends a test request to the Embeddings API to perform a health check.
+
+    Raises:
+        Exception: If the API call fails.
+    """
+    structlogger.info(
+        f"{log_source_function}.send_test_embeddings_api_request",
+        event_info=(
+            f"Sending a test Embeddings API request for the component - "
+            f"{log_source_component}."
+        ),
+        config=embedder.config,
+    )
+    try:
+        embedder.embed(["hello"])
+    except Exception as e:
+        structlogger.error(
+            f"{log_source_function}.send_test_llm_api_request_failed",
+            event_info="Test call to the Embeddings API failed.",
+            error=e,
+        )
+        print_error_and_exit(
+            f"Test call to the Embeddings API failed for component - "
+            f"{log_source_component}. Error: {e}"
+        )
+
+
+def is_api_health_check_enabled() -> bool:
+    """Determines whether the API health check is enabled based on an environment
+    variable.
+
+    Returns:
+        bool: True if the API health check is enabled, False otherwise.
+    """
+    return (
+        os.getenv(
+            LLM_API_HEALTH_CHECK_ENV_VAR, LLM_API_HEALTH_CHECK_DEFAULT_VALUE
+        ).lower()
+        == "true"
+    )

rasa/shared/utils/health_check/llm_health_check_mixin.py

@@ -0,0 +1,31 @@
+from typing import Optional, Dict, Any
+
+
+class LLMHealthCheckMixin:
+    """Mixin class that provides methods for performing llm health checks during
+    training and inference within components.
+
+    This mixin offers static methods that wrap the following health check functions:
+    - `perform_llm_health_check`
+    """
+
+    @staticmethod
+    def perform_llm_health_check(
+        custom_llm_config: Optional[Dict[str, Any]],
+        default_llm_config: Dict[str, Any],
+        log_source_method: str,
+        log_source_component: str,
+    ) -> None:
+        """Wraps the `perform_llm_health_check` function to enable
+        tracing and instrumentation.
+        """
+        from rasa.shared.utils.health_check.health_check import (
+            perform_llm_health_check,
+        )
+
+        perform_llm_health_check(
+            custom_llm_config,
+            default_llm_config,
+            log_source_method,
+            log_source_component,
+        )