MegaParce is now MegaFarce
new file: Dockerfile new file: LICENSE new file: README.md new file: app.py new file: docker-compose.yml new file: pyproject.toml new file: requirements.lock
This commit is contained in:
commit
aab74dff66
7 changed files with 1787 additions and 0 deletions
811
app.py
Normal file
811
app.py
Normal file
|
|
@ -0,0 +1,811 @@
|
|||
# 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"})\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)
|
||||
Loading…
Add table
Add a link
Reference in a new issue