ragbits-document-search 1.4.0.dev202601310254__py3-none-any.whl

ragbits/document_search/documents/element.py

@@ -0,0 +1,208 @@
+import hashlib
+import uuid
+from abc import ABC, abstractmethod
+from typing import Any, ClassVar
+
+from pydantic import BaseModel, computed_field
+
+from ragbits.core.utils.pydantic import SerializableBytes
+from ragbits.core.vector_stores.base import VectorStoreEntry
+from ragbits.document_search.documents.document import DocumentMeta
+
+
+class ElementLocation(BaseModel):
+    """
+    An object representing position of chunk within document.
+    """
+
+    page_number: int | None = None
+    coordinates: dict | None = None
+
+
+class Element(BaseModel, ABC):
+    """
+    An object representing an element in a document.
+    """
+
+    element_type: str
+    document_meta: DocumentMeta
+    location: ElementLocation | None = None
+    score: float | None = None
+
+    _elements_registry: ClassVar[dict[str, type["Element"]]] = {}
+
+    @property
+    def id(self) -> str:
+        """
+        Retrieve the ID of the element, primarily used to represent the element's data.
+
+        Returns:
+            str: string representing element
+        """
+        id_components = self.get_id_components()
+        return "&".join(f"{k}={v}" for k, v in id_components.items())
+
+    def get_id_components(self) -> dict[str, str]:
+        """
+        Creates a dictionary of key value pairs of id components
+
+        Returns:
+            dict: a dictionary
+        """
+        id_components = {
+            "meta": self.document_meta.id,
+            "type": self.element_type,
+            "key": str(self.key),
+            "text": str(self.text_representation),
+            "location": str(self.location),
+        }
+        return id_components
+
+    @computed_field  # type: ignore[prop-decorator]
+    @property
+    def key(self) -> str | None:
+        """
+        Get the representation of the element for embedding.
+
+        Returns:
+            The representation for embedding.
+        """
+        return self.text_representation
+
+    @computed_field  # type: ignore[prop-decorator]
+    @property
+    @abstractmethod
+    def text_representation(self) -> str | None:
+        """
+        Get the text representation of the element.
+
+        Returns:
+            The text representation.
+        """
+
+    @property
+    def image_representation(self) -> bytes | None:
+        """
+        Get the image representation of the element.
+
+        Returns:
+            The image representation.
+        """
+        return None
+
+    @classmethod
+    def __pydantic_init_subclass__(cls, **kwargs: Any) -> None:  # noqa: ANN401
+        element_type_default = cls.model_fields["element_type"].default
+        if element_type_default is None:
+            raise ValueError("Element type must be defined")
+        Element._elements_registry[element_type_default] = cls
+
+    @classmethod
+    def from_vector_db_entry(cls, db_entry: VectorStoreEntry, score: float | None = None) -> "Element":
+        """
+        Create an element from a vector database entry.
+
+        Args:
+            db_entry: The vector database entry.
+            score: The score of the element retrieved from the vector database or reranker.
+
+        Returns:
+            The element.
+        """
+        element_type = db_entry.metadata["element_type"]
+        element_cls = Element._elements_registry[element_type]
+        if "embedding_type" in db_entry.metadata:
+            del db_entry.metadata["embedding_type"]
+
+        element = element_cls(**db_entry.metadata)
+        element.score = score
+        return element
+
+    def to_vector_db_entry(self) -> VectorStoreEntry:
+        """
+        Create a vector database entry from the element.
+
+        Returns:
+            The vector database entry
+        """
+        id_components = [
+            self.id,
+        ]
+        vector_store_entry_id = uuid.uuid5(uuid.NAMESPACE_OID, ";".join(id_components))
+        metadata = self.model_dump(exclude={"id", "key"})
+        metadata["document_meta"]["source"]["id"] = self.document_meta.source.id
+
+        return VectorStoreEntry(
+            id=vector_store_entry_id, text=self.key, image_bytes=self.image_representation, metadata=metadata
+        )
+
+
+class TextElement(Element):
+    """
+    An object representing a text element in a document.
+    """
+
+    element_type: str = "text"
+    content: str
+
+    @computed_field  # type: ignore[prop-decorator]
+    @property
+    def text_representation(self) -> str:
+        """
+        Get the text representation of the element.
+
+        Returns:
+            The text representation.
+        """
+        return self.content
+
+
+class ImageElement(Element):
+    """
+    An object representing an image element in a document.
+    """
+
+    element_type: str = "image"
+    image_bytes: SerializableBytes
+    description: str | None = None
+    ocr_extracted_text: str | None = None
+
+    @computed_field  # type: ignore[prop-decorator]
+    @property
+    def text_representation(self) -> str | None:
+        """
+        Get the text representation of the element.
+
+        Returns:
+            The text representation.
+        """
+        if not self.description and not self.ocr_extracted_text:
+            return None
+
+        repr = ""
+        if self.description:
+            repr += f"Description: {self.description}\n"
+        if self.ocr_extracted_text:
+            repr += f"Extracted text: {self.ocr_extracted_text}"
+        return repr
+
+    @property
+    def image_representation(self) -> bytes:
+        """
+        Get the image representation of the element.
+
+        Returns:
+            The image representation.
+        """
+        return self.image_bytes
+
+    def get_id_components(self) -> dict[str, str]:
+        """
+        Creates a dictionary of key value pairs of id components
+
+        Returns:
+            dict: a dictionary
+        """
+        id_components = super().get_id_components()
+        id_components["image_hash"] = hashlib.sha256(self.image_bytes).hexdigest()
+        return id_components

ragbits/document_search/ingestion/enrichers/__init__.py

@@ -0,0 +1,5 @@
+from ragbits.document_search.ingestion.enrichers.base import ElementEnricher
+from ragbits.document_search.ingestion.enrichers.image import ImageElementEnricher
+from ragbits.document_search.ingestion.enrichers.router import ElementEnricherRouter
+
+__all__ = ["ElementEnricher", "ElementEnricherRouter", "ImageElementEnricher"]

ragbits/document_search/ingestion/enrichers/base.py

@@ -0,0 +1,64 @@
+from abc import ABC, abstractmethod
+from types import ModuleType, UnionType
+from typing import ClassVar, Generic, TypeVar, get_args, get_origin
+
+from ragbits.core.utils.config_handling import WithConstructionConfig
+from ragbits.document_search.documents.element import Element
+from ragbits.document_search.ingestion import enrichers
+from ragbits.document_search.ingestion.enrichers.exceptions import EnricherElementNotSupportedError
+
+ElementT = TypeVar("ElementT", bound=Element)
+
+
+class ElementEnricher(Generic[ElementT], WithConstructionConfig, ABC):
+    """
+    Base class for element enrichers, responsible for providing additional information about elements.
+
+    Enrichers operate on raw elements and are used to fill in missing fields that could not be filled in during parsing.
+    They usually deal with summarizing text or describing images.
+    """
+
+    default_module: ClassVar[ModuleType | None] = enrichers
+    configuration_key: ClassVar[str] = "enricher"
+
+    @abstractmethod
+    async def enrich(self, elements: list[ElementT]) -> list[ElementT]:
+        """
+        Enrich elements.
+
+        Args:
+            elements: The elements to be enriched.
+
+        Returns:
+            The list of enriched elements.
+
+        Raises:
+            EnricherError: If the enrichment of the elements failed.
+        """
+
+    @classmethod
+    def validate_element_type(cls, element_type: type[Element]) -> None:
+        """
+        Check if the enricher supports the element type.
+
+        Args:
+            element_type: The element type to validate against the enricher.
+
+        Raises:
+            EnricherElementNotSupportedError: If the element type is not supported.
+        """
+        expected_element_type = cls.__orig_bases__[0].__args__[0]  # type: ignore
+
+        # Check if expected_element_type is a Union and if element_type is in that Union
+        if (
+            (origin := get_origin(expected_element_type))
+            and origin == UnionType
+            and element_type in get_args(expected_element_type)
+        ):
+            return
+
+        # Check if element_type matches expected_element_type exactly
+        if element_type == expected_element_type:
+            return
+
+        raise EnricherElementNotSupportedError(enricher_name=cls.__name__, element_type=element_type)

ragbits/document_search/ingestion/enrichers/exceptions.py

@@ -0,0 +1,32 @@
+from ragbits.document_search.documents.element import Element
+
+
+class EnricherError(Exception):
+    """
+    Class for all exceptions raised by the element enricher and router.
+    """
+
+    def __init__(self, message: str) -> None:
+        super().__init__(message)
+        self.message = message
+
+
+class EnricherNotFoundError(EnricherError):
+    """
+    Raised when no enricher was found for the element type.
+    """
+
+    def __init__(self, element_type: type[Element]) -> None:
+        super().__init__(f"No enricher found for the element type {element_type}")
+        self.element_type = element_type
+
+
+class EnricherElementNotSupportedError(EnricherError):
+    """
+    Raised when the element type is not supported by the enricher.
+    """
+
+    def __init__(self, enricher_name: str, element_type: type[Element]) -> None:
+        super().__init__(f"Element type {element_type} is not supported by the {enricher_name}")
+        self.enricher_name = enricher_name
+        self.element_type = element_type

ragbits/document_search/ingestion/enrichers/image.py

@@ -0,0 +1,107 @@
+from pydantic import BaseModel
+
+from ragbits.core.llms.base import LLM, LLMType
+from ragbits.core.llms.factory import get_preferred_llm
+from ragbits.core.prompt import Attachment, Prompt
+from ragbits.core.utils.config_handling import ObjectConstructionConfig, import_by_path
+from ragbits.document_search.documents.element import ImageElement
+from ragbits.document_search.ingestion.enrichers.base import ElementEnricher
+
+
+class ImageDescriberInput(BaseModel):
+    """
+    Input data for an image describer prompt.
+    """
+
+    image: Attachment
+
+
+class ImageDescriberOutput(BaseModel):
+    """
+    Output data for an image describer prompt.
+    """
+
+    description: str
+
+
+class ImageDescriberPrompt(Prompt[ImageDescriberInput, ImageDescriberOutput]):
+    """
+    Prompt for describing image elements using LLM.
+    """
+
+    user_prompt = "Describe the content of the image."
+
+
+class ImageElementEnricher(ElementEnricher[ImageElement]):
+    """
+    Enricher that describes image elements using LLM.
+    """
+
+    def __init__(
+        self,
+        llm: LLM | None = None,
+        prompt: type[Prompt[ImageDescriberInput, ImageDescriberOutput]] | None = None,
+    ) -> None:
+        """
+        Initialize the ImageElementEnricher instance.
+
+        Args:
+            llm: The language model to use for describing images.
+            prompt: The prompt class to use.
+        """
+        self._llm = llm or get_preferred_llm(llm_type=LLMType.VISION)
+        self._prompt = prompt or ImageDescriberPrompt
+
+    async def enrich(self, elements: list[ImageElement]) -> list[ImageElement]:
+        """
+        Enrich image elements with additional description of the image.
+
+        Args:
+            elements: The elements to be enriched.
+
+        Returns:
+            The list of enriched elements.
+
+        Raises:
+            EnricherElementNotSupportedError: If the element type is not supported.
+            LLMError: If LLM generation fails.
+        """
+        responses: list[ImageDescriberOutput] = []
+        for element in elements:
+            self.validate_element_type(type(element))
+            image = Attachment(data=element.image_bytes)
+            prompt = self._prompt(ImageDescriberInput(image=image))
+            responses.append(await self._llm.generate(prompt))
+
+        return [
+            ImageElement(
+                document_meta=element.document_meta,
+                description=response.description,
+                image_bytes=element.image_bytes,
+                ocr_extracted_text=element.ocr_extracted_text,
+            )
+            for element, response in zip(elements, responses, strict=True)
+        ]
+
+    @classmethod
+    def from_config(cls, config: dict) -> "ImageElementEnricher":
+        """
+        Create an `ImageElementEnricher` instance from a configuration dictionary.
+
+        Args:
+            config: The dictionary containing the configuration settings.
+
+        Returns:
+            The initialized instance of `ImageElementEnricher`.
+
+        Raises:
+            ValidationError: If the configuration doesn't follow the expected format.
+            InvalidConfigError: If llm or prompt can't be found or are not the correct type.
+        """
+        config["llm"] = (
+            LLM.subclass_from_config(ObjectConstructionConfig.model_validate(config["llm"]))
+            if "llm" in config
+            else None
+        )
+        config["prompt"] = import_by_path(config["prompt"]) if "prompt" in config else None
+        return super().from_config(config)

ragbits/document_search/ingestion/enrichers/router.py

@@ -0,0 +1,86 @@
+from collections.abc import Mapping
+from typing import ClassVar
+
+from typing_extensions import Self
+
+from ragbits.core.utils.config_handling import ObjectConstructionConfig, WithConstructionConfig, import_by_path
+from ragbits.document_search.documents import element
+from ragbits.document_search.documents.element import Element
+from ragbits.document_search.ingestion.enrichers.base import ElementEnricher
+from ragbits.document_search.ingestion.enrichers.exceptions import EnricherNotFoundError
+
+_DEFAULT_ENRICHERS: dict[type[Element], ElementEnricher] = {}
+
+
+class ElementEnricherRouter(WithConstructionConfig):
+    """
+    The class responsible for routing the element to the correct enricher based on the element type.
+    """
+
+    configuration_key: ClassVar[str] = "enricher_router"
+
+    _enrichers: Mapping[type[Element], ElementEnricher]
+
+    def __init__(
+        self,
+        enrichers: Mapping[type[Element], ElementEnricher] | None = None,
+    ) -> None:
+        """
+        Initialize the ElementEnricherRouter instance.
+
+        Args:
+            enrichers: The mapping of element types and their enrichers. To override default enrichers.
+        """
+        self._enrichers = {**_DEFAULT_ENRICHERS, **enrichers} if enrichers else _DEFAULT_ENRICHERS
+
+    def __contains__(self, element_type: type[Element]) -> bool:
+        """
+        Check if there is an enricher defined of the given element type.
+
+        Args:
+            element_type: The element type.
+
+        Returns:
+            True if the enricher is defined for the element, otherwise False.
+        """
+        return element_type in self._enrichers
+
+    @classmethod
+    def from_config(cls, config: dict[str, ObjectConstructionConfig]) -> Self:
+        """
+        Initialize the class with the provided configuration.
+
+        Args:
+            config: A dictionary containing configuration details for the class.
+
+        Returns:
+            The ElementEnricherRouter.
+
+        Raises:
+            InvalidConfigError: If any of the provided parsers cannot be initialized.
+        """
+        enrichers: dict[type[Element], ElementEnricher] = {
+            import_by_path(element_type, element): ElementEnricher.subclass_from_config(enricher_config)
+            for element_type, enricher_config in config.items()
+        }
+        return super().from_config({"enrichers": enrichers})
+
+    def get(self, element_type: type[Element]) -> ElementEnricher:
+        """
+        Get the enricher for the element.
+
+        Args:
+            element_type: The element type.
+
+        Returns:
+            The enricher for processing the element.
+
+        Raises:
+            EnricherNotFoundError: If no enricher is found for the element type.
+        """
+        enricher = self._enrichers.get(element_type)
+
+        if isinstance(enricher, ElementEnricher):
+            return enricher
+
+        raise EnricherNotFoundError(element_type)

ragbits/document_search/ingestion/parsers/__init__.py

@@ -0,0 +1,9 @@
+from ragbits.document_search.ingestion.parsers.base import DocumentParser, ImageDocumentParser, TextDocumentParser
+from ragbits.document_search.ingestion.parsers.router import DocumentParserRouter
+
+__all__ = [
+    "DocumentParser",
+    "DocumentParserRouter",
+    "ImageDocumentParser",
+    "TextDocumentParser",
+]

ragbits/document_search/ingestion/parsers/base.py

@@ -0,0 +1,97 @@
+from abc import ABC, abstractmethod
+from types import ModuleType
+from typing import ClassVar
+
+from ragbits.core.utils.config_handling import WithConstructionConfig
+from ragbits.document_search.documents.document import Document, DocumentType
+from ragbits.document_search.documents.element import Element, ImageElement, TextElement
+from ragbits.document_search.ingestion import parsers
+from ragbits.document_search.ingestion.parsers.exceptions import ParserDocumentNotSupportedError
+
+
+class DocumentParser(WithConstructionConfig, ABC):
+    """
+    Base class for document parsers, responsible for converting the document into a list of elements.
+    """
+
+    default_module: ClassVar[ModuleType | None] = parsers
+    configuration_key: ClassVar[str] = "parser"
+
+    supported_document_types: set[DocumentType] = set()
+
+    @abstractmethod
+    async def parse(self, document: Document) -> list[Element]:
+        """
+        Parse the document.
+
+        Args:
+            document: The document to parse.
+
+        Returns:
+            The list of elements extracted from the document.
+
+        Raises:
+            ParserError: If the parsing of the document failed.
+        """
+
+    @classmethod
+    def validate_document_type(cls, document_type: DocumentType) -> None:
+        """
+        Check if the parser supports the document type.
+
+        Args:
+            document_type: The document type to validate against the parser.
+
+        Raises:
+            ParserDocumentNotSupportedError: If the document type is not supported.
+        """
+        if document_type not in cls.supported_document_types:
+            raise ParserDocumentNotSupportedError(parser_name=cls.__name__, document_type=document_type)
+
+
+class TextDocumentParser(DocumentParser):
+    """
+    Simple parser that maps a text to the text element.
+    """
+
+    supported_document_types = {DocumentType.TXT, DocumentType.MD}
+
+    async def parse(self, document: Document) -> list[Element]:
+        """
+        Parse the document.
+
+        Args:
+            document: The document to parse.
+
+        Returns:
+            List with an text element with the text content.
+
+        Raises:
+            ParserDocumentNotSupportedError: If the document type is not supported by the parser.
+        """
+        self.validate_document_type(document.metadata.document_type)
+        return [TextElement(content=document.local_path.read_text(), document_meta=document.metadata)]
+
+
+class ImageDocumentParser(DocumentParser):
+    """
+    Simple parser that maps an image to the image element.
+    """
+
+    supported_document_types = {DocumentType.JPG, DocumentType.PNG}
+
+    async def parse(self, document: Document) -> list[Element]:
+        """
+        Parse the document.
+
+        Args:
+            document: The document to parse.
+
+        Returns:
+            List with an image element with the image content.
+
+        Raises:
+            ParserDocumentNotSupportedError: If the document type is not supported by the parser.
+        """
+        self.validate_document_type(document.metadata.document_type)
+        return [ImageElement(image_bytes=document.local_path.read_bytes(), document_meta=document.metadata)]