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

rasa/shared/providers/embedding/huggingface_local_embedding_client.py

@@ -0,0 +1,191 @@
+import asyncio
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+import structlog
+
+from rasa.shared.constants import (
+    HUGGINGFACE_LOCAL_EMBEDDING_CACHING_FOLDER,
+)
+from rasa.shared.providers._configs.huggingface_local_embedding_client_config import (
+    HuggingFaceLocalEmbeddingClientConfig,
+)
+from rasa.shared.providers.embedding.embedding_response import EmbeddingResponse
+
+structlogger = structlog.get_logger()
+
+
+class HuggingFaceLocalEmbeddingClient:
+    """This client facilitates the computation of embeddings using locally
+    stored HuggingFace transformer models. It utilizes the
+    `sentence-transformers` package.
+
+    To use, you should have the `sentence_transformers` python package
+    installed.
+
+    Parameters:
+        model (str): Identifier for the model or path to a directory containing
+            model files.
+        cache_folder (Optional[str]): Path to the folder for caching models.
+            Default is defined by `HUGGINGFACE_LOCAL_EMBEDDING_CACHING_FOLDER`.
+        model_kwargs (Optional[Dict[str, Any]]): Additional keyword arguments
+            for model configuration.
+        encode_kwargs (Optional[Dict[str, Any]]): Parameters to control the
+            encoding process.
+        multi_process (bool): Flag to enable or disable multiprocessing for
+            embedding generation.
+        show_progress (bool): Whether to show a progress bar during embedding
+            generation.
+    """
+
+    def __init__(
+        self,
+        model: str,
+        cache_folder: Optional[str] = str(HUGGINGFACE_LOCAL_EMBEDDING_CACHING_FOLDER),
+        model_kwargs: Optional[Dict[str, Any]] = None,
+        encode_kwargs: Optional[Dict[str, Any]] = None,
+        multi_process: bool = False,
+        show_progress: bool = True,
+        **kwargs: Any,
+    ):
+        super().__init__(**kwargs)
+        self._model_name_or_path = model
+        self._multi_process = multi_process
+        self._cache_folder = cache_folder
+        self._model_kwargs = model_kwargs or {}
+        self._encode_kwargs = encode_kwargs or {}
+        self._show_progress = show_progress
+        self._init_client()
+
+    def _init_client(self) -> None:
+        self._validate_if_sentence_transformers_installed()
+
+        import sentence_transformers
+
+        self._client = sentence_transformers.SentenceTransformer(
+            model_name_or_path=self._model_name_or_path,
+            cache_folder=self._cache_folder,
+            **self._model_kwargs,
+        )
+
+    @property
+    def model(self) -> str:
+        return self._model_name_or_path
+
+    @property
+    def config(self) -> Dict:
+        config = HuggingFaceLocalEmbeddingClientConfig(
+            model=self._model_name_or_path,
+            cache_folder=self._cache_folder,
+            show_progress=self._show_progress,
+            model_kwargs=self._model_kwargs,
+            encode_kwargs=self._encode_kwargs,
+            multi_process=self._multi_process,
+        )
+        return config.to_dict()
+
+    @classmethod
+    def from_config(cls, config: Dict[str, Any]) -> "HuggingFaceLocalEmbeddingClient":
+        huggingface_config = HuggingFaceLocalEmbeddingClientConfig.from_dict(config)
+        return cls(
+            model=huggingface_config.model,
+            cache_folder=huggingface_config.cache_folder,
+            model_kwargs=huggingface_config.model_kwargs,
+            encode_kwargs=huggingface_config.encode_kwargs,
+            show_progress=huggingface_config.show_progress,
+        )
+
+    @staticmethod
+    def _validate_if_sentence_transformers_installed() -> None:
+        """In order to use the local embeddings from HuggingFace,
+        sentence-transformer must be installed.
+
+        Raises:
+            ImportError: If 'sentence-transformers' is not installed.
+        """
+        try:
+            import sentence_transformers  # noqa: F401
+        except ImportError as exc:
+            message = (
+                "Could not import sentence_transformers python package. "
+                "Please install it with 'pip install sentence-transformers'."
+            )
+            structlogger.error(
+                "huggingface_local_embedding_client.sentence_transformers_not_installed",
+                message=message,
+            )
+            raise ImportError(message) from exc
+
+    def embed(self, documents: List[str]) -> EmbeddingResponse:
+        """Compute doc embeddings using a HuggingFace transformer model.
+
+        Args:
+            documents: The list of documents to embed.
+
+        Returns:
+            List of embeddings, one for each text.
+        """
+        self._validate_if_sentence_transformers_installed()
+        documents = list(map(lambda x: x.replace("\n", " "), documents))
+
+        if self._multi_process:
+            embeddings = self._embed_with_multiprocessing(documents)
+        else:
+            embeddings = self._embed_without_multiprocessing(documents)
+
+        return EmbeddingResponse(
+            data=embeddings.tolist(),
+            model=self._model_name_or_path,
+        )
+
+    def _embed_without_multiprocessing(self, documents: List[str]) -> np.ndarray:
+        embeddings: np.ndarray = self._client.encode(
+            documents,
+            show_progress_bar=self._show_progress,
+            convert_to_numpy=True,
+            **self._encode_kwargs,
+        )
+        return embeddings
+
+    def _embed_with_multiprocessing(self, documents: List[str]) -> np.ndarray:
+        from sentence_transformers import SentenceTransformer
+
+        pool = self._client.start_multi_process_pool()
+        embeddings = self._client.encode_multi_process(documents, pool)
+        SentenceTransformer.stop_multi_process_pool(pool)
+        return embeddings
+
+    async def aembed(self, documents: List[str]) -> EmbeddingResponse:
+        """Asynchronous Embed search docs.
+
+        Args:
+            documents: List of documents to embed.
+
+        Returns:
+            List of embeddings.
+        """
+        loop = asyncio.get_running_loop()
+        # Using run_in_executor to execute the synchronous embedding function
+        # in a separate thread
+        embeddings = await loop.run_in_executor(None, self.embed, documents)
+        return embeddings
+
+    def validate_documents(self, documents: List[str]) -> None:
+        for doc in documents:
+            if not isinstance(doc, str):
+                message = "All documents must be strings."
+                structlogger.error(
+                    message=message,
+                    doc=doc,
+                )
+                raise ValueError(message)
+            if not doc.strip():
+                message = "Documents cannot be empty or whitespace."
+                structlogger.error(
+                    message=message,
+                    doc=doc,
+                )
+                raise ValueError(message)
+
+    def validate_client_setup(self, *args, **kwargs) -> None:  # type: ignore
+        pass

rasa/shared/providers/embedding/litellm_router_embedding_client.py

@@ -0,0 +1,138 @@
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, List
+
+import structlog
+
+from rasa.shared.exceptions import ProviderClientAPIException
+from rasa.shared.providers._configs.litellm_router_client_config import (
+    LiteLLMRouterClientConfig,
+)
+from rasa.shared.providers.embedding._base_litellm_embedding_client import (
+    _BaseLiteLLMEmbeddingClient,
+)
+from rasa.shared.providers.embedding.embedding_response import EmbeddingResponse
+from rasa.shared.providers.router._base_litellm_router_client import (
+    _BaseLiteLLMRouterClient,
+)
+from rasa.shared.utils.io import suppress_logs
+
+structlogger = structlog.get_logger()
+
+
+class LiteLLMRouterEmbeddingClient(
+    _BaseLiteLLMRouterClient, _BaseLiteLLMEmbeddingClient
+):
+    """A client for interfacing with LiteLLM Router Embedding endpoints.
+
+    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],
+        **kwargs: Any,
+    ):
+        super().__init__(
+            model_group_id, model_configurations, router_settings, **kwargs
+        )
+
+    @classmethod
+    def from_config(cls, config: Dict[str, Any]) -> LiteLLMRouterEmbeddingClient:
+        """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,
+            **client_config.extra_parameters,
+        )
+
+    @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:
+            response = self.router_client.embedding(
+                input=documents, **self._embedding_fn_args
+            )
+            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:
+            response = await self.router_client.aembedding(
+                input=documents, **self._embedding_fn_args
+            )
+            return self._format_response(response)
+        except Exception as e:
+            raise ProviderClientAPIException(
+                message="Failed to embed documents", original_exception=e
+            )
+
+    @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._model_group_id,
+        }

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