rasa-pro 3.12.0.dev1__py3-none-any.whl

rasa/shared/providers/llm/rasa_llm_client.py

@@ -0,0 +1,120 @@
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+
+import structlog
+
+from rasa.shared.constants import (
+    OPENAI_PROVIDER,
+    RASA_PROVIDER,
+)
+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
+
+structlogger = structlog.get_logger()
+
+
+class RasaLLMClient(_BaseLiteLLMClient):
+    """A client for interfacing with a Rasa-Hosted LLM endpoint that uses
+
+    Parameters:
+        model (str): The model or deployment name.
+        api_base (str): The base URL of the API endpoint.
+        kwargs: Any: Additional configuration parameters that can include, but
+            are not limited to model parameters and lite-llm specific
+            parameters. These parameters will be passed to the
+            completion/acompletion calls. To see what it can include, visit:
+
+    Raises:
+        ProviderClientValidationError: If validation of the client setup fails.
+        ProviderClientAPIException: If the API request fails.
+    """
+
+    def __init__(
+        self,
+        model: str,
+        api_base: str,
+        **kwargs: Any,
+    ):
+        super().__init__()  # type: ignore
+        self._model = model
+        self._api_base = api_base
+        self._use_chat_completions_endpoint = True
+        self._extra_parameters = kwargs or {}
+
+    @property
+    def model(self) -> str:
+        return self._model
+
+    @property
+    def api_base(self) -> Optional[str]:
+        """
+        Returns the base API URL for the openai llm client.
+        """
+        return self._api_base
+
+    @property
+    def provider(self) -> str:
+        """
+        Returns the provider name for the self hosted llm client.
+
+        Returns:
+            String representing the provider name.
+        """
+        return RASA_PROVIDER
+
+    @property
+    def _litellm_model_name(self) -> str:
+        return f"{OPENAI_PROVIDER}/{self._model}"
+
+    @property
+    def _litellm_extra_parameters(self) -> Dict[str, Any]:
+        return self._extra_parameters
+
+    @property
+    def config(self) -> dict:
+        return RasaLLMClientConfig(
+            model=self._model,
+            api_base=self._api_base,
+            extra_parameters=self._extra_parameters,
+        ).to_dict()
+
+    @property
+    def _completion_fn_args(self) -> Dict[str, Any]:
+        """Returns the completion arguments for invoking a call through
+        LiteLLM's completion functions.
+        """
+        fn_args = super()._completion_fn_args
+        fn_args.update(
+            {
+                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:
+        try:
+            client_config = RasaLLMClientConfig.from_dict(config)
+        except ValueError as e:
+            message = "Cannot instantiate a client from the passed configuration."
+            structlogger.error(
+                "rasa_llm_client.from_config.error",
+                message=message,
+                config=config,
+                original_error=e,
+            )
+            raise
+        return cls(
+            model=client_config.model,
+            api_base=client_config.api_base,
+            **client_config.extra_parameters,
+        )

rasa/shared/providers/llm/self_hosted_llm_client.py

@@ -0,0 +1,276 @@
+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 rasa.shared.constants import (
+    API_KEY,
+    SELF_HOSTED_VLLM_API_KEY_ENV_VAR,
+    SELF_HOSTED_VLLM_PREFIX,
+)
+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
+
+structlogger = structlog.get_logger()
+
+
+class SelfHostedLLMClient(_BaseLiteLLMClient):
+    """A client for interfacing with Self Hosted LLM endpoints that uses
+
+    Parameters:
+        model (str): The model or deployment name.
+        provider (str): The provider of the model.
+        api_base (str): The base URL of the API endpoint.
+        api_type (Optional[str]): The type of the API endpoint.
+        api_version (Optional[str]): The version of the API endpoint.
+        use_chat_completions_endpoint (Optional[bool]): Whether to use the chat
+            completions endpoint for completions. Defaults to True.
+        kwargs: Any: Additional configuration parameters that can include, but
+            are not limited to model parameters and lite-llm specific
+            parameters. These parameters will be passed to the
+            completion/acompletion calls. To see what it can include, visit:
+
+    Raises:
+        ProviderClientValidationError: If validation of the client setup fails.
+        ProviderClientAPIException: If the API request fails.
+    """
+
+    def __init__(
+        self,
+        provider: str,
+        model: str,
+        api_base: str,
+        api_type: Optional[str] = None,
+        api_version: Optional[str] = None,
+        use_chat_completions_endpoint: Optional[bool] = True,
+        **kwargs: Any,
+    ):
+        super().__init__()  # type: ignore
+        self._provider = provider
+        self._model = model
+        self._api_base = api_base
+        self._api_type = api_type
+        self._api_version = api_version
+        self._use_chat_completions_endpoint = use_chat_completions_endpoint
+        self._extra_parameters = kwargs or {}
+        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:
+        try:
+            client_config = SelfHostedLLMClientConfig.from_dict(config)
+        except ValueError as e:
+            message = "Cannot instantiate a client from the passed configuration."
+            structlogger.error(
+                "self_hosted_llm_client.from_config.error",
+                message=message,
+                config=config,
+                original_error=e,
+            )
+            raise
+
+        return cls(
+            model=client_config.model,
+            provider=client_config.provider,
+            api_base=client_config.api_base,
+            api_type=client_config.api_type,
+            api_version=client_config.api_version,
+            use_chat_completions_endpoint=client_config.use_chat_completions_endpoint,
+            **client_config.extra_parameters,
+        )
+
+    @property
+    def provider(self) -> str:
+        """
+        Returns the provider name for the self hosted llm client.
+
+        Returns:
+            String representing the provider name.
+        """
+        return self._provider
+
+    @property
+    def model(self) -> str:
+        """
+        Returns the model name for the self hosted llm client.
+
+        Returns:
+            String representing the model name.
+        """
+        return self._model
+
+    @property
+    def api_base(self) -> str:
+        """
+        Returns the base URL for the API endpoint.
+
+        Returns:
+            String representing the base URL.
+        """
+        return self._api_base
+
+    @property
+    def api_type(self) -> Optional[str]:
+        """
+        Returns the type of the API endpoint. Currently only OpenAI is supported.
+
+        Returns:
+            String representing the API type.
+        """
+        return self._api_type
+
+    @property
+    def api_version(self) -> Optional[str]:
+        """
+        Returns the version of the API endpoint.
+
+        Returns:
+            String representing the API version.
+        """
+        return self._api_version
+
+    @property
+    def config(self) -> Dict:
+        """
+        Returns the configuration for the self hosted llm client.
+        Returns:
+            Dictionary containing the configuration.
+        """
+        config = SelfHostedLLMClientConfig(
+            model=self._model,
+            provider=self._provider,
+            api_base=self._api_base,
+            api_type=self._api_type,
+            api_version=self._api_version,
+            use_chat_completions_endpoint=self._use_chat_completions_endpoint,
+            extra_parameters=self._extra_parameters,
+        )
+        return config.to_dict()
+
+    @property
+    def _litellm_model_name(self) -> str:
+        """Returns the value of LiteLLM's model parameter to be used in
+        completion/acompletion in LiteLLM format:
+
+        <hosted_vllm>/<model or deployment name>
+        """
+        if self.model and f"{SELF_HOSTED_VLLM_PREFIX}/" not in self.model:
+            return f"{SELF_HOSTED_VLLM_PREFIX}/{self.model}"
+        return self.model
+
+    @property
+    def _litellm_extra_parameters(self) -> Dict[str, Any]:
+        """Returns optional configuration parameters specific
+        to the client provider and deployed model.
+        """
+        return self._extra_parameters
+
+    @property
+    def _completion_fn_args(self) -> Dict[str, Any]:
+        """Returns the completion arguments for invoking a call through
+        LiteLLM's completion functions.
+        """
+        fn_args = super()._completion_fn_args
+        fn_args.update(
+            {
+                LITE_LLM_API_BASE_FIELD: self.api_base,
+                LITE_LLM_API_VERSION_FIELD: self.api_version,
+            }
+        )
+        return fn_args
+
+    @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 = 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 atext_completion(prompt=prompt, **self._completion_fn_args)
+            return self._format_text_completion_response(response)
+        except Exception as e:
+            raise ProviderClientAPIException(e)
+
+    async def acompletion(self, messages: Union[List[str], str]) -> LLMResponse:
+        """Asynchronous completion of the model with the given 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 acompletion method is called. Otherwise, the
+        atext_completion method is called.
+
+        Args:
+            messages: The messages to be used for completion.
+
+        Returns:
+            The completion response.
+        """
+        if self._use_chat_completions_endpoint:
+            return await super().acompletion(messages)
+        return await self._atext_completion(messages)
+
+    def completion(self, messages: Union[List[str], str]) -> LLMResponse:
+        """Completion of the model with the given 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: The messages to be used for completion.
+
+        Returns:
+            The completion response.
+        """
+        if self._use_chat_completions_endpoint:
+            return super().completion(messages)
+        return self._text_completion(messages)
+
+    @staticmethod
+    def _apply_dummy_api_key_if_missing() -> None:
+        if not os.getenv(SELF_HOSTED_VLLM_API_KEY_ENV_VAR):
+            os.environ[SELF_HOSTED_VLLM_API_KEY_ENV_VAR] = (
+                "dummy_self_hosted_llm_api_key"
+            )

rasa/shared/providers/mappings.py

@@ -0,0 +1,94 @@
+from typing import Dict, Type, Optional
+
+from rasa.shared.constants import (
+    AZURE_OPENAI_PROVIDER,
+    HUGGINGFACE_LOCAL_EMBEDDING_PROVIDER,
+    OPENAI_PROVIDER,
+    SELF_HOSTED_PROVIDER,
+    RASA_PROVIDER,
+    SELF_HOSTED_VLLM_PREFIX,
+)
+from rasa.shared.providers.embedding.azure_openai_embedding_client import (
+    AzureOpenAIEmbeddingClient,
+)
+from rasa.shared.providers.embedding.default_litellm_embedding_client import (
+    DefaultLiteLLMEmbeddingClient,
+)
+from rasa.shared.providers.embedding.embedding_client import EmbeddingClient
+from rasa.shared.providers.embedding.huggingface_local_embedding_client import (
+    HuggingFaceLocalEmbeddingClient,
+)
+from rasa.shared.providers.embedding.openai_embedding_client import (
+    OpenAIEmbeddingClient,
+)
+from rasa.shared.providers.llm.azure_openai_llm_client import AzureOpenAILLMClient
+from rasa.shared.providers.llm.default_litellm_llm_client import DefaultLiteLLMClient
+from rasa.shared.providers.llm.llm_client import LLMClient
+from rasa.shared.providers.llm.openai_llm_client import OpenAILLMClient
+from rasa.shared.providers.llm.self_hosted_llm_client import SelfHostedLLMClient
+from rasa.shared.providers.llm.rasa_llm_client import RasaLLMClient
+from rasa.shared.providers._configs.azure_openai_client_config import (
+    AzureOpenAIClientConfig,
+)
+from rasa.shared.providers._configs.default_litellm_client_config import (
+    DefaultLiteLLMClientConfig,
+)
+from rasa.shared.providers._configs.huggingface_local_embedding_client_config import (
+    HuggingFaceLocalEmbeddingClientConfig,
+)
+from rasa.shared.providers._configs.openai_client_config import OpenAIClientConfig
+from rasa.shared.providers._configs.self_hosted_llm_client_config import (
+    SelfHostedLLMClientConfig,
+)
+from rasa.shared.providers._configs.rasa_llm_client_config import RasaLLMClientConfig
+
+from rasa.shared.providers._configs.client_config import ClientConfig
+
+_provider_to_llm_client_mapping: Dict[str, Type[LLMClient]] = {
+    OPENAI_PROVIDER: OpenAILLMClient,
+    AZURE_OPENAI_PROVIDER: AzureOpenAILLMClient,
+    SELF_HOSTED_PROVIDER: SelfHostedLLMClient,
+    RASA_PROVIDER: RasaLLMClient,
+}
+
+_provider_to_embedding_client_mapping: Dict[str, Type[EmbeddingClient]] = {
+    OPENAI_PROVIDER: OpenAIEmbeddingClient,
+    AZURE_OPENAI_PROVIDER: AzureOpenAIEmbeddingClient,
+    HUGGINGFACE_LOCAL_EMBEDDING_PROVIDER: HuggingFaceLocalEmbeddingClient,
+}
+
+_provider_to_client_config_class_mapping: Dict[str, Type] = {
+    OPENAI_PROVIDER: OpenAIClientConfig,
+    AZURE_OPENAI_PROVIDER: AzureOpenAIClientConfig,
+    HUGGINGFACE_LOCAL_EMBEDDING_PROVIDER: HuggingFaceLocalEmbeddingClientConfig,
+    SELF_HOSTED_PROVIDER: SelfHostedLLMClientConfig,
+    RASA_PROVIDER: RasaLLMClientConfig,
+}
+
+
+_provider_to_prefix_mapping: Dict[str, str] = {
+    # Specify the provider name as the key and its corresponding prefix as the value
+    # for providers where the prefix differs from the provider name.
+    SELF_HOSTED_PROVIDER: SELF_HOSTED_VLLM_PREFIX,
+    RASA_PROVIDER: OPENAI_PROVIDER,
+}
+
+
+def get_llm_client_from_provider(provider: Optional[str]) -> Type[LLMClient]:
+    return _provider_to_llm_client_mapping.get(provider, DefaultLiteLLMClient)
+
+
+def get_embedding_client_from_provider(provider: str) -> Type[EmbeddingClient]:
+    return _provider_to_embedding_client_mapping.get(
+        provider, DefaultLiteLLMEmbeddingClient
+    )
+
+
+def get_client_config_class_from_provider(provider: str) -> Type[ClientConfig]:
+    return _provider_to_client_config_class_mapping.get(
+        provider, DefaultLiteLLMClientConfig
+    )
+
+
+def get_prefix_from_provider(provider: str) -> str:
+    return _provider_to_prefix_mapping.get(provider, provider)

rasa/shared/providers/router/_base_litellm_router_client.py

@@ -0,0 +1,185 @@
+from __future__ import annotations
+
+import os
+from typing import Any, Dict, List
+
+import structlog
+from litellm import Router
+
+from rasa.shared.constants import (
+    API_KEY,
+    LITELLM_PARAMS_KEY,
+    MODEL_CONFIG_KEY,
+    MODEL_GROUP_ID_CONFIG_KEY,
+    MODEL_LIST_KEY,
+    ROUTER_CONFIG_KEY,
+    SELF_HOSTED_VLLM_API_KEY_ENV_VAR,
+    SELF_HOSTED_VLLM_PREFIX,
+    USE_CHAT_COMPLETIONS_ENDPOINT_CONFIG_KEY,
+)
+from rasa.shared.exceptions import ProviderClientValidationError
+from rasa.shared.providers._configs.litellm_router_client_config import (
+    LiteLLMRouterClientConfig,
+)
+from rasa.shared.utils.io import resolve_environment_variables
+
+structlogger = structlog.get_logger()
+
+
+class _BaseLiteLLMRouterClient:
+    """An abstract base class for LiteLLM Router clients.
+
+    This class defines the interface and common functionality for all the router clients
+    based on LiteLLM.
+
+    The class is made private to prevent it from being part of the public-facing
+    interface, as it serves as an internal base class for specific implementations of
+    router clients that are based on LiteLLM router implementation.
+
+    Parameters:
+        model_group_id (str): The model group ID.
+        model_configurations (List[Dict[str, Any]]): The list of model configurations.
+        router_settings (Dict[str, Any]): The router settings.
+        kwargs (Optional[Dict[str, Any]]): Additional configuration parameters.
+
+    Raises:
+        ProviderClientValidationError: If validation of the client setup fails.
+    """
+
+    def __init__(
+        self,
+        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()
+            )
+            self._router_client = Router(
+                model_list=resolved_model_configurations, **router_settings
+            )
+        except Exception as e:
+            event_info = "Cannot instantiate a router client."
+            structlogger.error(
+                "_base_litellm_router_client.init.error",
+                event_info=event_info,
+                model_group_id=model_group_id,
+                models=model_configurations,
+                router=router_settings,
+                original_error=e,
+            )
+            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.
+
+        Args:
+            config: (Dict[str, Any]) The configuration dictionary.
+
+        Returns:
+            LiteLLMRouterLLMClient: The instantiated LiteLLM Router LLM client.
+
+        Raises:
+            ValueError: If the configuration is invalid.
+        """
+        try:
+            client_config = LiteLLMRouterClientConfig.from_dict(config)
+        except (KeyError, ValueError) as e:
+            message = "Cannot instantiate a client from the passed configuration."
+            structlogger.error(
+                "litellm_router_llm_client.from_config.error",
+                message=message,
+                config=config,
+                original_error=e,
+            )
+            raise
+
+        return cls(
+            model_group_id=client_config.model_group_id,
+            model_configurations=client_config.litellm_model_list,
+            router_settings=client_config.litellm_router_settings,
+            use_chat_completions_endpoint=client_config.use_chat_completions_endpoint,
+            **client_config.extra_parameters,
+        )
+
+    @property
+    def model_group_id(self) -> str:
+        """Returns the model group ID for the LiteLLM Router client."""
+        return self._model_group_id
+
+    @property
+    def model_configurations(self) -> List[Dict[str, Any]]:
+        """Returns the model configurations for the LiteLLM Router client."""
+        return self._model_configurations
+
+    @property
+    def router_settings(self) -> Dict[str, Any]:
+        """Returns the router settings for the LiteLLM Router client."""
+        return self._router_settings
+
+    @property
+    def router_client(self) -> Router:
+        """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]:
+        """
+        Returns the extra parameters for the LiteLLM Router client.
+
+        Returns:
+            Dictionary containing the model parameters.
+        """
+        return self._extra_parameters
+
+    @property
+    def config(self) -> Dict:
+        """Returns the configuration for the LiteLLM Router client in LiteLLM format."""
+        return {
+            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,
+        }
+
+    def _resolve_env_vars_in_model_configurations(self) -> List:
+        model_configuration_with_resolved_keys = []
+        for model_configuration in self.model_configurations:
+            resolved_model_configuration = resolve_environment_variables(
+                model_configuration
+            )
+            model_configuration_with_resolved_keys.append(resolved_model_configuration)
+        return model_configuration_with_resolved_keys