coco-framework 0.1.0__py3-none-any.whl

coco_core/documents/compressor.py

@@ -0,0 +1,74 @@
+"""Document compressor."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel
+
+from coco_core.runnables import run_in_executor
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from coco_core.callbacks import Callbacks
+    from coco_core.documents import Document
+
+
+class BaseDocumentCompressor(BaseModel, ABC):
+    """Base class for document compressors.
+
+    This abstraction is primarily used for post-processing of retrieved documents.
+
+    `Document` objects matching a given query are first retrieved.
+
+    Then the list of documents can be further processed.
+
+    For example, one could re-rank the retrieved documents using an LLM.
+
+    !!! note
+        Users should favor using a `RunnableLambda` instead of sub-classing from this
+        interface.
+
+    """
+
+    @abstractmethod
+    def compress_documents(
+        self,
+        documents: Sequence[Document],
+        query: str,
+        callbacks: Callbacks | None = None,
+    ) -> Sequence[Document]:
+        """Compress retrieved documents given the query context.
+
+        Args:
+            documents: The retrieved `Document` objects.
+            query: The query context.
+            callbacks: Optional `Callbacks` to run during compression.
+
+        Returns:
+            The compressed documents.
+
+        """
+
+    async def acompress_documents(
+        self,
+        documents: Sequence[Document],
+        query: str,
+        callbacks: Callbacks | None = None,
+    ) -> Sequence[Document]:
+        """Async compress retrieved documents given the query context.
+
+        Args:
+            documents: The retrieved `Document` objects.
+            query: The query context.
+            callbacks: Optional `Callbacks` to run during compression.
+
+        Returns:
+            The compressed documents.
+
+        """
+        return await run_in_executor(
+            None, self.compress_documents, documents, query, callbacks
+        )

coco_core/documents/transformers.py

@@ -0,0 +1,79 @@
+"""Document transformers."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
+
+from coco_core.runnables.config import run_in_executor
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from coco_core.documents import Document
+
+
+class BaseDocumentTransformer(ABC):
+    """Abstract base class for document transformation.
+
+    A document transformation takes a sequence of `Document` objects and returns a
+    sequence of transformed `Document` objects.
+
+    Example:
+        ```python
+        class EmbeddingsRedundantFilter(BaseDocumentTransformer, BaseModel):
+            embeddings: Embeddings
+            similarity_fn: Callable = cosine_similarity
+            similarity_threshold: float = 0.95
+
+            class Config:
+                arbitrary_types_allowed = True
+
+            def transform_documents(
+                self, documents: Sequence[Document], **kwargs: Any
+            ) -> Sequence[Document]:
+                stateful_documents = get_stateful_documents(documents)
+                embedded_documents = _get_embeddings_from_stateful_docs(
+                    self.embeddings, stateful_documents
+                )
+                included_idxs = _filter_similar_embeddings(
+                    embedded_documents,
+                    self.similarity_fn,
+                    self.similarity_threshold,
+                )
+                return [stateful_documents[i] for i in sorted(included_idxs)]
+
+            async def atransform_documents(
+                self, documents: Sequence[Document], **kwargs: Any
+            ) -> Sequence[Document]:
+                raise NotImplementedError
+        ```
+    """
+
+    @abstractmethod
+    def transform_documents(
+        self, documents: Sequence[Document], **kwargs: Any
+    ) -> Sequence[Document]:
+        """Transform a list of documents.
+
+        Args:
+            documents: A sequence of `Document` objects to be transformed.
+
+        Returns:
+            A sequence of transformed `Document` objects.
+        """
+
+    async def atransform_documents(
+        self, documents: Sequence[Document], **kwargs: Any
+    ) -> Sequence[Document]:
+        """Asynchronously transform a list of documents.
+
+        Args:
+            documents: A sequence of `Document` objects to be transformed.
+
+        Returns:
+            A sequence of transformed `Document` objects.
+        """
+        return await run_in_executor(
+            None, self.transform_documents, documents, **kwargs
+        )

coco_core/embeddings/__init__.py

@@ -0,0 +1,31 @@
+"""Embeddings."""
+
+from typing import TYPE_CHECKING
+
+from coco_core._import_utils import import_attr
+
+if TYPE_CHECKING:
+    from coco_core.embeddings.embeddings import Embeddings
+    from coco_core.embeddings.fake import (
+        DeterministicFakeEmbedding,
+        FakeEmbeddings,
+    )
+
+__all__ = ("DeterministicFakeEmbedding", "Embeddings", "FakeEmbeddings")
+
+_dynamic_imports = {
+    "Embeddings": "embeddings",
+    "DeterministicFakeEmbedding": "fake",
+    "FakeEmbeddings": "fake",
+}
+
+
+def __getattr__(attr_name: str) -> object:
+    module_name = _dynamic_imports.get(attr_name)
+    result = import_attr(attr_name, module_name, __spec__.parent)
+    globals()[attr_name] = result
+    return result
+
+
+def __dir__() -> list[str]:
+    return list(__all__)

coco_core/embeddings/embeddings.py

@@ -0,0 +1,78 @@
+"""**Embeddings** interface."""
+
+from abc import ABC, abstractmethod
+
+from coco_core.runnables.config import run_in_executor
+
+
+class Embeddings(ABC):
+    """Interface for embedding models.
+
+    This is an interface meant for implementing text embedding models.
+
+    Text embedding models are used to map text to a vector (a point in n-dimensional
+    space).
+
+    Texts that are similar will usually be mapped to points that are close to each
+    other in this space. The exact details of what's considered "similar" and how
+    "distance" is measured in this space are dependent on the specific embedding model.
+
+    This abstraction contains a method for embedding a list of documents and a method
+    for embedding a query text. The embedding of a query text is expected to be a single
+    vector, while the embedding of a list of documents is expected to be a list of
+    vectors.
+
+    Usually the query embedding is identical to the document embedding, but the
+    abstraction allows treating them independently.
+
+    In addition to the synchronous methods, this interface also provides asynchronous
+    versions of the methods.
+
+    By default, the asynchronous methods are implemented using the synchronous methods;
+    however, implementations may choose to override the asynchronous methods with
+    an async native implementation for performance reasons.
+    """
+
+    @abstractmethod
+    def embed_documents(self, texts: list[str]) -> list[list[float]]:
+        """Embed search docs.
+
+        Args:
+            texts: List of text to embed.
+
+        Returns:
+            List of embeddings.
+        """
+
+    @abstractmethod
+    def embed_query(self, text: str) -> list[float]:
+        """Embed query text.
+
+        Args:
+            text: Text to embed.
+
+        Returns:
+            Embedding.
+        """
+
+    async def aembed_documents(self, texts: list[str]) -> list[list[float]]:
+        """Asynchronous Embed search docs.
+
+        Args:
+            texts: List of text to embed.
+
+        Returns:
+            List of embeddings.
+        """
+        return await run_in_executor(None, self.embed_documents, texts)
+
+    async def aembed_query(self, text: str) -> list[float]:
+        """Asynchronous Embed query text.
+
+        Args:
+            text: Text to embed.
+
+        Returns:
+            Embedding.
+        """
+        return await run_in_executor(None, self.embed_query, text)

coco_core/embeddings/fake.py

@@ -0,0 +1,129 @@
+"""Module contains a few fake embedding models for testing purposes."""
+
+# Please do not add additional fake embedding model implementations here.
+import contextlib
+import hashlib
+
+from pydantic import BaseModel
+from typing_extensions import override
+
+from coco_core.embeddings import Embeddings
+
+with contextlib.suppress(ImportError):
+    import numpy as np
+
+
+class FakeEmbeddings(Embeddings, BaseModel):
+    """Fake embedding model for unit testing purposes.
+
+    This embedding model creates embeddings by sampling from a normal distribution.
+
+    !!! danger "Toy model"
+        Do not use this outside of testing, as it is not a real embedding model.
+
+    Instantiate:
+        ```python
+        from coco_core.embeddings import FakeEmbeddings
+
+        embed = FakeEmbeddings(size=100)
+        ```
+
+    Embed single text:
+        ```python
+        input_text = "The meaning of life is 42"
+        vector = embed.embed_query(input_text)
+        print(vector[:3])
+        ```
+        ```python
+        [-0.700234640213188, -0.581266257710429, -1.1328482266445354]
+        ```
+
+    Embed multiple texts:
+        ```python
+        input_texts = ["Document 1...", "Document 2..."]
+        vectors = embed.embed_documents(input_texts)
+        print(len(vectors))
+        # The first 3 coordinates for the first vector
+        print(vectors[0][:3])
+        ```
+        ```python
+        2
+        [-0.5670477847544458, -0.31403828652395727, -0.5840547508955257]
+        ```
+    """
+
+    size: int
+    """The size of the embedding vector."""
+
+    def _get_embedding(self) -> list[float]:
+        return list(np.random.default_rng().normal(size=self.size))
+
+    @override
+    def embed_documents(self, texts: list[str]) -> list[list[float]]:
+        return [self._get_embedding() for _ in texts]
+
+    @override
+    def embed_query(self, text: str) -> list[float]:
+        return self._get_embedding()
+
+
+class DeterministicFakeEmbedding(Embeddings, BaseModel):
+    """Deterministic fake embedding model for unit testing purposes.
+
+    This embedding model creates embeddings by sampling from a normal distribution
+    with a seed based on the hash of the text.
+
+    !!! danger "Toy model"
+        Do not use this outside of testing, as it is not a real embedding model.
+
+    Instantiate:
+        ```python
+        from coco_core.embeddings import DeterministicFakeEmbedding
+
+        embed = DeterministicFakeEmbedding(size=100)
+        ```
+
+    Embed single text:
+        ```python
+        input_text = "The meaning of life is 42"
+        vector = embed.embed_query(input_text)
+        print(vector[:3])
+        ```
+        ```python
+        [-0.700234640213188, -0.581266257710429, -1.1328482266445354]
+        ```
+
+    Embed multiple texts:
+        ```python
+        input_texts = ["Document 1...", "Document 2..."]
+        vectors = embed.embed_documents(input_texts)
+        print(len(vectors))
+        # The first 3 coordinates for the first vector
+        print(vectors[0][:3])
+        ```
+        ```python
+        2
+        [-0.5670477847544458, -0.31403828652395727, -0.5840547508955257]
+        ```
+    """
+
+    size: int
+    """The size of the embedding vector."""
+
+    def _get_embedding(self, seed: int) -> list[float]:
+        # set the seed for the random generator
+        rng = np.random.default_rng(seed)
+        return list(rng.normal(size=self.size))
+
+    @staticmethod
+    def _get_seed(text: str) -> int:
+        """Get a seed for the random generator, using the hash of the text."""
+        return int(hashlib.sha256(text.encode("utf-8")).hexdigest(), 16) % 10**8
+
+    @override
+    def embed_documents(self, texts: list[str]) -> list[list[float]]:
+        return [self._get_embedding(seed=self._get_seed(_)) for _ in texts]
+
+    @override
+    def embed_query(self, text: str) -> list[float]:
+        return self._get_embedding(seed=self._get_seed(text))

coco_core/env.py

@@ -0,0 +1,22 @@
+"""Utilities for getting information about the runtime environment."""
+
+import platform
+from functools import lru_cache
+
+from coco_core import __version__
+
+
+@lru_cache(maxsize=1)
+def get_runtime_environment() -> dict[str, str]:
+    """Get information about the Coco runtime environment.
+
+    Returns:
+        A dictionary with information about the runtime environment.
+    """
+    return {
+        "library_version": __version__,
+        "library": "trustcoco-core",
+        "platform": platform.platform(),
+        "runtime": "python",
+        "runtime_version": platform.python_version(),
+    }

coco_core/example_selectors/__init__.py

@@ -0,0 +1,47 @@
+"""Example selectors.
+
+**Example selector** implements logic for selecting examples to include them in prompts.
+This allows us to select examples that are most relevant to the input.
+"""
+
+from typing import TYPE_CHECKING
+
+from coco_core._import_utils import import_attr
+
+if TYPE_CHECKING:
+    from coco_core.example_selectors.base import BaseExampleSelector
+    from coco_core.example_selectors.length_based import (
+        LengthBasedExampleSelector,
+    )
+    from coco_core.example_selectors.semantic_similarity import (
+        MaxMarginalRelevanceExampleSelector,
+        SemanticSimilarityExampleSelector,
+        sorted_values,
+    )
+
+__all__ = (
+    "BaseExampleSelector",
+    "LengthBasedExampleSelector",
+    "MaxMarginalRelevanceExampleSelector",
+    "SemanticSimilarityExampleSelector",
+    "sorted_values",
+)
+
+_dynamic_imports = {
+    "BaseExampleSelector": "base",
+    "LengthBasedExampleSelector": "length_based",
+    "MaxMarginalRelevanceExampleSelector": "semantic_similarity",
+    "SemanticSimilarityExampleSelector": "semantic_similarity",
+    "sorted_values": "semantic_similarity",
+}
+
+
+def __getattr__(attr_name: str) -> object:
+    module_name = _dynamic_imports.get(attr_name)
+    result = import_attr(attr_name, module_name, __spec__.parent)
+    globals()[attr_name] = result
+    return result
+
+
+def __dir__() -> list[str]:
+    return list(__all__)

coco_core/example_selectors/base.py

@@ -0,0 +1,60 @@
+"""Interface for selecting examples to include in prompts."""
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+from coco_core.runnables import run_in_executor
+
+
+class BaseExampleSelector(ABC):
+    """Interface for selecting examples to include in prompts."""
+
+    @abstractmethod
+    def add_example(self, example: dict[str, str]) -> Any:
+        """Add new example to store.
+
+        Args:
+            example: A dictionary with keys as input variables
+                and values as their values.
+
+        Returns:
+            Any return value.
+        """
+
+    async def aadd_example(self, example: dict[str, str]) -> Any:
+        """Async add new example to store.
+
+        Args:
+            example: A dictionary with keys as input variables
+                and values as their values.
+
+        Returns:
+            Any return value.
+        """
+        return await run_in_executor(None, self.add_example, example)
+
+    @abstractmethod
+    def select_examples(self, input_variables: dict[str, str]) -> list[dict[str, Any]]:
+        """Select which examples to use based on the inputs.
+
+        Args:
+            input_variables: A dictionary with keys as input variables
+                and values as their values.
+
+        Returns:
+            A list of examples.
+        """
+
+    async def aselect_examples(
+        self, input_variables: dict[str, str]
+    ) -> list[dict[str, Any]]:
+        """Async select which examples to use based on the inputs.
+
+        Args:
+            input_variables: A dictionary with keys as input variables
+                and values as their values.
+
+        Returns:
+            A list of examples.
+        """
+        return await run_in_executor(None, self.select_examples, input_variables)

coco_core/example_selectors/length_based.py

@@ -0,0 +1,131 @@
+"""Select examples based on length."""
+
+import re
+from collections.abc import Callable
+from typing import Any
+
+from pydantic import BaseModel, Field, model_validator
+from typing_extensions import Self
+
+from coco_core.example_selectors.base import BaseExampleSelector
+from coco_core.prompts.prompt import PromptTemplate
+
+
+def _get_length_based(text: str) -> int:
+    return len(re.split(r"\n| ", text))
+
+
+class LengthBasedExampleSelector(BaseExampleSelector, BaseModel):
+    r"""Select examples based on length.
+
+    Example:
+        ```python
+        from coco_core.example_selectors import LengthBasedExampleSelector
+        from coco_core.prompts import PromptTemplate
+
+        # Define examples
+        examples = [
+            {"input": "happy", "output": "sad"},
+            {"input": "tall", "output": "short"},
+            {"input": "fast", "output": "slow"},
+        ]
+
+        # Create prompt template
+        example_prompt = PromptTemplate(
+            input_variables=["input", "output"],
+            template="Input: {input}\nOutput: {output}",
+        )
+
+        # Create selector with max length constraint
+        selector = LengthBasedExampleSelector(
+            examples=examples,
+            example_prompt=example_prompt,
+            max_length=50,  # Maximum prompt length
+        )
+
+        # Select examples for a new input
+        selected = selector.select_examples({"input": "large", "output": "tiny"})
+        # Returns examples that fit within max_length constraint
+        ```
+    """
+
+    examples: list[dict[str, Any]]
+    """A list of the examples that the prompt template expects."""
+
+    example_prompt: PromptTemplate
+    """Prompt template used to format the examples."""
+
+    get_text_length: Callable[[str], int] = _get_length_based
+    """Function to measure prompt length. Defaults to word count."""
+
+    max_length: int = 2048
+    """Max length for the prompt, beyond which examples are cut."""
+
+    example_text_lengths: list[int] = Field(default_factory=list)
+    """Length of each example."""
+
+    def add_example(self, example: dict[str, str]) -> None:
+        """Add new example to list.
+
+        Args:
+            example: A dictionary with keys as input variables
+                and values as their values.
+        """
+        self.examples.append(example)
+        string_example = self.example_prompt.format(**example)
+        self.example_text_lengths.append(self.get_text_length(string_example))
+
+    async def aadd_example(self, example: dict[str, str]) -> None:
+        """Async add new example to list.
+
+        Args:
+            example: A dictionary with keys as input variables
+                and values as their values.
+        """
+        self.add_example(example)
+
+    @model_validator(mode="after")
+    def post_init(self) -> Self:
+        """Validate that the examples are formatted correctly."""
+        if self.example_text_lengths:
+            return self
+        string_examples = [self.example_prompt.format(**eg) for eg in self.examples]
+        self.example_text_lengths = [self.get_text_length(eg) for eg in string_examples]
+        return self
+
+    def select_examples(self, input_variables: dict[str, str]) -> list[dict[str, Any]]:
+        """Select which examples to use based on the input lengths.
+
+        Args:
+            input_variables: A dictionary with keys as input variables
+                and values as their values.
+
+        Returns:
+            A list of examples to include in the prompt.
+        """
+        inputs = " ".join(input_variables.values())
+        remaining_length = self.max_length - self.get_text_length(inputs)
+        i = 0
+        examples = []
+        while remaining_length > 0 and i < len(self.examples):
+            new_length = remaining_length - self.example_text_lengths[i]
+            if new_length < 0:
+                break
+            examples.append(self.examples[i])
+            remaining_length = new_length
+            i += 1
+        return examples
+
+    async def aselect_examples(
+        self, input_variables: dict[str, str]
+    ) -> list[dict[str, Any]]:
+        """Async select which examples to use based on the input lengths.
+
+        Args:
+            input_variables: A dictionary with keys as input variables
+                and values as their values.
+
+        Returns:
+            A list of examples to include in the prompt.
+        """
+        return self.select_examples(input_variables)