r2rpy/py/core/providers/ingestion/unstructured/base.py

import asyncio
import base64
import logging
import os
import time
from copy import copy
from io import BytesIO
from typing import Any, AsyncGenerator

import httpx
from unstructured_client import UnstructuredClient
from unstructured_client.models import operations, shared

from core import parsers
from core.base import (
    AsyncParser,
    ChunkingStrategy,
    Document,
    DocumentChunk,
    DocumentType,
    RecursiveCharacterTextSplitter,
)
from core.base.abstractions import R2RSerializable
from core.base.providers.ingestion import IngestionConfig, IngestionProvider
from core.providers.ocr import MistralOCRProvider
from core.utils import generate_extraction_id

from ...database import PostgresDatabaseProvider
from ...llm import (
    LiteLLMCompletionProvider,
    OpenAICompletionProvider,
    R2RCompletionProvider,
)

logger = logging.getLogger()


class FallbackElement(R2RSerializable):
    text: str
    metadata: dict[str, Any]


class UnstructuredIngestionConfig(IngestionConfig):
    combine_under_n_chars: int = 128
    max_characters: int = 500
    new_after_n_chars: int = 1500
    overlap: int = 64

    coordinates: bool | None = None
    encoding: str | None = None  # utf-8
    extract_image_block_types: list[str] | None = None
    gz_uncompressed_content_type: str | None = None
    hi_res_model_name: str | None = None
    include_orig_elements: bool | None = None
    include_page_breaks: bool | None = None

    languages: list[str] | None = None
    multipage_sections: bool | None = None
    ocr_languages: list[str] | None = None
    # output_format: Optional[str] = "application/json"
    overlap_all: bool | None = None
    pdf_infer_table_structure: bool | None = None

    similarity_threshold: float | None = None
    skip_infer_table_types: list[str] | None = None
    split_pdf_concurrency_level: int | None = None
    split_pdf_page: bool | None = None
    starting_page_number: int | None = None
    strategy: str | None = None
    chunking_strategy: str | ChunkingStrategy | None = None  # type: ignore
    unique_element_ids: bool | None = None
    xml_keep_tags: bool | None = None

    def to_ingestion_request(self):
        import json

        x = json.loads(self.json())
        x.pop("extra_fields", None)
        x.pop("provider", None)
        x.pop("excluded_parsers", None)

        x = {k: v for k, v in x.items() if v is not None}
        return x


class UnstructuredIngestionProvider(IngestionProvider):
    R2R_FALLBACK_PARSERS = {
        DocumentType.GIF: [parsers.ImageParser],  # type: ignore
        DocumentType.JPEG: [parsers.ImageParser],  # type: ignore
        DocumentType.JPG: [parsers.ImageParser],  # type: ignore
        DocumentType.PNG: [parsers.ImageParser],  # type: ignore
        DocumentType.SVG: [parsers.ImageParser],  # type: ignore
        DocumentType.HEIC: [parsers.ImageParser],  # type: ignore
        DocumentType.MP3: [parsers.AudioParser],  # type: ignore
        DocumentType.JSON: [parsers.JSONParser],  # type: ignore
        DocumentType.HTML: [parsers.HTMLParser],  # type: ignore
        DocumentType.XLS: [parsers.XLSParser],  # type: ignore
        DocumentType.XLSX: [parsers.XLSXParser],  # type: ignore
        #DocumentType.DOC: [parsers.DOCParser],  # type: ignore
        DocumentType.PPT: [parsers.PPTParser],  # type: ignore
        DocumentType.CSV: [parsers.CSVParserAdvanced],  # type: ignore
    }

    EXTRA_PARSERS = {
        #DocumentType.CSV: {"advanced": parsers.CSVParserAdvanced},  # type: ignore
        DocumentType.PDF: {
            "ocr": parsers.OCRPDFParser,  # type: ignore
            "unstructured": parsers.PDFParserUnstructured,  # type: ignore
            "zerox": parsers.VLMPDFParser,  # type: ignore
        }
        #DocumentType.XLSX: {"advanced": parsers.XLSXParserAdvanced},  # type: ignore
    }

    IMAGE_TYPES = {
        DocumentType.GIF,
        DocumentType.HEIC,
        DocumentType.JPG,
        DocumentType.JPEG,
        DocumentType.PNG,
        DocumentType.SVG,
    }

    def __init__(
        self,
        config: UnstructuredIngestionConfig,
        database_provider: PostgresDatabaseProvider,
        llm_provider: (
            LiteLLMCompletionProvider
            | OpenAICompletionProvider
            | R2RCompletionProvider
        ),
        ocr_provider: MistralOCRProvider,
    ):
        super().__init__(config, database_provider, llm_provider)
        self.config: UnstructuredIngestionConfig = config
        self.database_provider: PostgresDatabaseProvider = database_provider
        self.llm_provider: (
            LiteLLMCompletionProvider
            | OpenAICompletionProvider
            | R2RCompletionProvider
        ) = llm_provider
        self.ocr_provider: MistralOCRProvider = ocr_provider

        self.client: UnstructuredClient | httpx.AsyncClient
        #config.provider = "unstructured_api"
        if config.provider == "unstructured_api":
            try:
                self.unstructured_api_auth = os.environ["UNSTRUCTURED_API_KEY"]
            except KeyError as e:
                raise ValueError(
                    "UNSTRUCTURED_API_KEY environment variable is not set"
                ) from e

            self.unstructured_api_url = os.environ.get(
                "UNSTRUCTURED_API_URL",
                "https://api.unstructuredapp.io/general/v0/general",
            )

            self.client = UnstructuredClient(
                api_key_auth=self.unstructured_api_auth,
                server_url=self.unstructured_api_url,
            )
            self.shared = shared
            self.operations = operations

        else:
            try:
                self.local_unstructured_url = os.environ[
                    "UNSTRUCTURED_SERVICE_URL"
                ]
            except KeyError as e:
                raise ValueError(
                    "UNSTRUCTURED_SERVICE_URL environment variable is not set"
                ) from e

            self.client = httpx.AsyncClient()

        self.parsers: dict[DocumentType, AsyncParser] = {}
        self._initialize_parsers()
 
    def _initialize_parsers(self):
        for doc_type, parsers in self.R2R_FALLBACK_PARSERS.items():
            for parser in parsers:
                if (
                    doc_type not in self.config.excluded_parsers
                    and doc_type not in self.parsers
                ):
                    # will choose the first parser in the list
                    self.parsers[doc_type] = parser(
                        config=self.config,
                        database_provider=self.database_provider,
                        llm_provider=self.llm_provider,
                    )
        # TODO - Reduce code duplication between Unstructured & R2R
        for doc_type, parser_names in self.config.extra_parsers.items():
            if not isinstance(parser_names, list):
                parser_names = [parser_names]

            for parser_name in parser_names:
                parser_key = f"{parser_name}_{str(doc_type)}"

                try:
                    self.parsers[parser_key] = self.EXTRA_PARSERS[doc_type][
                        parser_name
                    ](
                        config=self.config,
                        database_provider=self.database_provider,
                        llm_provider=self.llm_provider,
                        ocr_provider=self.ocr_provider,
                    )
                    logger.info(
                        f"Initialized extra parser {parser_name} for {doc_type}"
                    )
                except KeyError as e:
                    logger.error(
                        f"Parser {parser_name} for document type {doc_type} not found: {e}"
                    )

    async def parse_fallback(
        self,
        file_content: bytes,
        ingestion_config: dict,
        parser_name: str,
    ) -> AsyncGenerator[FallbackElement, None]:
        contents = []
        async for chunk in self.parsers[parser_name].ingest(  # type: ignore
            file_content, **ingestion_config
        ):  # type: ignore
            if isinstance(chunk, dict) and chunk.get("content"):
                contents.append(chunk)
            elif chunk:  # Handle string output for backward compatibility
                contents.append({"content": chunk})

        if not contents:
            logging.warning(
                "No valid text content was extracted during parsing"
            )
            return

        logging.info(f"Fallback ingestion with config = {ingestion_config}")

        vlm_ocr_one_page_per_chunk = ingestion_config.get(
            "vlm_ocr_one_page_per_chunk", True
        )

        iteration = 0
        for content_item in contents:
            text = content_item["content"]

            if vlm_ocr_one_page_per_chunk and parser_name.startswith(
                ("zerox_", "ocr_")
            ):
                # Use one page per chunk for OCR/VLM
                metadata = {"chunk_id": iteration}
                if "page_number" in content_item:
                    metadata["page_number"] = content_item["page_number"]

                yield FallbackElement(
                    text=text or "No content extracted.",
                    metadata=metadata,
                )
                iteration += 1
                await asyncio.sleep(0)
            else:
                # Use regular text splitting
                loop = asyncio.get_event_loop()
                splitter = RecursiveCharacterTextSplitter(
                    chunk_size=ingestion_config["new_after_n_chars"],
                    chunk_overlap=ingestion_config["overlap"],
                )
                chunks = await loop.run_in_executor(
                    None, splitter.create_documents, [text]
                )

                for text_chunk in chunks:
                    metadata = {"chunk_id": iteration}
                    if "page_number" in content_item:
                        metadata["page_number"] = content_item["page_number"]

                    yield FallbackElement(
                        text=text_chunk.page_content,
                        metadata=metadata,
                    )
                    iteration += 1
                    await asyncio.sleep(0)

    async def parse(
        self,
        file_content: bytes,
        document: Document,
        ingestion_config_override: dict,
    ) -> AsyncGenerator[DocumentChunk, None]:
        ingestion_config = copy(
            {
                **self.config.to_ingestion_request(),
                **(ingestion_config_override or {}),
            }
        )
        # cleanup extra fields
        ingestion_config.pop("provider", None)
        ingestion_config.pop("excluded_parsers", None)

        t0 = time.time()
        parser_overrides = ingestion_config_override.get(
            "parser_overrides", {}
        )
        elements = []
        #parser_overrides = {"pdf": "unstructured"}

        # TODO - Cleanup this approach to be less hardcoded
        # TODO - Remove code duplication between Unstructured & R2R
        logger.info(f"Parser overrides: {parser_overrides}")
        logger.info(f"R2R fallback parsers is: {document.document_type.value}")
        logger.info(f"R2R fallback parsers is: {self.EXTRA_PARSERS.keys()}")
        logger.info(f"R2R fallback parsers is: {document.document_type.value in self.EXTRA_PARSERS.keys()}")
        logger.info(f"R2R fallback parsers is: {document.document_type.value in parser_overrides}")
        #if document.document_type.value in parser_overrides:
        if document.document_type.value in parser_overrides:
            logger.info(
                f"Using parser_override for {document.document_type} with input value {parser_overrides[document.document_type.value]}"
            )
            if parser_overrides[document.document_type.value] == "zerox":
                async for element in self.parse_fallback(
                    file_content,
                    ingestion_config=ingestion_config,
                    parser_name=f"zerox_{DocumentType.PDF.value}",
                ):
                    logger.warning(
                        f"Using parser_override for {document.document_type}"
                    )
                    elements.append(element)
            elif parser_overrides[document.document_type.value] == "ocr":
                async for element in self.parse_fallback(
                    file_content,
                    ingestion_config=ingestion_config,
                    parser_name=f"ocr_{DocumentType.PDF.value}",
                ):
                    logger.warning(
                        f"Using OCR parser_override for {document.document_type}"
                    )
                    elements.append(element)

        elif document.document_type in self.R2R_FALLBACK_PARSERS.keys():
            logger.info(
                f"Parsing {document.document_type}: {document.id} with fallback parser"
            )
            try:
                async for element in self.parse_fallback(
                    file_content,
                    ingestion_config=ingestion_config,
                    parser_name=document.document_type,
                ):
                    elements.append(element)
            except Exception as e:
                logger.error(f"Error parsing {document.document_type}: {e}")
                raise e
        else:
            logger.info(
                f"Parsing {document.document_type}: {document.id} with unstructured"
            )

            file_io = BytesIO(file_content)
            logger.info(f"Provider is: {self.config.provider}")
            # TODO - Include check on excluded parsers here.
            if self.config.provider == "unstructured_api":
                logger.info(f"Using API to parse document {document.id}")
                files = self.shared.Files(
                    content=file_io.read(),
                    file_name=document.metadata.get("title", "unknown_file"),
                )

                ingestion_config.pop("app", None)
                ingestion_config.pop("extra_parsers", None)

                req = self.operations.PartitionRequest(
                    partition_parameters=self.shared.PartitionParameters(
                        files=files,
                        **ingestion_config,
                    )
                )
                elements = await self.client.general.partition_async(  # type: ignore
                    request=req
                )
                elements = list(elements.elements)  # type: ignore

            else:
                logger.info(
                    f"Using local unstructured fastapi server to parse document {document.id}"
                )
                # Base64 encode the file content
                encoded_content = base64.b64encode(file_io.read()).decode(
                    "utf-8"
                )

                logger.info(
                    f"Sending a request to {self.local_unstructured_url}/partition"
                )
                #ingestion_config["strategy"] = "hi_res"
                print(ingestion_config)
                response = await self.client.post(
                    f"{self.local_unstructured_url}/partition",
                    json={
                        "file_content": encoded_content,  # Use encoded string
                        "ingestion_config": ingestion_config,
                        "filename": document.metadata.get("title", None),
                    },
                    timeout=3600,  # Adjust timeout as needed
                )

                if response.status_code != 200:
                    logger.error(f"Error partitioning file: {response.text}")
                    raise ValueError(
                        f"Error partitioning file: {response.text}"
                    )
                elements = response.json().get("elements", [])

        iteration = 0  # if there are no chunks
        for iteration, element in enumerate(elements):
            if isinstance(element, FallbackElement):
                text = element.text
                metadata = copy(document.metadata)
                metadata.update(element.metadata)
            else:
                element_dict = (
                    element if isinstance(element, dict) else element.to_dict()
                )
                text = element_dict.get("text", "")
                if text == "":
                    continue

                metadata = copy(document.metadata)
                for key, value in element_dict.items():
                    if key == "metadata":
                        for k, v in value.items():
                            if k not in metadata and k != "orig_elements":
                                metadata[f"unstructured_{k}"] = v

            # indicate that the document was chunked using unstructured
            # nullifies the need for chunking in the pipeline
            metadata["partitioned_by_unstructured"] = True
            metadata["chunk_order"] = iteration
            # creating the text extraction
            yield DocumentChunk(
                id=generate_extraction_id(document.id, iteration),
                document_id=document.id,
                owner_id=document.owner_id,
                collection_ids=document.collection_ids,
                data=text,
                metadata=metadata,
            )

        logger.debug(
            f"Parsed document with id={document.id}, title={document.metadata.get('title', None)}, "
            f"user_id={document.metadata.get('user_id', None)}, metadata={document.metadata} "
            f"into {iteration + 1} extractions in t={time.time() - t0:.2f} seconds."
        )

    def get_parser_for_document_type(self, doc_type: DocumentType) -> str:
        return "unstructured_api"