rasa-pro 3.9.17__py3-none-any.whl → 3.10.3__py3-none-any.whl

rasa/shared/providers/embedding/openai_embedding_client.py

@@ -0,0 +1,172 @@
+import os
+from typing import Any, Dict, Optional
+
+import structlog
+
+from rasa.shared.constants import (
+    OPENAI_API_BASE_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.embedding._base_litellm_embedding_client import (
+    _BaseLiteLLMEmbeddingClient,
+)
+
+structlogger = structlog.get_logger()
+
+
+class OpenAIEmbeddingClient(_BaseLiteLLMEmbeddingClient):
+    """A client for interfacing with OpenAI Embeddings.
+
+    Parameters:
+        model (str): The OpenAI model name.
+        api_base (Optional[str]): Optional, the base URL for the API endpoints.
+            If not provided, it will be set via environment variables.
+        api_type (Optional[str]): Optional, the type of the API to use.
+            If not provided, it will be set via environment variable.
+        api_version (Optional[str]): Optional, 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 embedding model deployment.
+
+    Raises:
+        ProviderClientValidationError: If validation of the client setup fails.
+    """
+
+    def __init__(
+        self,
+        model: str,
+        api_base: Optional[str] = None,
+        api_type: Optional[str] = None,
+        api_version: Optional[str] = None,
+        **kwargs: Any,
+    ):
+        super().__init__()  # type: ignore
+        self._model = model
+        self._api_base = api_base or os.environ.get(OPENAI_API_BASE_ENV_VAR)
+        self._api_type = api_type or os.environ.get(OPENAI_API_TYPE_ENV_VAR)
+        self._api_version = api_version or os.environ.get(OPENAI_API_VERSION_ENV_VAR)
+        self._extra_parameters = kwargs or {}
+        self.validate_client_setup()
+
+    @classmethod
+    def from_config(cls, config: Dict[str, Any]) -> "OpenAIEmbeddingClient":
+        """
+        Initializes the client from given configuration.
+
+        Args:
+            config (Dict[str, Any]): Configuration.
+
+        Raises:
+            ValueError: Raised in cases of invalid configuration:
+                - If any of the required configuration keys are missing.
+                - If `api_type` has a value different from `azure`.
+
+        Returns:
+            OpenAIEmbeddingClient: Initialized client.
+        """
+        try:
+            openai_config = OpenAIClientConfig.from_dict(config)
+        except ValueError as e:
+            message = "Cannot instantiate a client from the passed configuration."
+            structlogger.error(
+                "openai_embedding_client.from_config.error",
+                message=message,
+                config=config,
+                original_error=e,
+            )
+            raise
+
+        return cls(
+            model=openai_config.model,
+            api_base=openai_config.api_base,
+            api_type=openai_config.api_type,
+            api_version=openai_config.api_version,
+            **openai_config.extra_parameters,
+        )
+
+    @property
+    def config(self) -> Dict:
+        """
+        Returns the configuration for the openai embedding client.
+
+        Returns:
+            Dictionary containing the configuration.
+        """
+        config = OpenAIClientConfig(
+            model=self.model,
+            api_base=self.api_base,
+            api_type=self.api_type,
+            api_version=self.api_version,
+            extra_parameters=self._extra_parameters,
+        )
+        return config.to_dict()
+
+    @property
+    def model(self) -> str:
+        """
+        Returns the model name for the openai embedding client.
+
+        Returns:
+            String representing the model name.
+        """
+        return self._model
+
+    @property
+    def api_base(self) -> Optional[str]:
+        """
+        Returns the base API URL for the openai embedding client.
+
+        Returns:
+            String representing the base API URL.
+        """
+        return self._api_base
+
+    @property
+    def api_type(self) -> Optional[str]:
+        """
+        Returns the API type for the openai embedding client.
+
+        Returns:
+            String representing the API version.
+        """
+        return self._api_type
+
+    @property
+    def api_version(self) -> Optional[str]:
+        """
+        Returns the API version for the openai embedding client.
+
+        Returns:
+            String representing the API version.
+        """
+        return self._api_version
+
+    @property
+    def _litellm_extra_parameters(self) -> Dict[str, Any]:
+        """
+        Returns the model parameters for the openai embedding client.
+
+        Returns:
+            Dictionary containing the model parameters.
+        """
+        return self._extra_parameters
+
+    @property
+    def _embedding_fn_args(self) -> Dict[str, Any]:
+        return {
+            "model": self._litellm_model_name,
+            "api_base": self.api_base,
+            "api_type": self.api_type,
+            "api_version": self.api_version,
+            **self._litellm_extra_parameters,
+        }
+
+    @property
+    def _litellm_model_name(self) -> str:
+        """Returns the LiteLLM model name for the openai embedding client."""
+        if self._model and f"{OPENAI_PROVIDER}/" not in self._model:
+            return f"{OPENAI_PROVIDER}/{self._model}"
+        return self._model

rasa/shared/providers/llm/_base_litellm_client.py

@@ -0,0 +1,227 @@
+from abc import abstractmethod
+from typing import Dict, List, Any, Union
+
+import logging
+import structlog
+from litellm import completion, acompletion, validate_environment
+
+from rasa.shared.exceptions import (
+    ProviderClientAPIException,
+    ProviderClientValidationError,
+)
+from rasa.shared.providers._ssl_verification_utils import (
+    ensure_ssl_certificates_for_litellm_non_openai_based_clients,
+    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
+
+structlogger = structlog.get_logger()
+
+_VALIDATE_ENVIRONMENT_MISSING_KEYS_KEY = "missing_keys"
+
+# Suppress LiteLLM info and debug logs - Global level.
+logging.getLogger("LiteLLM").setLevel(logging.WARNING)
+
+
+class _BaseLiteLLMClient:
+    """
+    An abstract base class for LiteLLM clients.
+
+    This class defines the interface and common functionality for all 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 clients that are currently based on
+    LiteLLM.
+
+    By keeping it private, we ensure that only the derived, concrete
+    implementations are exposed to users, maintaining a cleaner and
+    more controlled API surface.
+    """
+
+    def __init__(self):  # type: ignore
+        self._ensure_certificates()
+
+    @classmethod
+    @abstractmethod
+    def from_config(cls, config: Dict[str, Any]) -> "_BaseLiteLLMClient":
+        pass
+
+    @property
+    @abstractmethod
+    def config(self) -> dict:
+        """Returns the configuration for that the llm client
+        in dictionary form.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def _litellm_model_name(self) -> str:
+        """Returns the value of LiteLLM's model parameter to be used in
+        completion/acompletion in LiteLLM format:
+
+        <provider>/<model or deployment name>
+        """
+        pass
+
+    @property
+    def _litellm_extra_parameters(self) -> Dict[str, Any]:
+        """Returns a dictionary of extra parameters which include model
+        parameters as well as LiteLLM specific input parameters.
+
+        By default, this returns an empty dictionary (no extra parameters).
+        """
+        return {}
+
+    @property
+    def _completion_fn_args(self) -> dict:
+        return {
+            **self._litellm_extra_parameters,
+            "model": self._litellm_model_name,
+            # Since all providers covered by LiteLLM use the OpenAI format, but
+            # not all support every OpenAI parameter, raise an exception if
+            # provider/model uses unsupported parameter
+            "drop_params": False,
+        }
+
+    def validate_client_setup(self) -> None:
+        """Perform client validation. By default only environment variables
+        are validated.
+
+        Raises:
+            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)
+        if missing_environment_variables := validation_info.get(
+            _VALIDATE_ENVIRONMENT_MISSING_KEYS_KEY
+        ):
+            event_info = (
+                f"Environment variables: {missing_environment_variables} "
+                f"not set. Required for API calls."
+            )
+            structlogger.error(
+                "base_litellm_client.validate_environment_variables",
+                event_info=event_info,
+                missing_environment_variables=missing_environment_variables,
+            )
+            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)
+
+    @suppress_logs(log_level=logging.WARNING)
+    def completion(self, messages: Union[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.
+        Returns:
+            List of message completions.
+        Raises:
+            ProviderClientAPIException: If the API request fails.
+        """
+        try:
+            formatted_messages = self._format_messages(messages)
+            response = completion(
+                messages=formatted_messages, **self._completion_fn_args
+            )
+            return self._format_response(response)
+        except Exception as e:
+            raise ProviderClientAPIException(e)
+
+    @suppress_logs(log_level=logging.WARNING)
+    async def acompletion(self, messages: Union[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.
+        Returns:
+            List of message completions.
+        Raises:
+            ProviderClientAPIException: If the API request fails.
+        """
+        try:
+            formatted_messages = self._format_messages(messages)
+            response = await acompletion(
+                messages=formatted_messages, **self._completion_fn_args
+            )
+            return self._format_response(response)
+        except Exception as e:
+            raise ProviderClientAPIException(e)
+
+    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]
+
+    def _format_response(self, response: Any) -> LLMResponse:
+        """Parses the LiteLLM response to Rasa format."""
+        formatted_response = LLMResponse(
+            id=response.id,
+            created=response.created,
+            choices=[choice.message.content for choice in response.choices],
+            model=response.model,
+        )
+        if (
+            response.model_extra
+            and (usage := response.model_extra.get("usage")) is not None
+        ):
+            # We use `.get()` for accessing litellm.utils.Usage attributes.
+            # litellm.utils.Usage does not set the attributes if
+            # `prompt_tokens` or `completion_tokens` are absent (None).
+            prompt_tokens = (
+                num_tokens
+                if isinstance(num_tokens := usage.get("prompt_tokens", 0), (int, float))
+                else 0
+            )
+            completion_tokens = (
+                num_tokens
+                if isinstance(
+                    num_tokens := usage.get("completion_tokens", 0), (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
+        client initialization.
+
+        LiteLLM may utilize `openai` clients or other providers that require
+        SSL verification settings through the `SSL_VERIFY` / `SSL_CERTIFICATE`
+        environment variables or the `litellm.ssl_verify` /
+        `litellm.ssl_certificate` global settings.
+
+        This method ensures proper SSL configuration for both cases.
+        """
+        ensure_ssl_certificates_for_litellm_non_openai_based_clients()
+        ensure_ssl_certificates_for_litellm_openai_based_clients()