ragbits-core 0.0.1__py3-none-any.whl

ragbits/core/embeddings/base.py

@@ -0,0 +1,19 @@
+from abc import ABC, abstractmethod
+
+
+class Embeddings(ABC):
+    """
+    Abstract client for communication with embedding models.
+    """
+
+    @abstractmethod
+    async def embed_text(self, data: list[str]) -> list[list[float]]:
+        """
+        Creates embeddings for the given strings.
+
+        Args:
+            data: List of strings to get embeddings for.
+
+        Returns:
+            List of embeddings for the given strings.
+        """

ragbits/core/embeddings/exceptions.py

@@ -0,0 +1,36 @@
+class EmbeddingError(Exception):
+    """
+    Base class for all exceptions raised by the EmbeddingClient.
+    """
+
+    def __init__(self, message: str) -> None:
+        super().__init__(message)
+        self.message = message
+
+
+class EmbeddingConnectionError(EmbeddingError):
+    """
+    Raised when there is an error connecting to the embedding API.
+    """
+
+    def __init__(self, message: str = "Connection error.") -> None:
+        super().__init__(message)
+
+
+class EmbeddingStatusError(EmbeddingError):
+    """
+    Raised when an API response has a status code of 4xx or 5xx.
+    """
+
+    def __init__(self, message: str, status_code: int) -> None:
+        super().__init__(message)
+        self.status_code = status_code
+
+
+class EmbeddingResponseError(EmbeddingError):
+    """
+    Raised when an API response has an invalid schema.
+    """
+
+    def __init__(self, message: str = "Data returned by API invalid for expected schema.") -> None:
+        super().__init__(message)

ragbits/core/embeddings/litellm.py

@@ -0,0 +1,85 @@
+from typing import Optional
+
+try:
+    import litellm
+
+    HAS_LITELLM = True
+except ImportError:
+    HAS_LITELLM = False
+
+from ragbits.core.embeddings.base import Embeddings
+from ragbits.core.embeddings.exceptions import EmbeddingConnectionError, EmbeddingResponseError, EmbeddingStatusError
+
+
+class LiteLLMEmbeddings(Embeddings):
+    """
+    Client for creating text embeddings using LiteLLM API.
+    """
+
+    def __init__(
+        self,
+        model: str = "text-embedding-3-small",
+        options: Optional[dict] = None,
+        api_base: Optional[str] = None,
+        api_key: Optional[str] = None,
+        api_version: Optional[str] = None,
+    ) -> None:
+        """
+        Constructs the LiteLLMEmbeddingClient.
+
+        Args:
+            model: Name of the [LiteLLM supported model](https://docs.litellm.ai/docs/embedding/supported_embedding)\
+                to be used. Default is "text-embedding-3-small".
+            options: Additional options to pass to the LiteLLM API.
+            api_base: The API endpoint you want to call the model with.
+            api_key: API key to be used. API key to be used. If not specified, an environment variable will be used,
+                for more information, follow the instructions for your specific vendor in the\
+                [LiteLLM documentation](https://docs.litellm.ai/docs/embedding/supported_embedding).
+            api_version: The API version for the call.
+
+        Raises:
+            ImportError: If the 'litellm' extra requirements are not installed.
+        """
+        if not HAS_LITELLM:
+            raise ImportError("You need to install the 'litellm' extra requirements to use LiteLLM embeddings models")
+
+        super().__init__()
+        self.model = model
+        self.options = options or {}
+        self.api_base = api_base
+        self.api_key = api_key
+        self.api_version = api_version
+
+    async def embed_text(self, data: list[str]) -> list[list[float]]:
+        """
+        Creates embeddings for the given strings.
+
+        Args:
+            data: List of strings to get embeddings for.
+
+        Returns:
+            List of embeddings for the given strings.
+
+        Raises:
+            EmbeddingConnectionError: If there is a connection error with the embedding API.
+            EmbeddingStatusError: If the embedding API returns an error status code.
+            EmbeddingResponseError: If the embedding API response is invalid.
+        """
+
+        try:
+            response = await litellm.aembedding(
+                input=data,
+                model=self.model,
+                api_base=self.api_base,
+                api_key=self.api_key,
+                api_version=self.api_version,
+                **self.options,
+            )
+        except litellm.openai.APIConnectionError as exc:
+            raise EmbeddingConnectionError() from exc
+        except litellm.openai.APIStatusError as exc:
+            raise EmbeddingStatusError(exc.message, exc.status_code) from exc
+        except litellm.openai.APIResponseValidationError as exc:
+            raise EmbeddingResponseError() from exc
+
+        return [embedding["embedding"] for embedding in response.data]

ragbits/core/embeddings/local.py

@@ -0,0 +1,81 @@
+from typing import Iterator, Optional
+
+try:
+    import torch
+    import torch.nn.functional as F
+    from transformers import AutoModel, AutoTokenizer
+
+    HAS_LOCAL_EMBEDDINGS = True
+except ImportError:
+    HAS_LOCAL_EMBEDDINGS = False
+
+from ragbits.core.embeddings.base import Embeddings
+
+
+class LocalEmbeddings(Embeddings):
+    """
+    Class for interaction with any encoder available in HuggingFace.
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        api_key: Optional[str] = None,
+    ) -> None:
+        """
+        Constructs a new local LLM instance.
+
+        Args:
+            model_name: Name of the model to use.
+            api_key: The API key for Hugging Face authentication.
+
+        Raises:
+            ImportError: If the 'local' extra requirements are not installed.
+        """
+        if not HAS_LOCAL_EMBEDDINGS:
+            raise ImportError("You need to install the 'local' extra requirements to use local embeddings models")
+
+        super().__init__()
+
+        self.hf_api_key = api_key
+        self.model_name = model_name
+
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model = AutoModel.from_pretrained(self.model_name, token=self.hf_api_key).to(self.device)
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name, token=self.hf_api_key)
+
+    async def embed_text(self, data: list[str], batch_size: int = 1) -> list[list[float]]:
+        """
+        Calls the appropriate encoder endpoint with the given data and options.
+
+        Args:
+            data: List of strings to get embeddings for.
+            batch_size: Batch size.
+
+        Returns:
+            List of embeddings for the given strings.
+        """
+        embeddings = []
+        for batch in self._batch(data, batch_size):
+            batch_dict = self.tokenizer(
+                batch, max_length=self.tokenizer.model_max_length, padding=True, truncation=True, return_tensors="pt"
+            ).to(self.device)
+            with torch.no_grad():
+                outputs = self.model(**batch_dict)
+                batch_embeddings = self._average_pool(outputs.last_hidden_state, batch_dict["attention_mask"])
+                batch_embeddings = F.normalize(batch_embeddings, p=2, dim=1)
+            embeddings.extend(batch_embeddings.to("cpu").tolist())
+
+        torch.cuda.empty_cache()
+        return embeddings
+
+    @staticmethod
+    def _batch(data: list[str], batch_size: int) -> Iterator[list[str]]:
+        length = len(data)
+        for ndx in range(0, length, batch_size):
+            yield data[ndx : min(ndx + batch_size, length)]
+
+    @staticmethod
+    def _average_pool(last_hidden_states: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+        last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
+        return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]

ragbits/core/llms/__init__.py

@@ -0,0 +1,5 @@
+from .base import LLM
+from .litellm import LiteLLM
+from .local import LocalLLM
+
+__all__ = ["LLM", "LiteLLM", "LocalLLM"]

ragbits/core/llms/base.py

@@ -0,0 +1,121 @@
+from abc import ABC, abstractmethod
+from functools import cached_property
+from typing import Generic, Optional, Type, cast, overload
+
+from ragbits.core.prompt.base import BasePrompt, BasePromptWithParser, OutputT
+
+from .clients.base import LLMClient, LLMClientOptions, LLMOptions
+
+
+class LLM(Generic[LLMClientOptions], ABC):
+    """
+    Abstract class for interaction with Large Language Model.
+    """
+
+    _options_cls: Type[LLMClientOptions]
+
+    def __init__(self, model_name: str, default_options: Optional[LLMOptions] = None) -> None:
+        """
+        Constructs a new LLM instance.
+
+        Args:
+            model_name: Name of the model to be used.
+            default_options: Default options to be used.
+
+        Raises:
+            TypeError: If the subclass is missing the '_options_cls' attribute.
+        """
+        self.model_name = model_name
+        self.default_options = default_options or self._options_cls()
+
+    def __init_subclass__(cls) -> None:
+        if not hasattr(cls, "_options_cls"):
+            raise TypeError(f"Class {cls.__name__} is missing the '_options_cls' attribute")
+
+    @cached_property
+    @abstractmethod
+    def client(self) -> LLMClient:
+        """
+        Client for the LLM.
+        """
+
+    def count_tokens(self, prompt: BasePrompt) -> int:
+        """
+        Counts tokens in the prompt.
+
+        Args:
+            prompt: Formatted prompt template with conversation and response parsing configuration.
+
+        Returns:
+            Number of tokens in the prompt.
+        """
+        return sum(len(message["content"]) for message in prompt.chat)
+
+    async def generate_raw(
+        self,
+        prompt: BasePrompt,
+        *,
+        options: Optional[LLMOptions] = None,
+    ) -> str:
+        """
+        Prepares and sends a prompt to the LLM and returns the raw response (without parsing).
+
+        Args:
+            prompt: Formatted prompt template with conversation.
+            options: Options to use for the LLM client.
+
+        Returns:
+            Raw text response from LLM.
+        """
+        options = (self.default_options | options) if options else self.default_options
+
+        response = await self.client.call(
+            conversation=prompt.chat,
+            options=options,
+            json_mode=prompt.json_mode,
+            output_schema=prompt.output_schema(),
+        )
+
+        return response
+
+    @overload
+    async def generate(
+        self,
+        prompt: BasePromptWithParser[OutputT],
+        *,
+        options: Optional[LLMOptions] = None,
+    ) -> OutputT:
+        ...
+
+    @overload
+    async def generate(
+        self,
+        prompt: BasePrompt,
+        *,
+        options: Optional[LLMOptions] = None,
+    ) -> OutputT:
+        ...
+
+    async def generate(
+        self,
+        prompt: BasePrompt,
+        *,
+        options: Optional[LLMOptions] = None,
+    ) -> OutputT:
+        """
+        Prepares and sends a prompt to the LLM and returns response parsed to the
+        output type of the prompt (if available).
+
+        Args:
+            prompt: Formatted prompt template with conversation and optional response parsing configuration.
+            options: Options to use for the LLM client.
+
+        Returns:
+            Text response from LLM.
+        """
+        response = await self.generate_raw(prompt, options=options)
+
+        if isinstance(prompt, BasePromptWithParser):
+            return prompt.parse_response(response)
+
+        return cast(OutputT, response)

ragbits/core/llms/clients/__init__.py

@@ -0,0 +1,12 @@
+from .base import LLMClient, LLMOptions
+from .litellm import LiteLLMClient, LiteLLMOptions
+from .local import LocalLLMClient, LocalLLMOptions
+
+__all__ = [
+    "LLMClient",
+    "LLMOptions",
+    "LiteLLMClient",
+    "LiteLLMOptions",
+    "LocalLLMClient",
+    "LocalLLMOptions",
+]

ragbits/core/llms/clients/base.py

@@ -0,0 +1,86 @@
+from abc import ABC, abstractmethod
+from dataclasses import asdict, dataclass
+from typing import Any, ClassVar, Dict, Generic, Optional, Type, TypeVar
+
+from pydantic import BaseModel
+
+from ragbits.core.prompt import ChatFormat
+
+from ..types import NotGiven
+
+LLMClientOptions = TypeVar("LLMClientOptions", bound="LLMOptions")
+
+
+@dataclass
+class LLMOptions(ABC):
+    """
+    Abstract dataclass that represents all available LLM call options.
+    """
+
+    _not_given: ClassVar[Any] = None
+
+    def __or__(self, other: "LLMOptions") -> "LLMOptions":
+        """
+        Merges two LLMOptions, prioritizing non-NOT_GIVEN values from the 'other' object.
+        """
+        self_dict = asdict(self)
+        other_dict = asdict(other)
+
+        updated_dict = {
+            key: other_dict.get(key, self_dict[key])
+            if not isinstance(other_dict.get(key), NotGiven)
+            else self_dict[key]
+            for key in self_dict
+        }
+
+        return self.__class__(**updated_dict)
+
+    def dict(self) -> Dict[str, Any]:
+        """
+        Creates a dictionary representation of the LLMOptions instance.
+        If a value is None, it will be replaced with a provider-specific not-given sentinel.
+
+        Returns:
+            A dictionary representation of the LLMOptions instance.
+        """
+        options = asdict(self)
+        return {
+            key: self._not_given if value is None or isinstance(value, NotGiven) else value
+            for key, value in options.items()
+        }
+
+
+class LLMClient(Generic[LLMClientOptions], ABC):
+    """
+    Abstract client for a direct communication with LLM.
+    """
+
+    def __init__(self, model_name: str) -> None:
+        """
+        Constructs a new LLMClient instance.
+
+        Args:
+            model_name: Name of the model to be used.
+        """
+        self.model_name = model_name
+
+    @abstractmethod
+    async def call(
+        self,
+        conversation: ChatFormat,
+        options: LLMClientOptions,
+        json_mode: bool = False,
+        output_schema: Optional[Type[BaseModel] | Dict] = None,
+    ) -> str:
+        """
+        Calls LLM inference API.
+
+        Args:
+            conversation: List of dicts with "role" and "content" keys, representing the chat history so far.
+            options: Additional settings used by LLM.
+            json_mode: Force the response to be in JSON format.
+            output_schema: Schema for structured response (either Pydantic model or a JSON schema).
+
+        Returns:
+            Response string from LLM.
+        """

ragbits/core/llms/clients/exceptions.py

@@ -0,0 +1,36 @@
+class LLMError(Exception):
+    """
+    Base class for all exceptions raised by the LLMClient.
+    """
+
+    def __init__(self, message: str) -> None:
+        super().__init__(message)
+        self.message = message
+
+
+class LLMConnectionError(LLMError):
+    """
+    Raised when there is an error connecting to the LLM API.
+    """
+
+    def __init__(self, message: str = "Connection error.") -> None:
+        super().__init__(message)
+
+
+class LLMStatusError(LLMError):
+    """
+    Raised when an API response has a status code of 4xx or 5xx.
+    """
+
+    def __init__(self, message: str, status_code: int) -> None:
+        super().__init__(message)
+        self.status_code = status_code
+
+
+class LLMResponseError(LLMError):
+    """
+    Raised when an API response has an invalid schema.
+    """
+
+    def __init__(self, message: str = "Data returned by API invalid for expected schema.") -> None:
+        super().__init__(message)

ragbits/core/llms/clients/litellm.py

@@ -0,0 +1,129 @@
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Type, Union
+
+from pydantic import BaseModel
+
+try:
+    import litellm
+
+    HAS_LITELLM = True
+except ImportError:
+    HAS_LITELLM = False
+
+
+from ragbits.core.prompt import ChatFormat
+
+from ..types import NOT_GIVEN, NotGiven
+from .base import LLMClient, LLMOptions
+from .exceptions import LLMConnectionError, LLMResponseError, LLMStatusError
+
+
+@dataclass
+class LiteLLMOptions(LLMOptions):
+    """
+    Dataclass that represents all available LLM call options for the LiteLLM client.
+    Each of them is described in the [LiteLLM documentation](https://docs.litellm.ai/docs/completion/input).
+    """
+
+    frequency_penalty: Union[Optional[float], NotGiven] = NOT_GIVEN
+    max_tokens: Union[Optional[int], NotGiven] = NOT_GIVEN
+    n: Union[Optional[int], NotGiven] = NOT_GIVEN
+    presence_penalty: Union[Optional[float], NotGiven] = NOT_GIVEN
+    seed: Union[Optional[int], NotGiven] = NOT_GIVEN
+    stop: Union[Optional[Union[str, List[str]]], NotGiven] = NOT_GIVEN
+    temperature: Union[Optional[float], NotGiven] = NOT_GIVEN
+    top_p: Union[Optional[float], NotGiven] = NOT_GIVEN
+    mock_response: Union[Optional[str], NotGiven] = NOT_GIVEN
+
+
+class LiteLLMClient(LLMClient[LiteLLMOptions]):
+    """
+    Client for the LiteLLM that supports calls to 100+ LLMs APIs, including OpenAI, Anthropic, VertexAI,
+    Hugging Face and others.
+    """
+
+    _options_cls = LiteLLMOptions
+
+    def __init__(
+        self,
+        model_name: str,
+        *,
+        base_url: Optional[str] = None,
+        api_key: Optional[str] = None,
+        api_version: Optional[str] = None,
+        use_structured_output: bool = False,
+    ) -> None:
+        """
+        Constructs a new LiteLLMClient instance.
+
+        Args:
+            model_name: Name of the model to use.
+            base_url: Base URL of the LLM API.
+            api_key: API key used to authenticate with the LLM API.
+            api_version: API version of the LLM API.
+            use_structured_output: Whether to request a structured output from the model. Default is False.
+
+        Raises:
+            ImportError: If the 'litellm' extra requirements are not installed.
+        """
+        if not HAS_LITELLM:
+            raise ImportError("You need to install the 'litellm' extra requirements to use LiteLLM models")
+
+        super().__init__(model_name)
+        self.base_url = base_url
+        self.api_key = api_key
+        self.api_version = api_version
+        self.use_structured_output = use_structured_output
+
+    async def call(
+        self,
+        conversation: ChatFormat,
+        options: LiteLLMOptions,
+        json_mode: bool = False,
+        output_schema: Optional[Type[BaseModel] | Dict] = None,
+    ) -> str:
+        """
+        Calls the appropriate LLM endpoint with the given prompt and options.
+
+        Args:
+            conversation: List of dicts with "role" and "content" keys, representing the chat history so far.
+            options: Additional settings used by the LLM.
+            json_mode: Force the response to be in JSON format.
+            output_schema: Output schema for requesting a specific response format.
+            Only used if the client has been initialized with `use_structured_output=True`.
+
+        Returns:
+            Response string from LLM.
+
+        Raises:
+            LLMConnectionError: If there is a connection error with the LLM API.
+            LLMStatusError: If the LLM API returns an error status code.
+            LLMResponseError: If the LLM API response is invalid.
+        """
+        supported_params = litellm.get_supported_openai_params(model=self.model_name)
+
+        response_format = None
+        if supported_params is not None and "response_format" in supported_params:
+            if output_schema is not None and self.use_structured_output:
+                response_format = output_schema
+            elif json_mode:
+                response_format = {"type": "json_object"}
+
+        try:
+            response = await litellm.acompletion(
+                messages=conversation,
+                model=self.model_name,
+                base_url=self.base_url,
+                api_key=self.api_key,
+                api_version=self.api_version,
+                response_format=response_format,
+                **options.dict(),
+            )
+        except litellm.openai.APIConnectionError as exc:
+            raise LLMConnectionError() from exc
+        except litellm.openai.APIStatusError as exc:
+            raise LLMStatusError(exc.message, exc.status_code) from exc
+        except litellm.openai.APIResponseValidationError as exc:
+            raise LLMResponseError() from exc
+
+        return response.choices[0].message.content