coco-framework 0.1.0__py3-none-any.whl

coco_core/chat_sessions.py

@@ -0,0 +1,19 @@
+"""**Chat Sessions** are a collection of messages and function calls."""
+
+from collections.abc import Sequence
+from typing import Any, TypedDict
+
+from coco_core.messages import BaseMessage
+
+
+class ChatSession(TypedDict, total=False):
+    """Chat Session.
+
+    Chat Session represents a single conversation, channel, or other group of messages.
+    """
+
+    messages: Sequence[BaseMessage]
+    """A sequence of the Coco chat messages loaded from the source."""
+
+    functions: Sequence[dict[str, Any]]
+    """A sequence of the function calling specs for the messages."""

coco_core/cross_encoders.py

@@ -0,0 +1,18 @@
+"""Cross Encoder interface."""
+
+from abc import ABC, abstractmethod
+
+
+class BaseCrossEncoder(ABC):
+    """Interface for cross encoder models."""
+
+    @abstractmethod
+    def score(self, text_pairs: list[tuple[str, str]]) -> list[float]:
+        """Score pairs' similarity.
+
+        Args:
+            text_pairs: List of pairs of texts.
+
+        Returns:
+            List of scores.
+        """

coco_core/document_loaders/__init__.py

@@ -0,0 +1,36 @@
+"""Document loaders."""
+
+from typing import TYPE_CHECKING
+
+from coco_core._import_utils import import_attr
+
+if TYPE_CHECKING:
+    from coco_core.document_loaders.base import BaseBlobParser, BaseLoader
+    from coco_core.document_loaders.blob_loaders import Blob, BlobLoader, PathLike
+
+__all__ = (
+    "BaseBlobParser",
+    "BaseLoader",
+    "Blob",
+    "BlobLoader",
+    "PathLike",
+)
+
+_dynamic_imports = {
+    "BaseBlobParser": "base",
+    "BaseLoader": "base",
+    "Blob": "blob_loaders",
+    "BlobLoader": "blob_loaders",
+    "PathLike": "blob_loaders",
+}
+
+
+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/document_loaders/base.py

@@ -0,0 +1,155 @@
+"""Abstract interface for document loader implementations."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+from coco_core.runnables import run_in_executor
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator, Iterator
+
+    from coco_text_splitters import TextSplitter
+
+    from coco_core.documents import Document
+    from coco_core.documents.base import Blob
+
+try:
+    from coco_text_splitters import RecursiveCharacterTextSplitter
+
+    _HAS_TEXT_SPLITTERS = True
+except ImportError:
+    _HAS_TEXT_SPLITTERS = False
+
+
+class BaseLoader(ABC):  # noqa: B024
+    """Interface for document loader.
+
+    Implementations should implement the lazy-loading method using generators to avoid
+    loading all documents into memory at once.
+
+    `load` is provided just for user convenience and should not be overridden.
+    """
+
+    # Sub-classes should not implement this method directly. Instead, they
+    # should implement the lazy load method.
+    def load(self) -> list[Document]:
+        """Load data into `Document` objects.
+
+        Returns:
+            The documents.
+        """
+        return list(self.lazy_load())
+
+    async def aload(self) -> list[Document]:
+        """Load data into `Document` objects.
+
+        Returns:
+            The documents.
+        """
+        return [document async for document in self.alazy_load()]
+
+    def load_and_split(
+        self, text_splitter: TextSplitter | None = None
+    ) -> list[Document]:
+        """Load `Document` and split into chunks. Chunks are returned as `Document`.
+
+        !!! danger
+
+            Do not override this method. It should be considered to be deprecated!
+
+        Args:
+            text_splitter: `TextSplitter` instance to use for splitting documents.
+
+                Defaults to `RecursiveCharacterTextSplitter`.
+
+        Raises:
+            ImportError: If `trustcoco-text-splitters` is not installed and no
+                `text_splitter` is provided.
+
+        Returns:
+            List of `Document` objects.
+        """
+        if text_splitter is None:
+            if not _HAS_TEXT_SPLITTERS:
+                msg = (
+                    "Unable to import from coco_text_splitters. Please specify "
+                    "text_splitter or install coco_text_splitters with "
+                    "`pip install -U trustcoco-text-splitters`."
+                )
+                raise ImportError(msg)
+
+            text_splitter_: TextSplitter = RecursiveCharacterTextSplitter()
+        else:
+            text_splitter_ = text_splitter
+        docs = self.load()
+        return text_splitter_.split_documents(docs)
+
+    # Attention: This method will be upgraded into an abstractmethod once it's
+    #            implemented in all the existing subclasses.
+    def lazy_load(self) -> Iterator[Document]:
+        """A lazy loader for `Document`.
+
+        Yields:
+            The `Document` objects.
+        """
+        if type(self).load != BaseLoader.load:
+            return iter(self.load())
+        msg = f"{self.__class__.__name__} does not implement lazy_load()"
+        raise NotImplementedError(msg)
+
+    async def alazy_load(self) -> AsyncIterator[Document]:
+        """A lazy loader for `Document`.
+
+        Yields:
+            The `Document` objects.
+        """
+        iterator = await run_in_executor(None, self.lazy_load)
+        done = object()
+        while True:
+            doc = await run_in_executor(None, next, iterator, done)
+            if doc is done:
+                break
+            yield doc  # type: ignore[misc]
+
+
+class BaseBlobParser(ABC):
+    """Abstract interface for blob parsers.
+
+    A blob parser provides a way to parse raw data stored in a blob into one or more
+    `Document` objects.
+
+    The parser can be composed with blob loaders, making it easy to reuse a parser
+    independent of how the blob was originally loaded.
+    """
+
+    @abstractmethod
+    def lazy_parse(self, blob: Blob) -> Iterator[Document]:
+        """Lazy parsing interface.
+
+        Subclasses are required to implement this method.
+
+        Args:
+            blob: `Blob` instance
+
+        Returns:
+            Generator of `Document` objects
+        """
+
+    def parse(self, blob: Blob) -> list[Document]:
+        """Eagerly parse the blob into a `Document` or list of `Document` objects.
+
+        This is a convenience method for interactive development environment.
+
+        Production applications should favor the `lazy_parse` method instead.
+
+        Subclasses should generally not over-ride this parse method.
+
+        Args:
+            blob: `Blob` instance
+
+        Returns:
+            List of `Document` objects
+        """
+        return list(self.lazy_parse(blob))

coco_core/document_loaders/blob_loaders.py

@@ -0,0 +1,38 @@
+"""Schema for Blobs and Blob Loaders.
+
+The goal is to facilitate decoupling of content loading from content parsing code. In
+addition, content loading code should provide a lazy loading interface by default.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+# Re-export Blob and PathLike for backwards compatibility
+from coco_core.documents.base import Blob, PathLike
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+
+class BlobLoader(ABC):
+    """Abstract interface for blob loaders implementation.
+
+    Implementer should be able to load raw content from a storage system according to
+    some criteria and return the raw content lazily as a stream of blobs.
+    """
+
+    @abstractmethod
+    def yield_blobs(
+        self,
+    ) -> Iterator[Blob]:
+        """A lazy loader for raw data represented by Coco's `Blob` object.
+
+        Yields:
+            `Blob` objects.
+        """
+
+
+# Re-export Blob and Pathlike for backwards compatibility
+__all__ = ["Blob", "BlobLoader", "PathLike"]

coco_core/documents/__init__.py

@@ -0,0 +1,55 @@
+"""Documents module for data retrieval and processing workflows.
+
+This module provides core abstractions for handling data in retrieval-augmented
+generation (RAG) pipelines, vector stores, and document processing workflows.
+
+!!! warning "Documents vs. message content"
+
+    This module is distinct from `coco_core.messages.content`, which provides
+    multimodal content blocks for **LLM chat I/O** (text, images, audio, etc. within
+    messages).
+
+    **Key distinction:**
+
+    - **Documents** (this module): For **data retrieval and processing workflows**
+        - Vector stores, retrievers, RAG pipelines
+        - Text chunking, embedding, and semantic search
+        - Example: Chunks of a PDF stored in a vector database
+
+    - **Content Blocks** (`messages.content`): For **LLM conversational I/O**
+        - Multimodal message content sent to/from models
+        - Tool calls, reasoning, citations within chat
+        - Example: An image sent to a vision model in a chat message (via
+            [`ImageContentBlock`][coco.messages.ImageContentBlock])
+
+    While both can represent similar data types (text, files), they serve different
+    architectural purposes in Coco applications.
+"""
+
+from typing import TYPE_CHECKING
+
+from coco_core._import_utils import import_attr
+
+if TYPE_CHECKING:
+    from coco_core.documents.base import Document
+    from coco_core.documents.compressor import BaseDocumentCompressor
+    from coco_core.documents.transformers import BaseDocumentTransformer
+
+__all__ = ("BaseDocumentCompressor", "BaseDocumentTransformer", "Document")
+
+_dynamic_imports = {
+    "Document": "base",
+    "BaseDocumentCompressor": "compressor",
+    "BaseDocumentTransformer": "transformers",
+}
+
+
+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/documents/base.py

@@ -0,0 +1,347 @@
+"""Base classes for media and documents.
+
+This module contains core abstractions for **data retrieval and processing workflows**:
+
+- `BaseMedia`: Base class providing `id` and `metadata` fields
+- `Blob`: Raw data loading (files, binary data) - used by document loaders
+- `Document`: Text content for retrieval (RAG, vector stores, semantic search)
+
+!!! note "Not for LLM chat messages"
+
+    These classes are for data processing pipelines, not LLM I/O. For multimodal
+    content in chat messages (images, audio in conversations), see
+    `coco.messages` content blocks instead.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import mimetypes
+from io import BufferedReader, BytesIO
+from pathlib import Path, PurePath
+from typing import TYPE_CHECKING, Any, Literal, cast
+
+from pydantic import ConfigDict, Field, model_validator
+
+from coco_core.load.serializable import Serializable
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+
+PathLike = str | PurePath
+
+
+class BaseMedia(Serializable):
+    """Base class for content used in retrieval and data processing workflows.
+
+    Provides common fields for content that needs to be stored, indexed, or searched.
+
+    !!! note
+
+        For multimodal content in **chat messages** (images, audio sent to/from LLMs),
+        use `coco.messages` content blocks instead.
+    """
+
+    # The ID field is optional at the moment.
+    # It will likely become required in a future major release after
+    # it has been adopted by enough VectorStore implementations.
+    id: str | None = Field(default=None, coerce_numbers_to_str=True)
+    """An optional identifier for the document.
+
+    Ideally this should be unique across the document collection and formatted
+    as a UUID, but this will not be enforced.
+    """
+
+    metadata: dict[Any, Any] = Field(default_factory=dict)
+    """Arbitrary metadata associated with the content."""
+
+
+class Blob(BaseMedia):
+    """Raw data abstraction for document loading and file processing.
+
+    Represents raw bytes or text, either in-memory or by file reference. Used
+    primarily by document loaders to decouple data loading from parsing.
+
+    Inspired by [Mozilla's `Blob`](https://developer.mozilla.org/en-US/docs/Web/API/Blob)
+
+    ???+ example "Initialize a blob from in-memory data"
+
+        ```python
+        from coco_core.documents import Blob
+
+        blob = Blob.from_data("Hello, world!")
+
+        # Read the blob as a string
+        print(blob.as_string())
+
+        # Read the blob as bytes
+        print(blob.as_bytes())
+
+        # Read the blob as a byte stream
+        with blob.as_bytes_io() as f:
+            print(f.read())
+        ```
+
+    ??? example "Load from memory and specify MIME type and metadata"
+
+        ```python
+        from coco_core.documents import Blob
+
+        blob = Blob.from_data(
+            data="Hello, world!",
+            mime_type="text/plain",
+            metadata={"source": "https://example.com"},
+        )
+        ```
+
+    ??? example "Load the blob from a file"
+
+        ```python
+        from coco_core.documents import Blob
+
+        blob = Blob.from_path("path/to/file.txt")
+
+        # Read the blob as a string
+        print(blob.as_string())
+
+        # Read the blob as bytes
+        print(blob.as_bytes())
+
+        # Read the blob as a byte stream
+        with blob.as_bytes_io() as f:
+            print(f.read())
+        ```
+    """
+
+    data: bytes | str | None = None
+    """Raw data associated with the `Blob`."""
+
+    mimetype: str | None = None
+    """MIME type, not to be confused with a file extension."""
+
+    encoding: str = "utf-8"
+    """Encoding to use if decoding the bytes into a string.
+
+    Uses `utf-8` as default encoding if decoding to string.
+    """
+
+    path: PathLike | None = None
+    """Location where the original content was found."""
+
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        frozen=True,
+    )
+
+    @property
+    def source(self) -> str | None:
+        """The source location of the blob as string if known otherwise none.
+
+        If a path is associated with the `Blob`, it will default to the path location.
+
+        Unless explicitly set via a metadata field called `'source'`, in which
+        case that value will be used instead.
+        """
+        if self.metadata and "source" in self.metadata:
+            return cast("str | None", self.metadata["source"])
+        return str(self.path) if self.path else None
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_blob_is_valid(cls, values: dict[str, Any]) -> Any:
+        """Verify that either data or path is provided."""
+        if "data" not in values and "path" not in values:
+            msg = "Either data or path must be provided"
+            raise ValueError(msg)
+        return values
+
+    def as_string(self) -> str:
+        """Read data as a string.
+
+        Raises:
+            ValueError: If the blob cannot be represented as a string.
+
+        Returns:
+            The data as a string.
+        """
+        if self.data is None and self.path:
+            return Path(self.path).read_text(encoding=self.encoding)
+        if isinstance(self.data, bytes):
+            return self.data.decode(self.encoding)
+        if isinstance(self.data, str):
+            return self.data
+        msg = f"Unable to get string for blob {self}"
+        raise ValueError(msg)
+
+    def as_bytes(self) -> bytes:
+        """Read data as bytes.
+
+        Raises:
+            ValueError: If the blob cannot be represented as bytes.
+
+        Returns:
+            The data as bytes.
+        """
+        if isinstance(self.data, bytes):
+            return self.data
+        if isinstance(self.data, str):
+            return self.data.encode(self.encoding)
+        if self.data is None and self.path:
+            return Path(self.path).read_bytes()
+        msg = f"Unable to get bytes for blob {self}"
+        raise ValueError(msg)
+
+    @contextlib.contextmanager
+    def as_bytes_io(self) -> Generator[BytesIO | BufferedReader, None, None]:
+        """Read data as a byte stream.
+
+        Raises:
+            NotImplementedError: If the blob cannot be represented as a byte stream.
+
+        Yields:
+            The data as a byte stream.
+        """
+        if isinstance(self.data, bytes):
+            yield BytesIO(self.data)
+        elif self.data is None and self.path:
+            with Path(self.path).open("rb") as f:
+                yield f
+        else:
+            msg = f"Unable to convert blob {self}"
+            raise NotImplementedError(msg)
+
+    @classmethod
+    def from_path(
+        cls,
+        path: PathLike,
+        *,
+        encoding: str = "utf-8",
+        mime_type: str | None = None,
+        guess_type: bool = True,
+        metadata: dict[Any, Any] | None = None,
+    ) -> Blob:
+        """Load the blob from a path like object.
+
+        Args:
+            path: Path-like object to file to be read
+            encoding: Encoding to use if decoding the bytes into a string
+            mime_type: If provided, will be set as the MIME type of the data
+            guess_type: If `True`, the MIME type will be guessed from the file
+                extension, if a MIME type was not provided
+            metadata: Metadata to associate with the `Blob`
+
+        Returns:
+            `Blob` instance
+        """
+        if mime_type is None and guess_type:
+            mimetype = mimetypes.guess_type(path)[0]
+        else:
+            mimetype = mime_type
+        # We do not load the data immediately, instead we treat the blob as a
+        # reference to the underlying data.
+        return cls(
+            data=None,
+            mimetype=mimetype,
+            encoding=encoding,
+            path=path,
+            metadata=metadata if metadata is not None else {},
+        )
+
+    @classmethod
+    def from_data(
+        cls,
+        data: str | bytes,
+        *,
+        encoding: str = "utf-8",
+        mime_type: str | None = None,
+        path: str | None = None,
+        metadata: dict[Any, Any] | None = None,
+    ) -> Blob:
+        """Initialize the `Blob` from in-memory data.
+
+        Args:
+            data: The in-memory data associated with the `Blob`
+            encoding: Encoding to use if decoding the bytes into a string
+            mime_type: If provided, will be set as the MIME type of the data
+            path: If provided, will be set as the source from which the data came
+            metadata: Metadata to associate with the `Blob`
+
+        Returns:
+            `Blob` instance
+        """
+        return cls(
+            data=data,
+            mimetype=mime_type,
+            encoding=encoding,
+            path=path,
+            metadata=metadata if metadata is not None else {},
+        )
+
+    def __repr__(self) -> str:
+        """Return the blob representation."""
+        str_repr = f"Blob {id(self)}"
+        if self.source:
+            str_repr += f" {self.source}"
+        return str_repr
+
+
+class Document(BaseMedia):
+    """Class for storing a piece of text and associated metadata.
+
+    !!! note
+
+        `Document` is for **retrieval workflows**, not chat I/O. For sending text
+        to an LLM in a conversation, use message types from `coco.messages`.
+
+    Example:
+        ```python
+        from coco_core.documents import Document
+
+        document = Document(
+            page_content="Hello, world!", metadata={"source": "https://example.com"}
+        )
+        ```
+    """
+
+    page_content: str
+    """String text."""
+
+    type: Literal["Document"] = "Document"
+
+    def __init__(self, page_content: str, **kwargs: Any) -> None:
+        """Pass page_content in as positional or named arg."""
+        # my-py is complaining that page_content is not defined on the base class.
+        # Here, we're relying on pydantic base class to handle the validation.
+        super().__init__(page_content=page_content, **kwargs)  # type: ignore[call-arg,unused-ignore]
+
+    @classmethod
+    def is_lc_serializable(cls) -> bool:
+        """Return `True` as this class is serializable."""
+        return True
+
+    @classmethod
+    def get_lc_namespace(cls) -> list[str]:
+        """Get the namespace of the Coco object.
+
+        Returns:
+            `["coco", "schema", "document"]`
+        """
+        return ["coco", "schema", "document"]
+
+    def __str__(self) -> str:
+        """Override `__str__` to restrict it to page_content and metadata.
+
+        Returns:
+            A string representation of the `Document`.
+        """
+        # The format matches pydantic format for __str__.
+        #
+        # The purpose of this change is to make sure that user code that feeds
+        # Document objects directly into prompts remains unchanged due to the addition
+        # of the id field (or any other fields in the future).
+        #
+        # This override will likely be removed in the future in favor of a more general
+        # solution of formatting content directly inside the prompts.
+        if self.metadata:
+            return f"page_content='{self.page_content}' metadata={self.metadata}"
+        return f"page_content='{self.page_content}'"