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

rasa/shared/providers/llm/default_litellm_llm_client.py

@@ -0,0 +1,108 @@
+from typing import Dict, Any
+
+from rasa.shared.constants import (
+    AWS_BEDROCK_PROVIDER,
+    AWS_SAGEMAKER_PROVIDER,
+)
+from rasa.shared.providers._configs.default_litellm_client_config import (
+    DefaultLiteLLMClientConfig,
+)
+from rasa.shared.providers._utils import validate_aws_setup_for_litellm_clients
+from rasa.shared.providers.llm._base_litellm_client import _BaseLiteLLMClient
+
+
+class DefaultLiteLLMClient(_BaseLiteLLMClient):
+    """A default client for interfacing with LiteLLM LLM endpoints.
+
+    Parameters:
+        model (str): The model or deployment name.
+        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:
+
+            https://docs.litellm.ai/docs/completion/input
+    Raises:
+        ProviderClientValidationError: If validation of the client setup fails.
+        ProviderClientAPIException: If the API request fails.
+    """
+
+    def __init__(self, provider: str, model: str, **kwargs: Any):
+        super().__init__()  # type: ignore
+        self._provider = provider
+        self._model = model
+        self._extra_parameters = kwargs
+        self.validate_client_setup()
+
+    @classmethod
+    def from_config(cls, config: Dict[str, Any]) -> "DefaultLiteLLMClient":
+        default_config = DefaultLiteLLMClientConfig.from_dict(config)
+        return cls(
+            model=default_config.model,
+            provider=default_config.provider,
+            # Pass the rest of the configuration as extra parameters
+            **default_config.extra_parameters,
+        )
+
+    @property
+    def provider(self) -> str:
+        return self._provider
+
+    @property
+    def model(self) -> str:
+        """
+
+        Returns:
+        """
+        return self._model
+
+    @property
+    def config(self) -> Dict:
+        """
+        Returns the configuration for the openai embedding client.
+        Returns:
+            Dictionary containing the configuration.
+        """
+        config = DefaultLiteLLMClientConfig(
+            model=self._model,
+            provider=self._provider,
+            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:
+
+        <provider>/<model or deployment name>
+        """
+        if self.model and f"{self.provider}/" not in self.model:
+            return f"{self.provider}/{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
+
+    def validate_client_setup(self) -> None:
+        # TODO: Temporarily change the environment variable validation for AWS setup
+        #       (Bedrock and SageMaker) until resolved by either:
+        #       1. An update from the LiteLLM package addressing the issue.
+        #       2. The implementation of a Bedrock client on our end.
+        #       ---
+        #       This fix ensures a consistent user experience for Bedrock (and
+        #       SageMaker) in Rasa by allowing AWS secrets to be provided as extra
+        #       parameters without triggering validation errors due to missing AWS
+        #       environment variables.
+        if self.provider.lower() in [AWS_BEDROCK_PROVIDER, AWS_SAGEMAKER_PROVIDER]:
+            validate_aws_setup_for_litellm_clients(
+                self._litellm_model_name,
+                self._litellm_extra_parameters,
+                "default_litellm_llm_client",
+            )
+        else:
+            super().validate_client_setup()

rasa/shared/providers/llm/litellm_router_llm_client.py

@@ -0,0 +1,182 @@
+from typing import Any, Dict, List, Union
+import logging
+import structlog
+
+from rasa.shared.exceptions import ProviderClientAPIException
+from rasa.shared.providers._configs.litellm_router_client_config import (
+    LiteLLMRouterClientConfig,
+)
+from rasa.shared.providers.llm._base_litellm_client import _BaseLiteLLMClient
+from rasa.shared.providers.llm.llm_response import LLMResponse
+from rasa.shared.providers.router._base_litellm_router_client import (
+    _BaseLiteLLMRouterClient,
+)
+from rasa.shared.utils.io import suppress_logs
+
+structlogger = structlog.get_logger()
+
+
+class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
+    """A client for interfacing with LiteLLM Router LLM 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]) -> "LiteLLMRouterLLMClient":
+        """Instantiates a LiteLLM Router LLM 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,
+        )
+
+    @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 = self.router_client.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 self.router_client.atext_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)
+    def completion(self, messages: Union[List[str], str]) -> LLMResponse:
+        """
+        Synchronously generate completions for given list of 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: List of messages or a single message to generate the
+                completion for.
+        Returns:
+            List of message completions.
+        Raises:
+            ProviderClientAPIException: If the API request fails.
+        """
+        if not self._use_chat_completions_endpoint:
+            return self._text_completion(messages)
+        try:
+            formatted_messages = self._format_messages(messages)
+            response = self.router_client.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.
+
+        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: List of messages or a single message to generate the
+                completion for.
+        Returns:
+            List of message completions.
+        Raises:
+            ProviderClientAPIException: If the API request fails.
+        """
+        if not self._use_chat_completions_endpoint:
+            return await self._atext_completion(messages)
+        try:
+            formatted_messages = self._format_messages(messages)
+            response = await self.router_client.acompletion(
+                messages=formatted_messages, **self._completion_fn_args
+            )
+            return self._format_response(response)
+        except Exception as e:
+            raise ProviderClientAPIException(e)
+
+    @property
+    def _completion_fn_args(self) -> Dict[str, Any]:
+        """Returns the completion arguments for invoking a call through
+        LiteLLM's completion functions.
+        """
+        return {
+            **self._litellm_extra_parameters,
+            "model": self.model_group_id,
+        }

rasa/shared/providers/llm/llm_client.py

@@ -0,0 +1,76 @@
+from typing import Protocol, Dict, List, runtime_checkable, Union
+
+from rasa.shared.providers.llm.llm_response import LLMResponse
+
+
+@runtime_checkable
+class LLMClient(Protocol):
+    """
+    Protocol for an LLM client that specifies the interface for interacting
+    with the API.
+    """
+
+    @classmethod
+    def from_config(cls, config: dict) -> "LLMClient":
+        """
+        Initializes the llm client with the given configuration.
+
+        This class method should be implemented to parse the given
+        configuration and create an instance of an llm client.
+        """
+        ...
+
+    @property
+    def config(self) -> Dict:
+        """
+        Returns the configuration for that the llm client is initialized with.
+
+        This property should be implemented to return a dictionary containing
+        the configuration settings for the llm client.
+        """
+        ...
+
+    def completion(self, messages: Union[List[str], str]) -> LLMResponse:
+        """
+        Synchronously generate completions for given list of messages.
+
+        This method should be implemented to take a list of messages (as
+        strings) and return a list of completions (as strings).
+
+        Args:
+            messages: List of messages or a single message to generate the
+                completion for.
+        Returns:
+            LLMResponse
+        """
+        ...
+
+    async def acompletion(self, messages: Union[List[str], str]) -> LLMResponse:
+        """
+        Asynchronously generate completions for given list of messages.
+
+        This method should be implemented to take a list of messages (as
+        strings) and return a list of completions (as strings).
+
+        Args:
+            messages: List of messages or a single message to generate the
+                completion for.
+        Returns:
+            LLMResponse
+        """
+        ...
+
+    def validate_client_setup(self, *args, **kwargs) -> None:  # type: ignore
+        """
+        Perform client setup validation.
+
+        This method should be implemented to validate whether the client can be
+        used with the parameters provided through configuration or environment
+        variables.
+
+        If there are any issues, the client should raise
+        ProviderClientValidationError.
+
+        If no validation is needed, this check can simply pass.
+        """
+        ...

rasa/shared/providers/llm/llm_response.py

@@ -0,0 +1,79 @@
+from dataclasses import dataclass, field, asdict
+from typing import Dict, List, Optional, Text, Any
+
+
+@dataclass
+class LLMUsage:
+    prompt_tokens: int
+    """Number of prompt tokens used to generate completion."""
+
+    completion_tokens: int
+    """Number of generated tokens."""
+
+    total_tokens: int = field(init=False)
+    """Total number of used tokens."""
+
+    def __post_init__(self) -> None:
+        self.total_tokens = self.prompt_tokens + self.completion_tokens
+
+    @classmethod
+    def from_dict(cls, data: Dict[Text, Any]) -> "LLMUsage":
+        """
+        Creates an LLMUsage object from a dictionary.
+        If any keys are missing, they will default to zero
+        or whatever default you prefer.
+        """
+        return cls(
+            prompt_tokens=data.get("prompt_tokens"),
+            completion_tokens=data.get("completion_tokens"),
+        )
+
+    def to_dict(self) -> dict:
+        """Converts the LLMUsage dataclass instance into a dictionary."""
+        return asdict(self)
+
+
+@dataclass
+class LLMResponse:
+    id: str
+    """A unique identifier for the completion."""
+
+    choices: List[str]
+    """The list of completion choices the model generated for the input prompt."""
+
+    created: int
+    """The Unix timestamp (in seconds) of when the completion was created."""
+
+    model: Optional[str] = None
+    """The model used for completion."""
+
+    usage: Optional[LLMUsage] = None
+    """An optional details about the token usage for the API call."""
+
+    additional_info: Optional[Dict] = None
+    """Optional dictionary for storing additional information related to the
+    completion that may not be covered by other fields."""
+
+    @classmethod
+    def from_dict(cls, data: Dict[Text, Any]) -> "LLMResponse":
+        """
+        Creates an LLMResponse from a dictionary.
+        """
+        usage_data = data.get("usage")
+        usage_obj = LLMUsage.from_dict(usage_data) if usage_data else None
+
+        return cls(
+            id=data["id"],
+            choices=data["choices"],
+            created=data["created"],
+            model=data.get("model"),
+            usage=usage_obj,
+            additional_info=data.get("additional_info"),
+        )
+
+    def to_dict(self) -> dict:
+        """Converts the LLMResponse dataclass instance into a dictionary."""
+        result = asdict(self)
+        if self.usage:
+            result["usage"] = self.usage.to_dict()
+        return result

rasa/shared/providers/llm/openai_llm_client.py

@@ -0,0 +1,155 @@
+import os
+import re
+from typing import Dict, Any, Optional
+
+import structlog
+
+from rasa.shared.constants import (
+    OPENAI_API_BASE_ENV_VAR,
+    OPENAI_API_VERSION_ENV_VAR,
+    OPENAI_API_TYPE_ENV_VAR,
+    OPENAI_PROVIDER,
+)
+from rasa.shared.providers._configs.openai_client_config import OpenAIClientConfig
+from rasa.shared.providers.llm._base_litellm_client import _BaseLiteLLMClient
+
+structlogger = structlog.get_logger()
+
+
+class OpenAILLMClient(_BaseLiteLLMClient):
+    """
+    A client for interfacing with OpenAI LLMs.
+
+    Parameters:
+        model (str): The OpenAI model name.
+        api_base (Optional[str]): Optional, the base URL for the API endpoints.
+            If not provided, it will try to be set via environment variables.
+        api_version (Optional[str]): Optional, the version of the API to use.
+            If not provided, it will try to be set via environment variable.
+        api_type: (Optional[str]): The api type. If not provided, it will be set via
+            environment variable (also optional).
+        kwargs (Optional[Dict[str, Any]]): Optional configuration parameters specific
+            to the model.
+
+    Raises:
+        ProviderClientValidationError: If validation of the client setup fails.
+    """
+
+    def __init__(
+        self,
+        model: str,
+        api_base: Optional[str] = None,
+        api_version: Optional[str] = None,
+        api_type: 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, None)
+        self._api_version = api_version or os.environ.get(
+            OPENAI_API_VERSION_ENV_VAR, None
+        )
+
+        # Not used by LiteLLM, here for backward compatibility
+        self._api_type = api_type or os.environ.get(OPENAI_API_TYPE_ENV_VAR)
+
+        self._extra_parameters = kwargs or {}
+        self.validate_client_setup()
+
+    @classmethod
+    def from_config(cls, config: Dict[str, Any]) -> "OpenAILLMClient":
+        """
+        Initializes the client from given configuration.
+
+        Args:
+            config (Dict[str, Any]): Configuration.
+
+        Raises:
+            KeyError: If any of the required configuration keys are missing.
+            ValueError: If `api_type` has a value different from `openai`.
+
+        Returns:
+            AzureOpenAILLMClient: Initialized client.
+        """
+        try:
+            openai_config = OpenAIClientConfig.from_dict(config)
+        except (KeyError, ValueError) as e:
+            message = "Cannot instantiate a client from the passed configuration."
+            structlogger.error(
+                "openai_llm_client.from_config.error",
+                message=message,
+                config=config,
+                original_error=e,
+            )
+            raise
+
+        return cls(
+            openai_config.model,
+            openai_config.api_base,
+            openai_config.api_version,
+            openai_config.api_type,
+            **openai_config.extra_parameters,
+        )
+
+    @property
+    def config(self) -> dict:
+        config = OpenAIClientConfig(
+            model=self.model,
+            api_type=self.api_type,
+            api_base=self.api_base,
+            api_version=self.api_version,
+            extra_parameters=self._litellm_extra_parameters,
+        )
+        return config.to_dict()
+
+    @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 api_version(self) -> Optional[str]:
+        """
+        Returns the API version for the OpenAI LLM client.
+        """
+        return self._api_version
+
+    @property
+    def api_type(self) -> Optional[str]:
+        return self._api_type
+
+    @property
+    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>
+        """
+        regex_patter = rf"^{OPENAI_PROVIDER}/"
+        if not re.match(regex_patter, self._model):
+            return f"{OPENAI_PROVIDER}/{self._model}"
+        return self._model
+
+    @property
+    def _litellm_extra_parameters(self) -> Dict[str, Any]:
+        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