ragbits-core 0.0.1__tar.gz

ragbits_core-0.0.1/.gitignore

@@ -0,0 +1,97 @@
+# Directories
+.vscode/
+.idea/
+.neptune/
+.pytest_cache/
+.mypy_cache/
+venv/
+__pycache__/
+**.egg-info/
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+env/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Sphinx documentation
+docs/_build/
+public/
+# autogenerated package license table
+docs/licenses_table.rst
+
+# license dump file
+licenses.txt
+
+# File formats
+*.onnx
+*.pyc
+*.pt
+*.pth
+*.pkl
+*.mar
+*.torchscript
+**/.ipynb_checkpoints
+**/dist/
+**/checkpoints/
+**/outputs/
+**/multirun/
+
+# Other env files
+.python-version
+pyvenv.cfg
+pip-selfcheck.json
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*,cover
+.hypothesis/
+
+# dotenv
+.env
+
+# coverage and pytest reports
+coverage.xml
+report.xml
+
+# CMake
+cmake-build-*/
+
+# Terraform
+**/.terraform.lock.hcl
+**/.terraform
+
+# benchmarks
+benchmarks/sql/data/
+
+# mkdocs generated files
+site/
+
+# build artifacts
+dist/

ragbits_core-0.0.1/PKG-INFO

@@ -0,0 +1,30 @@
+Metadata-Version: 2.3
+Name: ragbits-core
+Version: 0.0.1
+Summary: Building blocks for rapid development of GenAI applications
+Author-email: "deepsense.ai" <contact@deepsense.ai>
+License-Expression: MIT
+Keywords: GenAI,Generative AI,LLMs,Large Language Models,Prompt Management,RAG,Retrieval Augmented Generation
+Classifier: Development Status :: 1 - Planning
+Classifier: Environment :: Console
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Natural Language :: English
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.10
+Requires-Dist: jinja2>=3.1.4
+Requires-Dist: pydantic>=2.9.1
+Provides-Extra: litellm
+Requires-Dist: litellm~=1.46.0; extra == 'litellm'
+Provides-Extra: local
+Requires-Dist: numpy~=1.24.0; extra == 'local'
+Requires-Dist: torch~=2.2.1; extra == 'local'
+Requires-Dist: transformers~=4.44.2; extra == 'local'
+Description-Content-Type: text/markdown
+
+# Ragbits Core

ragbits_core-0.0.1/README.md

@@ -0,0 +1 @@
+# Ragbits Core

ragbits_core-0.0.1/examples/llm_example.py

@@ -0,0 +1,62 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "ragbits[litellm]",
+# ]
+# ///
+import asyncio
+
+from pydantic import BaseModel
+
+from ragbits.core.llms.litellm import LiteLLM
+from ragbits.core.prompt import Prompt
+
+
+class LoremPromptInput(BaseModel):
+    """
+    Input format for the LoremPrompt.
+    """
+
+    theme: str
+    pun_allowed: bool = False
+
+
+class LoremPromptOutput(BaseModel):
+    """
+    Output format for the LoremPrompt.
+    """
+
+    joke: str
+    joke_category: str
+
+
+class JokePrompt(Prompt[LoremPromptInput, LoremPromptOutput]):
+    """
+    A prompt that generates jokes.
+    """
+
+    system_prompt = """
+    You are a joke generator. The jokes you generate should be funny and not offensive. {% if not pun_allowed %}Also, make sure
+    that the jokes do not contain any puns.{% else %}You can use any type of joke, even if it contains puns.{% endif %}
+
+    Respond as json with two fields: joke and joke_category.
+    """
+
+    user_prompt = """
+     theme: {{ theme }}
+    """
+
+
+async def main():
+    """
+    Example of using the LiteLLM client with a Prompt class. Requires the OPENAI_API_KEY environment variable to be set.
+    """
+    llm = LiteLLM("gpt-4o-2024-08-06", use_structured_output=True)
+    prompt = JokePrompt(LoremPromptInput(theme="software developers", pun_allowed=True))
+    response = await llm.generate(prompt)
+    print(f"The LLM generated a (hopefully) funny {response.joke_category} joke:")
+    print(response.joke)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

ragbits_core-0.0.1/examples/prompt_example.py

@@ -0,0 +1,52 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "ragbits",
+# ]
+# ///
+from pydantic import BaseModel
+
+from ragbits.core.prompt import Prompt
+
+
+class LoremPromptInput(BaseModel):
+    """
+    Input format for the LoremPrompt.
+    """
+
+    theme: str
+    nsfw_allowed: bool = False
+
+
+class LoremPromptOutput(BaseModel):
+    """
+    Output format for the LoremPrompt.
+    """
+
+    text: str
+
+
+class LoremPrompt(Prompt[LoremPromptInput, LoremPromptOutput]):
+    """
+    A prompt that generates Lorem Ipsum text.
+    """
+
+    system_prompt = """
+    You are a helpful Lorem Ipsum generator. The kind of vocablurary that you use besides "Lorem Ipsum" depends
+    on the theme provided by the user. Make sure it is latin and not too long. {% if not nsfw_allowed %}Also, make sure
+    that the text is safe for work.{% else %}You can use any text, even if it is not safe for work.{% endif %}
+    """
+
+    user_prompt = """
+     theme: {{ theme }}
+    """
+
+
+if __name__ == "__main__":
+    lorem_prompt = LoremPrompt(LoremPromptInput(theme="business"))
+    lorem_prompt.add_assistant_message("Lorem Ipsum biznessum dolor copy machinum yearly reportum")
+    print("CHAT:")
+    print(lorem_prompt.chat)
+    print()
+    print("OUTPUT MODEL:")
+    print(lorem_prompt.output_schema())

ragbits_core-0.0.1/pyproject.toml

@@ -0,0 +1,68 @@
+[project]
+name = "ragbits-core"
+version = "0.0.1"
+description = "Building blocks for rapid development of GenAI applications"
+readme = "README.md"
+requires-python = ">=3.10"
+license = "MIT"
+authors = [
+    { name = "deepsense.ai", email = "contact@deepsense.ai"}
+]
+keywords = [
+    "Retrieval Augmented Generation",
+    "RAG",
+    "Large Language Models",
+    "LLMs",
+    "Generative AI",
+    "GenAI",
+    "Prompt Management"
+]
+classifiers = [
+    "Development Status :: 1 - Planning",
+    "Environment :: Console",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Natural Language :: English",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Software Development :: Libraries :: Python Modules",
+]
+dependencies = [
+    "jinja2>=3.1.4",
+    "pydantic>=2.9.1"
+]
+
+[project.optional-dependencies]
+litellm = [
+    "litellm~=1.46.0",
+]
+local = [
+    "torch~=2.2.1",
+    "transformers~=4.44.2",
+    "numpy~=1.24.0"
+]
+
+[tool.uv]
+dev-dependencies = [
+    "pre-commit~=3.8.0",
+    "pytest~=8.3.3",
+    "pytest-cov~=5.0.0",
+    "pytest-asyncio~=0.24.0",
+    "pip-licenses>=4.0.0,<5.0.0"
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.metadata]
+allow-direct-references = true
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/ragbits"]
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"

ragbits_core-0.0.1/src/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-0.0.1/src/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-0.0.1/src/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-0.0.1/src/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-0.0.1/src/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-0.0.1/src/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-0.0.1/src/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",
+]