# 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 pydantic import BaseModel, Field, ValidationError, validator
from unstructured.partition.auto import partition
# Load environment variables
load_dotenv()
# ---------------------- Langs --------------------------
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
class FileExtension(str, Enum):
"""Supported file extension enumeration."""
PDF = ".pdf"
DOCX = ".docx"
DOC = ".doc"
TXT = ".txt"
OTF = ".odt"
EPUB = ".epub"
HTML = ".html"
XML = ".xml"
CSV = ".csv"
XLSX = ".xlsx"
XLS = ".xls"
PPTX = ".pptx"
PPT = ".ppt"
JSON = ".json"
LOG = ".log"
MD = ".md"
MARKDOWN = ".markdown"
# --------------------- 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"})\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}
{html}
{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)