megafarce/app.py

# app.py

import asyncio
import base64
import enum
import io
import os
import re
import tempfile
import time
from enum import Enum
from pathlib import Path
from typing import Any, Dict, IO, List, Optional, Union

from abc import ABC, abstractmethod

import httpx
import nats
import psutil
import uvicorn
from dotenv import load_dotenv
from fastapi import (
    Depends,
    FastAPI,
    File,
    Form,
    HTTPException,
    UploadFile,
)
from langchain_anthropic import ChatAnthropic
from langchain_community.document_loaders import PlaywrightURLLoader
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from llama_index.core.schema import Document as LlamaDocument
from llama_parse import LlamaParse as _LlamaParse
from llama_parse.utils import Language, ResultType
#from megaparse_sdk.schema.languages import Language as MP_Language
from pydantic import BaseModel, Field, ValidationError, validator



# Load environment variables
load_dotenv()


# ---------------------- Langauges --------------------------

from enum import Enum


class Language(str, Enum):
    BAZA = "abq"
    ADYGHE = "ady"
    AFRIKAANS = "af"
    ANGIKA = "ang"
    ARABIC = "ar"
    ASSAMESE = "as"
    AVAR = "ava"
    AZERBAIJANI = "az"
    BELARUSIAN = "be"
    BULGARIAN = "bg"
    BIHARI = "bh"
    BHOJPURI = "bho"
    BENGALI = "bn"
    BOSNIAN = "bs"
    SIMPLIFIED_CHINESE = "ch_sim"
    TRADITIONAL_CHINESE = "ch_tra"
    CHECHEN = "che"
    CZECH = "cs"
    WELSH = "cy"
    DANISH = "da"
    DARGWA = "dar"
    GERMAN = "de"
    ENGLISH = "en"
    SPANISH = "es"
    ESTONIAN = "et"
    PERSIAN_FARSI = "fa"
    FRENCH = "fr"
    IRISH = "ga"
    GOAN_KONKANI = "gom"
    HINDI = "hi"
    CROATIAN = "hr"
    HUNGARIAN = "hu"
    INDONESIAN = "id"
    INGUSH = "inh"
    ICELANDIC = "is"
    ITALIAN = "it"
    JAPANESE = "ja"
    KABARDIAN = "kbd"
    KANNADA = "kn"
    KOREAN = "ko"
    KURDISH = "ku"
    LATIN = "la"
    LAK = "lbe"
    LEZGHIAN = "lez"
    LITHUANIAN = "lt"
    LATVIAN = "lv"
    MAGAHI = "mah"
    MAITHILI = "mai"
    MAORI = "mi"
    MONGOLIAN = "mn"
    MARATHI = "mr"
    MALAY = "ms"
    MALTESE = "mt"
    NEPALI = "ne"
    NEWARI = "new"
    DUTCH = "nl"
    NORWEGIAN = "no"
    OCCITAN = "oc"
    PALI = "pi"
    POLISH = "pl"
    PORTUGUESE = "pt"
    ROMANIAN = "ro"
    RUSSIAN = "ru"
    SERBIAN_CYRILLIC = "rs_cyrillic"
    SERBIAN_LATIN = "rs_latin"
    NAGPURI = "sck"
    SLOVAK = "sk"
    SLOVENIAN = "sl"
    ALBANIAN = "sq"
    SWEDISH = "sv"
    SWAHILI = "sw"
    TAMIL = "ta"
    TABASSARAN = "tab"
    TELUGU = "te"
    THAI = "th"
    TAJIK = "tjk"
    TAGALOG = "tl"
    TURKISH = "tr"
    UYGHUR = "ug"
    UKRAINIAN = "uk"
    URDU = "ur"
    UZBEK = "uz"
    VIETNAMESE = "vi"

MP_Language = Language

# --------------------- Exception Classes ---------------------

class HTTPModelNotSupported(HTTPException):
    def __init__(
        self,
        detail: str = "The requested model is not supported yet.",
        headers: Dict[str, Any] | None = None,
    ):
        super().__init__(status_code=501, detail=detail, headers=headers)


class HTTPFileNotFound(HTTPException):
    def __init__(
        self,
        message="The UploadFile.filename does not exist and is needed for this operation",
    ):
        super().__init__(status_code=404, detail=message)


class HTTPDownloadError(HTTPException):
    def __init__(self, file_name: str, message: str = "Failed to download the file"):
        message = f"{file_name} : {message}"
        super().__init__(status_code=400, detail=message)


class HTTPParsingException(HTTPException):
    def __init__(self, file_name: str, message: str = "Failed to parse the file"):
        message = f"{file_name} : {message}"
        super().__init__(status_code=500, detail=message)


class ParsingException(Exception):
    """Exception raised for errors in the parsing process."""

    def __init__(self, message: str = "An error occurred during parsing"):
        self.message = message
        super().__init__(self.message)

# --------------------- Pydantic Models ---------------------

class MarkDownType(str, Enum):
    """Markdown type enumeration."""

    TITLE = "Title"
    SUBTITLE = "Subtitle"
    HEADER = "Header"
    FOOTER = "Footer"
    NARRATIVE_TEXT = "NarrativeText"
    LIST_ITEM = "ListItem"
    TABLE = "Table"
    PAGE_BREAK = "PageBreak"
    IMAGE = "Image"
    FORMULA = "Formula"
    FIGURE_CAPTION = "FigureCaption"
    ADDRESS = "Address"
    EMAIL_ADDRESS = "EmailAddress"
    CODE_SNIPPET = "CodeSnippet"
    PAGE_NUMBER = "PageNumber"
    DEFAULT = "Default"
    UNDEFINED = "Undefined"


class ParserType(str, Enum):
    """Parser type enumeration."""

    UNSTRUCTURED = "unstructured"
    LLAMA_PARSER = "llama_parser"
    MEGAPARSE_VISION = "megaparse_vision"


class StrategyEnum(str, Enum):
    """Method to use for the conversion"""

    FAST = "fast"
    AUTO = "auto"
    HI_RES = "hi_res"


class SupportedModel(str, Enum):
    """Supported models enumeration."""

    GPT_4O = "gpt-4o"
    GPT_4O_MINI = "gpt-o1-mini"
    CLAUDE_3_5_SONNET = "claude-3-5-sonnet"
    CLAUDE_3_OPUS = "claude-3-opus"

    def __str__(self):
        return self.value

    @classmethod
    def is_supported(cls, model_name: str) -> bool:
        """Check if the model is supported."""
        return model_name in cls.__members__.values()


class APIOutputType(str, Enum):
    PARSE_OK = "parse_file_ok"
    PARSE_ERR = "parse_file_err"


class APIOutput(BaseModel):
    message: str
    result: str


class UploadFileConfig(BaseModel):
    method: ParserType = ParserType.UNSTRUCTURED
    strategy: StrategyEnum = StrategyEnum.AUTO
    check_table: bool = False
    language: Language = Language.ENGLISH
    parsing_instruction: Optional[str] = None
    model_name: SupportedModel = SupportedModel.GPT_4O

    @validator("model_name")
    def validate_model(cls, v):
        if not SupportedModel.is_supported(v.value):
            raise ValueError("Unsupported model selected.")
        return v


# --------------------- Parser Classes ---------------------

class BaseParser(ABC):
    """Mother Class for all the parsers [Unstructured, LlamaParse, MegaParseVision]"""

    @abstractmethod
    async def convert(
        self,
        file_path: str | Path | None = None,
        file: IO[bytes] | None = None,
        **kwargs,
    ) -> str:
        """
        Convert the given file to a specific format.

        Args:
            file_path (str | Path): The path to the file to be converted.
            **kwargs: Additional keyword arguments for the conversion process.

        Returns:
            str: The result of the conversion process.

        Raises:
            NotImplementedError: If the method is not implemented by a subclass.
        """
        raise NotImplementedError("Subclasses should implement this method")


class UnstructuredParser(BaseParser):
    def __init__(
        self, strategy=StrategyEnum.AUTO, model: Optional[BaseChatModel] = None, **kwargs
    ):
        self.strategy = strategy
        self.model = model

    # Function to convert element category to markdown format
    def convert_to_markdown(self, elements: List[Dict[str, Any]]) -> str:
        markdown_content = ""

        for el in elements:
            markdown_content += self.get_markdown_line(el)

        return markdown_content

    def get_markdown_line(self, el: Dict[str, Any]) -> str:
        element_type = el["type"]
        text = el["text"]
        metadata = el["metadata"]
        parent_id = metadata.get("parent_id", None)
        category_depth = metadata.get("category_depth", 0)

        # Markdown line defaults to empty
        markdown_line = ""

        # Element type-specific markdown content
        markdown_types = {
            "Title": f"## {text}\n\n" if parent_id else f"# {text}\n\n",
            "Subtitle": f"## {text}\n\n",
            "Header": f"{'#' * (category_depth + 1)} {text}\n\n",
            "Footer": f"#### {text}\n\n",
            "NarrativeText": f"{text}\n\n",
            "ListItem": f"- {text}\n",
            "Table": f"{text}\n\n",
            "PageBreak": "---\n\n",
            "Image": f"![Image]({el['metadata'].get('image_path', '')})\n\n",
            "Formula": f"$$ {text} $$\n\n",
            "FigureCaption": f"**Figure:** {text}\n\n",
            "Address": f"**Address:** {text}\n\n",
            "EmailAddress": f"**Email:** {text}\n\n",
            "CodeSnippet": f"```{el['metadata'].get('language', '')}\n{text}\n```\n\n",
            "PageNumber": "",  # Page number is not included in markdown
        }

        markdown_line = markdown_types.get(element_type, f"{text}\n\n")

        if element_type == "Table" and self.model:
            # FIXME: @Chloé - Add a modular table enhancement here - LVM
            prompt = ChatPromptTemplate.from_messages(
                [
                    (
                        "human",
                        """You are an expert in markdown tables, match this text and this html table to fill a md table. You answer with just the table in pure markdown, nothing else.
                        <TEXT>
                        {text}
                        </TEXT>
                        <HTML>
                        {html}
                        </HTML>
                        <PREVIOUS_TABLE>
                        {previous_table}
                        </PREVIOUS_TABLE>""",
                    ),
                ]
            )
            chain = prompt | self.model
            result = chain.invoke(
                {
                    "text": el["text"],
                    "html": metadata["text_as_html"],
                    "previous_table": "",
                }
            )
            content_str = (
                str(result.content)
                if not isinstance(result.content, str)
                else result.content
            )
            cleaned_content = re.sub(r"^```.*$\n?", "", content_str, flags=re.MULTILINE)
            markdown_line = f"[TABLE]\n{cleaned_content}\n[/TABLE]\n\n"

        return markdown_line

    async def convert(
        self,
        file_path: str | Path | None = None,
        file: IO[bytes] | None = None,
        **kwargs,
    ) -> str:
        # Partition the PDF
        elements = partition(
            filename=str(file_path) if file_path else None,
            file=file,
            strategy=self.strategy,
            skip_infer_table_types=[],
        )
        elements_dict = [el.to_dict() for el in elements]
        markdown_content = self.convert_to_markdown(elements_dict)
        return markdown_content


class LlamaParser(BaseParser):
    def __init__(
        self,
        api_key: str,
        verbose: bool = True,
        language: Language = Language.ENGLISH,
        parsing_instruction: Optional[str] = None,
        **kwargs,
    ):
        self.api_key = api_key
        self.verbose = verbose
        self.language = language
        if parsing_instruction:
            self.parsing_instruction = parsing_instruction
        else:
            self.parsing_instruction = """Do not take into account the page breaks (no --- between pages),
            do not repeat the header and the footer so the tables are merged if needed. Keep the same format for similar tables."""

    async def convert(
        self,
        file_path: str | Path | None = None,
        file: IO[bytes] | None = None,
        **kwargs,
    ) -> str:
        if not file_path:
            raise ValueError("File_path should be provided to run LlamaParser")

        llama_parser = _LlamaParse(
            api_key=self.api_key,
            result_type=ResultType.MD,
            gpt4o_mode=True,
            verbose=self.verbose,
            language=self.language,
            parsing_instruction=self.parsing_instruction,
        )

        documents: List[LlamaDocument] = await llama_parser.aload_data(str(file_path))
        parsed_md = ""
        for document in documents:
            text_content = document.text
            parsed_md = parsed_md + text_content

        return parsed_md


class MegaParseVision(BaseParser):
    def __init__(self, model: BaseChatModel, **kwargs):
        if hasattr(model, "model_name"):
            if not SupportedModel.is_supported(model.model_name):
                raise ValueError(
                    f"Invald model name, MegaParse vision only supports model that have vision capabilities. "
                    f"{model.model_name} is not supported."
                )
        self.model = model
        self.parsed_chunks: List[str] | None = None

    def process_file(self, file_path: str, image_format: str = "PNG") -> List[str]:
        """
        Process a PDF file and convert its pages to base64 encoded images.

        :param file_path: Path to the PDF file
        :param image_format: Format to save the images (default: PNG)
        :return: List of base64 encoded images
        """
        try:
            images = convert_from_path(file_path)
            images_base64 = []
            for image in images:
                buffered = io.BytesIO()
                image.save(buffered, format=image_format)
                image_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
                images_base64.append(image_base64)
            return images_base64
        except Exception as e:
            raise ValueError(f"Error processing PDF file: {str(e)}")

    def get_element(self, tag: Enum, chunk: str) -> List[str]:
        pattern = rf"\[{tag.value}\]([\s\S]*?)\[/{tag.value}\]"
        all_elmts = re.findall(pattern, chunk)
        if not all_elmts:
            print(f"No {tag.value} found in the chunk")
            return []
        return [elmt.strip() for elmt in all_elmts]

    async def send_to_mlm(self, images_data: List[str]) -> str:
        """
        Send images to the language model for processing.

        :param images_data: List of base64 encoded images
        :return: Processed content as a string
        """
        images_prompt = [
            {
                "type": "image_url",
                "image_url": {"url": f"data:image/jpeg;base64,{image_data}"},
            }
            for image_data in images_data
        ]
        message = {
            "content": [
                {
                    "type": "text",
                    "text": BASE_OCR_PROMPT,
                },
                *images_prompt,
            ]
        }
        response = await self.model.invoke([message])
        return str(response.content)

    def get_cleaned_content(self, parsed_file: str) -> str:
        """
        Get cleaned parsed file without any tags defined in TagEnum.

        This method removes all tags from TagEnum from the parsed file, formats the content,
        and handles the HEADER tag specially by keeping only the first occurrence.

        Args:
            parsed_file (str): The parsed file content with tags.

        Returns:
            str: The cleaned content without TagEnum tags.
        """
        tag_pattern = "|".join(map(re.escape, TagEnum.__members__.values()))
        tag_regex = rf"\[({tag_pattern})\](.*?)\[/\1\]"
        # handle the HEADER tag specially
        header_pattern = rf"\[{TagEnum.HEADER.value}\](.*?)\[/{TagEnum.HEADER.value}\]"
        headers = re.findall(header_pattern, parsed_file, re.DOTALL)
        if headers:
            first_header = headers[0].strip()
            # Remove all HEADER tags and their content
            parsed_file = re.sub(header_pattern, "", parsed_file, flags=re.DOTALL)
            # Add the first header back at the beginning
            parsed_file = f"{first_header}\n{parsed_file}"

        # Remove all other tags
        def remove_tag(match):
            return match.group(2)

        cleaned_content = re.sub(tag_regex, remove_tag, parsed_file, flags=re.DOTALL)

        cleaned_content = re.sub(r"^```.*$\n?", "", cleaned_content, flags=re.MULTILINE)
        cleaned_content = re.sub(r"\n\s*\n", "\n\n", cleaned_content)
        cleaned_content = cleaned_content.replace("|\n\n|", "|\n|")
        cleaned_content = cleaned_content.strip()

        return cleaned_content

    async def convert(
        self,
        file_path: str | Path | None = None,
        file: IO[bytes] | None = None,
        batch_size: int = 3,
        **kwargs,
    ) -> str:
        """
        Parse a PDF file and process its content using the language model.

        :param file_path: Path to the PDF file
        :param batch_size: Number of pages to process concurrently
        :return: List of processed content strings
        """
        if not file_path:
            raise ValueError("File_path should be provided to run MegaParseVision")

        if isinstance(file_path, Path):
            file_path = str(file_path)
        pdf_base64 = self.process_file(file_path)
        tasks = [
            self.send_to_mlm(pdf_base64[i : i + batch_size])
            for i in range(0, len(pdf_base64), batch_size)
        ]
        self.parsed_chunks = await asyncio.gather(*tasks)
        responses = self.get_cleaned_content("\n".join(self.parsed_chunks))
        return responses

# --------------------- MegaParse Class ---------------------

class MegaParse:
    def __init__(
        self,
        parser: BaseParser,
        format_checker: Optional[Any] = None,
    ) -> None:
        self.parser = parser
        self.format_checker = format_checker
        self.last_parsed_document: str = ""

    async def aload(
        self,
        file_path: Path | str | None = None,
        file: IO[bytes] | None = None,
        file_extension: str | None = "",
    ) -> str:
        if not (file_path or file):
            raise ValueError("Either file_path or file should be provided")
        if file_path and file:
            raise ValueError("Only one of file_path or file should be provided")

        if file_path:
            if isinstance(file_path, str):
                file_path = Path(file_path)
            file_extension = file_path.suffix
        elif file:
            if not file_extension:
                raise ValueError(
                    "file_extension should be provided when given file argument"
                )
            file.seek(0)

        try:
            FileExtension(file_extension)
        except ValueError:
            raise ValueError(f"Unsupported file extension: {file_extension}")

        if file_extension != ".pdf":
            if self.format_checker:
                raise ValueError(
                    f"Format Checker : Unsupported file extension: {file_extension}"
                )
            if not isinstance(self.parser, UnstructuredParser):
                raise ValueError(
                    f" Unsupported file extension : Parser {self.parser} do not support {file_extension}"
                )

        try:
            parsed_document: str = await self.parser.convert(
                file_path=file_path, file=file
            )
        except Exception as e:
            raise ParsingException(f"Error while parsing {file_path}: {e}")

        self.last_parsed_document = parsed_document
        return parsed_document

    def load(self, file_path: Path | str) -> str:
        if isinstance(file_path, str):
            file_path = Path(file_path)
        file_extension: str = file_path.suffix

        if file_extension != ".pdf":
            if self.format_checker:
                raise ValueError(
                    f"Format Checker : Unsupported file extension: {file_extension}"
                )
            if not isinstance(self.parser, UnstructuredParser):
                raise ValueError(
                    f"Parser {self.parser}: Unsupported file extension: {file_extension}"
                )

        try:
            loop = asyncio.get_event_loop()
            parsed_document: str = loop.run_until_complete(
                self.parser.convert(file_path=file_path)
            )
        except Exception as e:
            raise ValueError(f"Error while parsing {file_path}: {e}")

        self.last_parsed_document = parsed_document
        return parsed_document

    def save(self, file_path: Path | str) -> None:
        os.makedirs(os.path.dirname(file_path), exist_ok=True)
        with open(file_path, "w+") as f:
            f.write(self.last_parsed_document)

# --------------------- FastAPI App and Endpoints ---------------------

app = FastAPI()

playwright_loader = PlaywrightURLLoader(urls=[], remove_selectors=["header", "footer"])

def parser_builder_dep():
    return ParserBuilder()

def get_playwright_loader():
    return playwright_loader

@app.get("/healthz")
def healthz():
    return {"status": "ok"}

def _check_free_memory() -> bool:
    """Reject traffic when free memory is below minimum (default 2GB)."""
    mem = psutil.virtual_memory()
    memory_free_minimum = int(os.environ.get("MEMORY_FREE_MINIMUM_MB", 2048))

    if mem.available <= memory_free_minimum * 1024 * 1024:
        return False
    return True

@app.post(
    "/v1/file",
    response_model=APIOutput,
)
async def parse_file(
    file: UploadFile = File(...),
    method: ParserType = Form(ParserType.UNSTRUCTURED),
    strategy: StrategyEnum = Form(StrategyEnum.AUTO),
    check_table: bool = Form(False),
    language: MP_Language = Form(MP_Language.ENGLISH),
    parsing_instruction: Optional[str] = Form(None),
    model_name: SupportedModel = Form(SupportedModel.GPT_4O),
    parser_builder=Depends(parser_builder_dep),
) -> Dict[str, str]:
    if not _check_free_memory():
        raise HTTPException(
            status_code=503, detail="Service unavailable due to low memory"
        )
    model = None
    if model_name and check_table:
        if model_name.value.startswith("gpt"):
            model = ChatOpenAI(model=model_name.value, api_key=os.getenv("OPENAI_API_KEY"))  # type: ignore
        elif model_name.value.startswith("claude"):
            model = ChatAnthropic(
                model_name=model_name.value,
                api_key=os.getenv("ANTHROPIC_API_KEY"),  # type: ignore
                timeout=60,
                stop=None,
            )
        else:
            raise HTTPModelNotSupported()

    parser_config = {
        "method": method,
        "strategy": strategy,
        "model": model if model and check_table else None,
        "language": language,
        "parsing_instruction": parsing_instruction,
    }
    try:
        parser = ParserBuilder().build(parser_config)
        megaparse = MegaParse(parser=parser)
        if not file.filename:
            raise HTTPFileNotFound("No filename provided")
        _, extension = os.path.splitext(file.filename)
        file_bytes = await file.read()
        file_stream = io.BytesIO(file_bytes)
        result = await megaparse.aload(file=file_stream, file_extension=extension)
        return {"message": "File parsed successfully", "result": result}
    except ParsingException as e:
        print(e)
        raise HTTPParsingException(file.filename)
    except ValueError as e:
        print(e)
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        print(e)
        raise HTTPException(status_code=500, detail=str(e))

@app.post(
    "/v1/url",
    response_model=APIOutput,
)
async def upload_url(
    url: str, playwright_loader=Depends(get_playwright_loader)
) -> Dict[str, str]:
    playwright_loader.urls = [url]

    if url.endswith(".pdf"):
        # Download the file
        async with httpx.AsyncClient() as client:
            response = await client.get(url)
        if response.status_code != 200:
            raise HTTPDownloadError(url)

        with tempfile.NamedTemporaryFile(delete=False, suffix="pdf") as temp_file:
            temp_file.write(response.content)
            try:
                megaparse = MegaParse(
                    parser=UnstructuredParser(strategy=StrategyEnum.AUTO)
                )
                result = await megaparse.aload(temp_file.name)
                return {"message": "File parsed successfully", "result": result}
            except ParsingException:
                raise HTTPParsingException(url)
    else:
        data = await playwright_loader.aload()
        # Now turn the data into a string
        extracted_content = ""
        for page in data:
            extracted_content += page.page_content
        if not extracted_content:
            raise HTTPDownloadError(
                url,
                message="Failed to extract content from the website. Valid URL example : https://www.quivr.com",
            )
        return {
            "message": "Website content parsed successfully",
            "result": extracted_content,
        }

# --------------------- Parser Builder ---------------------

class ParserBuilder:
    parser_dict: Dict[str, BaseParser] = {
        "unstructured": UnstructuredParser,
        "llama_parser": LlamaParser,
        "megaparse_vision": MegaParseVision,
    }

    def build(self, config: Dict[str, Any]) -> BaseParser:
        """
        Build a parser based on the given configuration.

        Args:
            config (Dict): The configuration to be used for building the parser.

        Returns:
            BaseParser: The built parser.

        Raises:
            ValueError: If the configuration is invalid.
        """
        parser_class = self.parser_dict.get(config["method"])
        if not parser_class:
            raise ValueError(f"Unsupported parser method: {config['method']}")
        return parser_class(**config)

# --------------------- Runner ---------------------

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8001)