rasa-pro 3.9.18__py3-none-any.whl → 3.10.16__py3-none-any.whl

rasa/shared/providers/llm/self_hosted_llm_client.py

@@ -0,0 +1,293 @@
+from typing import Any, Dict, List, Optional, Union
+from litellm import (
+    text_completion,
+    atext_completion,
+)
+import logging
+import os
+import structlog
+
+from rasa.shared.constants import (
+    SELF_HOSTED_VLLM_PREFIX,
+    SELF_HOSTED_VLLM_API_KEY_ENV_VAR,
+)
+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.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 {}
+        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:
+
+        <openai>/<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(
+            {
+                "api_base": self.api_base,
+                "api_version": 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)
+
+    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):
+            os.environ[SELF_HOSTED_VLLM_API_KEY_ENV_VAR] = (
+                "dummy_self_hosted_llm_api_key"
+            )

rasa/shared/providers/mappings.py

@@ -0,0 +1,75 @@
+from typing import Dict, Type, Optional
+
+from rasa.shared.constants import (
+    AZURE_OPENAI_PROVIDER,
+    HUGGINGFACE_LOCAL_EMBEDDING_PROVIDER,
+    OPENAI_PROVIDER,
+    SELF_HOSTED_PROVIDER,
+)
+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._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.client_config import ClientConfig
+
+_provider_to_llm_client_mapping: Dict[str, Type[LLMClient]] = {
+    OPENAI_PROVIDER: OpenAILLMClient,
+    AZURE_OPENAI_PROVIDER: AzureOpenAILLMClient,
+    SELF_HOSTED_PROVIDER: SelfHostedLLMClient,
+}
+
+_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,
+}
+
+
+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
+    )

rasa/shared/utils/cli.py

@@ -1,3 +1,5 @@
+import math
+import shutil
 import sys
 from typing import Any, Text, NoReturn
 
@@ -70,3 +72,31 @@ def print_error_and_exit(message: Text, exit_code: int = 1) -> NoReturn:
     """
     print_error(message)
     sys.exit(exit_code)
+
+
+def pad(text: Text, char: Text = "=", min: int = 3) -> Text:
+    """Pad text to a certain length.
+
+    Uses `char` to pad the text to the specified length. If the text is longer
+    than the specified length, at least `min` are used.
+
+    The padding is applied to the left and right of the text (almost) equally.
+
+    Example:
+        >>> pad("Hello")
+        "========= Hello ========"
+        >>> pad("Hello", char="-")
+        "--------- Hello --------"
+
+    Args:
+        text: Text to pad.
+        min: Minimum length of the padding.
+        char: Character to pad with.
+
+    Returns:
+        Padded text.
+    """
+    width = shutil.get_terminal_size((80, 20)).columns
+    padding = max(width - len(text) - 2, min * 2)
+
+    return char * (padding // 2) + " " + text + " " + char * math.ceil(padding / 2)

rasa/shared/utils/io.py

@@ -1,12 +1,15 @@
 from collections import OrderedDict
+from functools import wraps
+from hashlib import md5
+import asyncio
 import errno
 import glob
-from hashlib import md5
 import json
+import logging
 import os
 import sys
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Text, Type, Union
+from typing import Any, cast, Callable, Dict, List, Optional, Text, Type, TypeVar, Union
 import warnings
 import random
 import string
@@ -137,6 +140,17 @@ def read_json_file(filename: Union[Text, Path]) -> Any:
         )
 
 
+def read_jsonl_file(file_path: Union[Text, Path]) -> List[Any]:
+    """Read JSONL from a file."""
+    content = read_file(file_path)
+    try:
+        return [json.loads(line) for line in content.splitlines()]
+    except ValueError as e:
+        raise FileIOException(
+            f"Failed to read JSONL from '{os.path.abspath(file_path)}'. Error: {e}"
+        )
+
+
 def list_directory(path: Text) -> List[Text]:
     """Returns all files and folders excluding hidden files.
 
@@ -413,3 +427,52 @@ def file_as_bytes(file_path: Text) -> bytes:
         raise FileNotFoundException(
             f"Failed to read file, " f"'{os.path.abspath(file_path)}' does not exist."
         )
+
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def suppress_logs(log_level: int = logging.WARNING) -> Callable[[F], F]:
+    """Decorator to suppress logs during the execution of a function.
+
+    Args:
+        log_level: The log level to set during the execution of the function.
+
+    Returns:
+        The decorated function.
+    """
+
+    def decorator(func: F) -> F:
+        @wraps(func)
+        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+            # Store the original logging level and set the new level.
+            original_logging_level = logging.getLogger().getEffectiveLevel()
+            logging.getLogger().setLevel(log_level)
+            try:
+                # Execute the async function.
+                result = await func(*args, **kwargs)
+            finally:
+                # Reset the logging level to the original level.
+                logging.getLogger().setLevel(original_logging_level)
+            return result
+
+        @wraps(func)
+        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+            # Store the original logging level and set the new level.
+            original_logging_level = logging.getLogger().getEffectiveLevel()
+            logging.getLogger().setLevel(log_level)
+            try:
+                # Execute the function.
+                result = func(*args, **kwargs)
+            finally:
+                # Reset the logging level to the original level.
+                logging.getLogger().setLevel(original_logging_level)
+            return result
+
+        # Determine if the function is async or not
+        if asyncio.iscoroutinefunction(func):
+            return cast(F, async_wrapper)
+        else:
+            return cast(F, sync_wrapper)
+
+    return decorator