rasa-pro 3.11.3a1.dev2__py3-none-any.whl

rasa/shared/providers/_utils.py

@@ -0,0 +1,79 @@
+import structlog
+
+from rasa.shared.constants import (
+    AWS_ACCESS_KEY_ID_ENV_VAR,
+    AWS_ACCESS_KEY_ID_CONFIG_KEY,
+    AWS_SECRET_ACCESS_KEY_ENV_VAR,
+    AWS_SECRET_ACCESS_KEY_CONFIG_KEY,
+    AWS_REGION_NAME_ENV_VAR,
+    AWS_REGION_NAME_CONFIG_KEY,
+    AWS_SESSION_TOKEN_CONFIG_KEY,
+    AWS_SESSION_TOKEN_ENV_VAR,
+)
+from rasa.shared.exceptions import ProviderClientValidationError
+from litellm import validate_environment
+from rasa.shared.providers.embedding._base_litellm_embedding_client import (
+    _VALIDATE_ENVIRONMENT_MISSING_KEYS_KEY,
+)
+
+structlogger = structlog.get_logger()
+
+
+def validate_aws_setup_for_litellm_clients(
+    litellm_model_name: str, litellm_call_kwargs: dict, source_log: str
+) -> None:
+    """Validates the AWS setup for LiteLLM clients to ensure all required
+    environment variables or corresponding call kwargs are set.
+
+    Args:
+        litellm_model_name (str): The name of the LiteLLM model being validated.
+        litellm_call_kwargs (dict): Additional keyword arguments passed to the client,
+            which may include configuration values for AWS credentials.
+        source_log (str): The source log identifier for structured logging.
+
+    Raises:
+        ProviderClientValidationError: If any required AWS environment variable
+            or corresponding configuration key is missing.
+    """
+
+    # Mapping of environment variable names to their corresponding config keys
+    envs_to_args = {
+        AWS_ACCESS_KEY_ID_ENV_VAR: AWS_ACCESS_KEY_ID_CONFIG_KEY,
+        AWS_SECRET_ACCESS_KEY_ENV_VAR: AWS_SECRET_ACCESS_KEY_CONFIG_KEY,
+        AWS_REGION_NAME_ENV_VAR: AWS_REGION_NAME_CONFIG_KEY,
+        AWS_SESSION_TOKEN_ENV_VAR: AWS_SESSION_TOKEN_CONFIG_KEY,
+    }
+
+    # Validate the environment setup for the model
+    validation_info = validate_environment(litellm_model_name)
+    missing_environment_variables = validation_info.get(
+        _VALIDATE_ENVIRONMENT_MISSING_KEYS_KEY, []
+    )
+    # Filter out missing environment variables that have been set trough arguments
+    # in extra parameters
+    missing_environment_variables = [
+        missing_env_var
+        for missing_env_var in missing_environment_variables
+        if litellm_call_kwargs.get(envs_to_args.get(missing_env_var)) is None
+    ]
+
+    if missing_environment_variables:
+        missing_environment_details = [
+            (
+                f"'{missing_env_var}' environment variable or "
+                f"'{envs_to_args.get(missing_env_var)}' config key"
+            )
+            for missing_env_var in missing_environment_variables
+        ]
+        event_info = (
+            f"The following environment variables or configuration keys are "
+            f"missing: "
+            f"{', '.join(missing_environment_details)}. "
+            f"These settings are required for API calls."
+        )
+        structlogger.error(
+            f"{source_log}.validate_aws_environment_variables",
+            event_info=event_info,
+            missing_environment_variables=missing_environment_variables,
+        )
+        raise ProviderClientValidationError(event_info)

rasa/shared/providers/embedding/_base_litellm_embedding_client.py

@@ -0,0 +1,243 @@
+import logging
+from abc import abstractmethod
+from typing import Any, Dict, List
+
+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,
+)
+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.embedding.embedding_response import (
+    EmbeddingResponse,
+    EmbeddingUsage,
+)
+from rasa.shared.utils.io import suppress_logs, resolve_environment_variables
+
+structlogger = structlog.get_logger()
+
+_VALIDATE_ENVIRONMENT_MISSING_KEYS_KEY = "missing_keys"
+
+
+class _BaseLiteLLMEmbeddingClient:
+    """An abstract base class for LiteLLM embedding 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()
+
+    @property
+    @abstractmethod
+    def config(self) -> dict:
+        """Returns the configuration for that the embedding client in dict form."""
+        pass
+
+    @property
+    @abstractmethod
+    def _litellm_model_name(self) -> str:
+        """Returns the model name in LiteLLM format based on the Provider/API type."""
+        pass
+
+    @property
+    @abstractmethod
+    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 _embedding_fn_args(self) -> Dict[str, Any]:
+        """Returns the arguments to be passed to the embedding function."""
+        return {
+            **self._litellm_extra_parameters,
+            "model": self._litellm_model_name,
+        }
+
+    def validate_client_setup(self) -> None:
+        """Perform client validation. By default only environment variables
+        are validated. Override this method to add more validation steps.
+
+        Raises:
+            ProviderClientValidationError if validation fails.
+        """
+        self._validate_environment_variables()
+
+    def _validate_environment_variables(self) -> None:
+        """Validate that the required environment variables are set."""
+        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
+        ):
+            event_info = (
+                f"Environment variables: {missing_environment_variables} "
+                f"not set. Required for API calls."
+            )
+            structlogger.error(
+                "base_litellm_embedding_client.validate_environment_variables",
+                event_info=event_info,
+                missing_environment_variables=missing_environment_variables,
+            )
+            raise ProviderClientValidationError(event_info)
+
+    def validate_documents(self, documents: List[str]) -> None:
+        """Validates a list of documents to ensure they are suitable for embedding.
+
+        Args:
+            documents: List of documents to be validated.
+
+        Raises:
+            ValueError: If any document is invalid.
+        """
+        for doc in documents:
+            if not isinstance(doc, str):
+                raise ValueError("All documents must be strings.")
+            if not doc.strip():
+                raise ValueError("Documents cannot be empty or whitespace.")
+
+    @suppress_logs(log_level=logging.WARNING)
+    def embed(self, documents: List[str]) -> EmbeddingResponse:
+        """Embeds a list of documents synchronously.
+
+        Args:
+            documents: List of documents to be embedded.
+
+        Returns:
+            List of embedding vectors.
+
+        Raises:
+            ProviderClientAPIException: If API calls raised an error.
+        """
+        self.validate_documents(documents)
+        try:
+            arguments = resolve_environment_variables(self._embedding_fn_args)
+            response = embedding(input=documents, **arguments)
+            return self._format_response(response)
+        except Exception as e:
+            raise ProviderClientAPIException(
+                message="Failed to embed documents", original_exception=e
+            )
+
+    @suppress_logs(log_level=logging.WARNING)
+    async def aembed(self, documents: List[str]) -> EmbeddingResponse:
+        """Embeds a list of documents asynchronously.
+
+        Args:
+            documents: List of documents to be embedded.
+
+        Returns:
+            List of embedding vectors.
+
+        Raises:
+            ProviderClientAPIException: If API calls raised an error.
+        """
+        self.validate_documents(documents)
+        try:
+            arguments = resolve_environment_variables(self._embedding_fn_args)
+            response = await aembedding(input=documents, **arguments)
+            return self._format_response(response)
+        except Exception as e:
+            raise ProviderClientAPIException(
+                message="Failed to embed documents", original_exception=e
+            )
+
+    def _format_response(
+        self, response: litellm.EmbeddingResponse
+    ) -> EmbeddingResponse:
+        """Parses the LiteLLM EmbeddingResponse to Rasa format.
+
+        Raises:
+            ValueError: If any response data is None.
+        """
+        # If data is not available (None), raise a ValueError
+        if response.data is None:
+            message = (
+                "Failed to embed documents. Received 'None' " "instead of embeddings."
+            )
+            structlogger.error(
+                "base_litellm_client.format_response.data_is_none",
+                message=message,
+                response=response.to_dict(),
+            )
+            raise ValueError(message)
+
+        # Sort the embeddings by the "index" key
+        response.data.sort(key=lambda x: x["index"])
+        # Extract the embedding vectors
+        embeddings = [data["embedding"] for data in response.data]
+        formatted_response = EmbeddingResponse(
+            data=embeddings,
+            model=response.model,
+        )
+
+        # Process additional usage information if available
+        if response.usage:
+            completion_tokens = (
+                response.usage.completion_tokens
+                if hasattr(response.usage, "completion_tokens")
+                else 0
+            )
+            prompt_tokens = (
+                response.usage.prompt_tokens
+                if hasattr(response.usage, "prompt_tokens")
+                else 0
+            )
+            total_tokens = (
+                response.usage.total_tokens
+                if hasattr(response.usage, "total_tokens")
+                else 0
+            )
+
+            formatted_response.usage = EmbeddingUsage(
+                completion_tokens=completion_tokens,
+                prompt_tokens=prompt_tokens,
+                total_tokens=total_tokens,
+            )
+
+        # Log the response with masked data for brevity
+        log_response = formatted_response.to_dict()
+        log_response["data"] = "Embedding response data not shown here for brevity."
+        structlogger.debug(
+            "base_litellm_client.formatted_response",
+            formatted_response=log_response,
+        )
+        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()

rasa/shared/providers/embedding/_langchain_embedding_client_adapter.py

@@ -0,0 +1,74 @@
+from typing import List
+
+from langchain_core.embeddings.embeddings import Embeddings
+
+from rasa.shared.providers.embedding.embedding_client import EmbeddingClient
+
+
+class _LangchainEmbeddingClientAdapter(Embeddings):
+    """
+    Temporary adapter to bridge differences between LiteLLM and LangChain.
+
+    Clients instantiated with `embedder_factory` follow our new EmbeddingClient
+    protocol, but `langchain`'s vector stores require an `Embeddings` type
+    client. This adapter extracts and returns the necessary part of the output
+    from our LiteLLM-based clients.
+
+    This adapter will be removed in ticket:
+    https://rasahq.atlassian.net/browse/ENG-1220
+    """
+
+    def __init__(self, client: EmbeddingClient):
+        self._client = client
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        """Embed search docs.
+
+        Args:
+            texts: List of text to embed.
+
+        Returns:
+            List of embeddings.
+        """
+        response = self._client.embed(documents=texts)
+        embedding_vector = response.data
+        return embedding_vector
+
+    def embed_query(self, text: str) -> List[float]:
+        """Embed query text.
+
+        Args:
+            text: Text to embed.
+
+        Returns:
+            Embedding.
+        """
+        response = self._client.embed(documents=[text])
+        embedding_vector = response.data[0]
+        return embedding_vector
+
+    async def aembed_documents(self, texts: List[str]) -> List[List[float]]:
+        """Asynchronous Embed search docs.
+
+        Args:
+            texts: List of text to embed.
+
+        Returns:
+            List of embeddings.
+        """
+        response = await self._client.aembed(documents=texts)
+        embedding_vector = response.data
+        return embedding_vector
+
+    async def aembed_query(self, text: str) -> List[float]:
+        """Asynchronous Embed query text.
+
+        Args:
+            text: Text to embed.
+
+        Returns:
+            Embedding.
+        """
+        response = await self._client.aembed(documents=[text])
+        embedding_vector = response.data[0]
+        return embedding_vector