Marker parsing v2 (#291)

* feat: ajout page endpoint parser

* wip

* wip 1

* change PyMyPDF parsing

* fix: chunking

* fix: wip

* wip

* wip

* change collection for document

* fix: tests and ocr signature

* WIP: sauvegarde avant pull

* Debug

* add parsing test

* remove key

* Update coverage badge

---------

Co-authored-by: Kundun78 <coralexandre055@gmail.com>
Co-authored-by: leoguillaume <leo.guillaume@modernisation.gouv.fr>
Co-authored-by: leoguillaume <leoguillaume@users.noreply.github.com>

Author: 3 people

.github/badges/coverage.json

@@ -1 +1 @@
-{"schemaVersion":1,"label":"coverage","message":"88.36%","color":"green"}
+{"schemaVersion":1,"label":"coverage","message":"87.38%","color":"green"}

.gitignore

@@ -208,4 +208,4 @@ uv.lock
 app/tests/cassettes
 
 .envrc
-.direnv/
+.direnv/

app/clients/parser/__init__.py

@@ -0,0 +1,4 @@
+from ._baseparserclient import BaseParserClient
+from ._markerparserclient import MarkerParserClient
+
+__all__ = ["BaseParserClient", "MarkerParserClient"]

app/clients/parser/_baseparserclient.py

@@ -0,0 +1,33 @@
+from abc import ABC, abstractmethod
+import importlib
+from typing import Optional, Type
+
+from fastapi import UploadFile
+
+from app.schemas.core.settings import ParserType
+from app.schemas.parse import Languages, ParsedDocument, ParsedDocumentOutputFormat
+
+
+class BaseParserClient(ABC):
+    SUPPORTED_FORMATS = []
+
+    @staticmethod
+    def import_module(type: ParserType) -> "Type[BaseParserClient]":
+        """
+        Import the module for the given parser type.
+        """
+        module = importlib.import_module(f"app.clients.parser._{type.value}parserclient")
+        return getattr(module, f"{type.capitalize()}ParserClient")
+
+    @abstractmethod
+    def parse(
+        self,
+        file: UploadFile,
+        output_format: Optional[ParsedDocumentOutputFormat] = None,
+        force_ocr: bool = False,
+        languages: Optional[Languages] = None,
+        page_range: Optional[str] = None,
+        paginate_output: Optional[bool] = None,
+        use_llm: Optional[bool] = None,
+    ) -> ParsedDocument:
+        pass

app/clients/parser/_markerparserclient.py

@@ -0,0 +1,98 @@
+from io import BytesIO
+import json
+from typing import List, Optional
+
+from fastapi import HTTPException
+import httpx
+import pymupdf
+
+from app.schemas.core.documents import FileType, ParserParams
+from app.schemas.parse import ParsedDocument, ParsedDocumentMetadata, ParsedDocumentPage
+
+from ._baseparserclient import BaseParserClient
+
+
+class MarkerParserClient(BaseParserClient):
+    """
+    Class to interact with the Marker PDF API for document analysis.
+    """
+
+    SUPPORTED_FORMATS = [FileType.PDF]
+
+    def __init__(self, api_url: str, api_key: Optional[str] = None, timeout=120, *args, **kwargs) -> None:
+        self.api_url = api_url
+        self.api_key = api_key
+        self.timeout = timeout
+        self.headers = {"Authorization": f"Bearer {self.api_key}"} if self.api_key else {}
+
+        # Keep health check synchronous in __init__
+        response = httpx.get(f"{self.api_url}/health", headers=self.headers, timeout=self.timeout)
+        assert response.status_code == 200, "Marker API is not reachable."
+
+    def convert_page_range(self, page_range: str, page_count: int) -> List[int]:
+        if page_range == "":
+            return [i for i in range(page_count)]
+
+        page_ranges = page_range.split(",")
+        pages = []
+        for page_range in page_ranges:
+            page_range = page_range.split("-")
+            if len(page_range) == 1:
+                pages.append(int(page_range[0]))
+            else:
+                for i in range(int(page_range[0]), int(page_range[1]) + 1):
+                    pages.append(i)
+
+        pages = list(set(pages))
+
+        return pages
+
+    async def parse(self, **params: ParserParams) -> ParsedDocument:
+        params = ParserParams(**params)
+        file_content = await params.file.read()
+
+        try:
+            # Correct way to open PDF from bytes with PyMuPDF
+            pdf = pymupdf.open(stream=file_content, filetype="pdf")
+            page_count = pdf.page_count
+        except Exception as e:
+            # Handle corrupted or invalid PDF files
+            raise HTTPException(status_code=400, detail=f"Invalid PDF file: {str(e)}")
+
+        data = []
+        payload = {
+            "output_format": params.output_format.value,
+            "force_ocr": params.force_ocr,
+            "languages": params.languages.value,
+            "paginate_output": params.paginate_output,
+            "use_llm": params.use_llm,
+        }
+        pages = self.convert_page_range(page_range=params.page_range, page_count=page_count)
+        async with httpx.AsyncClient() as client:
+            for i in pages:
+                # Create a fresh BytesIO object for each request to avoid stream consumption issues
+                files = {"file": (params.file.filename, BytesIO(file_content), "application/pdf")}
+                payload["page_range"] = str(i)
+
+                response = await client.post(
+                    url=f"{self.api_url}/marker/upload",
+                    files=files,
+                    data=payload,
+                    headers=self.headers,
+                    timeout=self.timeout,
+                )
+                if response.status_code != 200:
+                    raise HTTPException(status_code=response.status_code, detail=json.loads(response.text).get("detail", "Parsing failed."))
+
+                result = response.json()
+                if not result.get("success", False):
+                    raise HTTPException(status_code=500, detail=result.get("error", "Parsing failed."))
+
+                metadata = ParsedDocumentMetadata(document_name=params.file.filename, page=i, **result["metadata"])
+                data.append(ParsedDocumentPage(content=result["output"], images=result["images"], metadata=metadata))
+
+        # Close the PDF document to free memory
+        pdf.close()
+        document = ParsedDocument(data=data)
+
+        return document

app/clients/web_search/_bravewebsearchclient.py

@@ -1,5 +1,4 @@
 import logging
-import traceback
 from typing import List
 
 import httpx
@@ -25,8 +24,7 @@ async def search(self, query: str, n: int = 3) -> List[str]:
                 response = await client.get(url=self.URL, headers=self.headers, params=params)
                 results = response.json().get("web", {}).get("results", [])
         except Exception:
-            logger.error(msg="Brave Search API unreachable.")
-            logger.debug(msg=traceback.format_exc())
+            logger.exception(msg="Brave Search API unreachable.")
             results = []
 
         return [result["url"].lower() for result in results]

app/clients/web_search/_duckduckgowebsearchclient.py

@@ -1,5 +1,4 @@
 import logging
-import traceback
 from typing import List
 
 import httpx
@@ -29,8 +28,7 @@ async def search(self, query: str, n: int = 3) -> List[str]:
                 response = await client.get(url=self.URL, headers=self.headers, params=params, follow_redirects=True)
                 results = response.json().get("Results", [])[:n]
         except Exception:
-            logger.error(msg="DuckDuckGo API unreachable.")
-            logger.debug(msg=traceback.format_exc())
+            logger.exception(msg="DuckDuckGo API unreachable.")
             results = []
 
         return [result["FirstURL"].lower() for result in results]

app/endpoints/documents.py

@@ -1,21 +1,114 @@
-from typing import Optional, Union
+from typing import List, Literal, Optional, Union
 from uuid import UUID
 
-from fastapi import APIRouter, Depends, Path, Query, Request, Response, Security
+from fastapi import APIRouter, Depends, Path, Query, Request, Response, Security, UploadFile
 from fastapi.responses import JSONResponse
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from app.helpers import AccessController
-from app.schemas.documents import Document, Documents
+from app.schemas.documents import (
+    ChunkerName,
+    ChunkerNameForm,
+    ChunkMinSizeForm,
+    ChunkOverlapForm,
+    ChunkSizeForm,
+    CollectionForm,
+    Document,
+    DocumentResponse,
+    Documents,
+    IsSeparatorRegexForm,
+    LengthFunctionForm,
+    MetadataForm,
+    SeparatorsForm,
+)
+from app.schemas.parse import (
+    FileForm,
+    ForceOCRForm,
+    Languages,
+    LanguagesForm,
+    OutputFormatForm,
+    PageRangeForm,
+    PaginateOutputForm,
+    ParsedDocumentOutputFormat,
+    UseLLMForm,
+)
 from app.sql.session import get_db as get_session
 from app.utils.context import global_context, request_context
-from app.utils.exceptions import CollectionNotFoundException, DocumentNotFoundException
+from app.utils.exceptions import CollectionNotFoundException, DocumentNotFoundException, FileSizeLimitExceededException
 from app.utils.variables import ENDPOINT__DOCUMENTS
 
 router = APIRouter()
 
 
-@router.get(path=ENDPOINT__DOCUMENTS + "/{document:path}", dependencies=[Security(dependency=AccessController())], status_code=200, response_model=Document)  # fmt: off
+@router.post(path=ENDPOINT__DOCUMENTS, status_code=201, dependencies=[Security(dependency=AccessController())], response_model=DocumentResponse)
+async def create_document(
+    request: Request,
+    session: AsyncSession = Depends(get_session),
+    file: UploadFile = FileForm,
+    collection: int = CollectionForm,
+    paginate_output: Optional[bool] = PaginateOutputForm,
+    page_range: str = PageRangeForm,
+    languages: Optional[Languages] = LanguagesForm,
+    force_ocr: bool = ForceOCRForm,
+    output_format: ParsedDocumentOutputFormat = OutputFormatForm,
+    use_llm: Optional[bool] = UseLLMForm,
+    chunker_name: ChunkerName = ChunkerNameForm,
+    chunk_size: int = ChunkSizeForm,
+    chunk_overlap: int = ChunkOverlapForm,
+    length_function: Literal["len"] = LengthFunctionForm,
+    is_separator_regex: bool = IsSeparatorRegexForm,
+    separators: List[str] = SeparatorsForm,
+    chunk_min_size: int = ChunkMinSizeForm,
+    metadata: str = MetadataForm,
+) -> JSONResponse:
+    """
+    Parse a file and create a document.
+    """
+    if not global_context.documents:  # no vector store available
+        raise CollectionNotFoundException()
+
+    file_size = len(file.file.read())
+    if file_size > FileSizeLimitExceededException.MAX_CONTENT_SIZE:
+        raise FileSizeLimitExceededException()
+    file.file.seek(0)  # reset file pointer to the beginning of the file
+
+    length_function = len if length_function == "len" else length_function
+
+    document = await global_context.parser.parse_file(
+        file=file,
+        collection=collection,
+        paginate_output=paginate_output,
+        page_range=page_range,
+        languages=languages,
+        force_ocr=force_ocr,
+        output_format=output_format,
+        use_llm=use_llm,
+    )
+
+    document_id = await global_context.documents.create_document(
+        user_id=request_context.get().user_id,
+        session=session,
+        collection_id=collection,
+        document=document,
+        chunker_name=chunker_name,
+        chunk_size=chunk_size,
+        chunk_overlap=chunk_overlap,
+        length_function=length_function,
+        is_separator_regex=is_separator_regex,
+        separators=separators,
+        chunk_min_size=chunk_min_size,
+        metadata=metadata,
+    )
+
+    return JSONResponse(content=DocumentResponse(id=document_id).model_dump(), status_code=201)
+
+
+@router.get(
+    path=ENDPOINT__DOCUMENTS + "/{document:path}",
+    dependencies=[Security(dependency=AccessController())],
+    status_code=200,
+    response_model=Document,
+)
 async def get_document(
     request: Request,
     document: int = Path(description="The document ID"),
@@ -68,7 +161,7 @@ async def delete_document(
     session: AsyncSession = Depends(get_session),
 ) -> Response:
     """
-    Delete a document and relative collections.
+    Delete a document.
     """
     if not global_context.documents:  # no vector store available
         raise DocumentNotFoundException()