Use kreuzberg for document extraction + rate limit downloads

Dockerfile

@@ -1,7 +1,7 @@
 FROM python:3.11.3
 
 RUN apt-get update && \
-    apt-get install -y libreoffice libreoffice-writer libreoffice-calc libreoffice-impress && \
+    apt-get install -y libreoffice libreoffice-writer libreoffice-calc libreoffice-impress pandoc && \
     apt-get clean && rm -rf /var/lib/apt/lists/*
 
 RUN useradd -m -u 1000 user

api/docs.py

@@ -1,4 +1,5 @@
 import asyncio
+from aiolimiter import AsyncLimiter
 from pathlib import Path
 import traceback
 from typing import Literal, Tuple
@@ -19,6 +20,7 @@ from fastapi import Depends, HTTPException
 from dependencies import get_http_client, get_llm_router
 from fastapi.responses import StreamingResponse
 from litellm.router import Router
+from kreuzberg import ExtractionConfig, extract_bytes
 
 from schemas import DataRequest, DataResponse, DocRequirements, DocDownloadRequest, MeetingsRequest, MeetingsResponse, ExtractRequirementsRequest, ExtractRequirementsResponse
 
@@ -33,8 +35,11 @@ NSMAP = {
 
 # ================================== Converting of files to .txt ====================================
 
+KREUZBERG_CONFIG: ExtractionConfig = ExtractionConfig(
+    force_ocr=False, ocr_backend=None)
+
 # Unfortunately needs to be kept to 1, as libreoffice isn't built to support parallel instances
-CONVERSION_MUTEX = asyncio.Semaphore(1)
+LO_CONVERSION_MUTEX = asyncio.Semaphore(1)
 
 
 async def convert_file(contents: io.BytesIO, filename: str, input_ext: str, output_ext: str, filter: str = None) -> io.BytesIO:
@@ -50,7 +55,7 @@ async def convert_file(contents: io.BytesIO, filename: str, input_ext: str, outp
         filter: The conversion filter to use.
     """
 
-    await CONVERSION_MUTEX.acquire()
+    await LO_CONVERSION_MUTEX.acquire()
 
     with tempfile.TemporaryDirectory() as tmpdir:
         dir_path = Path(tmpdir)
@@ -91,7 +96,7 @@ async def convert_file(contents: io.BytesIO, filename: str, input_ext: str, outp
                 stderr=stderr
             )
 
-        CONVERSION_MUTEX.release()
+        LO_CONVERSION_MUTEX.release()
 
         with open(output_file_path, mode="rb") as out:
             out_bytes.write(out.read())
@@ -100,31 +105,35 @@ async def convert_file(contents: io.BytesIO, filename: str, input_ext: str, outp
         return out_bytes
 
 
+# Rate limit of FTP downloads per minute
+FTP_DOWNLOAD_RATE_LIMITER = AsyncLimiter(max_rate=120, time_period=60)
+
 async def get_doc_archive(url: str, client: AsyncClient) -> tuple[str, str, io.BytesIO]:
     """Récupère le docx depuis l'URL et le retourne un tuple (nom, extension, contenu)"""
 
-    if not url.endswith("zip"):
-        raise ValueError("URL doit pointer vers un fichier ZIP")
+    async with FTP_DOWNLOAD_RATE_LIMITER:
+        if not url.endswith("zip"):
+            raise ValueError("URL doit pointer vers un fichier ZIP")
 
-    doc_id = os.path.splitext(os.path.basename(url))[0]
-    resp = await client.get(url, headers={
-        "User-Agent": 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
-    })
+        doc_id = os.path.splitext(os.path.basename(url))[0]
+        resp = await client.get(url, headers={
+            "User-Agent": 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
+        })
 
-    resp.raise_for_status()
+        resp.raise_for_status()
 
-    with zipfile.ZipFile(io.BytesIO(resp.content)) as zf:
-        # there should be a single file per file
-        for entry in zf.infolist():
-            if entry.is_dir():
-                continue
+        with zipfile.ZipFile(io.BytesIO(resp.content)) as zf:
+            # there should be a single file per file
+            for entry in zf.infolist():
+                if entry.is_dir():
+                    continue
 
-            file_name = entry.filename
-            root, ext = os.path.splitext(file_name)
-            doc_bytes = zf.read(file_name)
-            return (root, ext.lower(), io.BytesIO(doc_bytes))
+                file_name = entry.filename
+                root, ext = os.path.splitext(file_name)
+                doc_bytes = zf.read(file_name)
+                return (root, ext.lower(), io.BytesIO(doc_bytes))
 
-    raise ValueError("Aucun fichier trouvé dans l'archive")
+        raise ValueError("Aucun fichier trouvé dans l'archive")
 
 
 def apply_docx_revisions(docx_zip: zipfile.ZipFile) -> io.BytesIO:
@@ -189,6 +198,13 @@ def apply_docx_revisions(docx_zip: zipfile.ZipFile) -> io.BytesIO:
     return output
 
 
+FORMAT_MIME_TYPES = {
+    ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+    ".pdf": "application/pdf",
+    ".pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation"
+}
+
+
 async def doc_to_txt(doc_id: str, url: str, client: AsyncClient) -> str:
     """
     Télécharge le TDoc spécifié et le convertit en texte.
@@ -197,31 +213,28 @@ async def doc_to_txt(doc_id: str, url: str, client: AsyncClient) -> str:
     # Grab the document archive
     filename, ext, bytes = await get_doc_archive(url, client)
 
-    final_bytes: io.BytesIO = None
+    final_text: str = None
     if ext == ".doc":
         logging.debug(f"Converting {filename} .doc --> .docx")
         docx_bytes = await convert_file(bytes, doc_id, "doc", "docx")
-        logging.debug(f"Converting {filename} .docx --> .txt")
-        final_bytes = await convert_file(docx_bytes, f"{doc_id}", "docx", "txt")
+        extracted_data = await extract_bytes(docx_bytes.read(), FORMAT_MIME_TYPES[".docx"], config=KREUZBERG_CONFIG)
+        final_text = extracted_data.content
     elif ext == ".docx":
-        logging.debug(f"Updating .docx revisions")
+        # Applying doc revisions to docx files (especially for pCR / draftCR files)
+        logging.debug(f"Updating .docx revisions for {doc_id}.")
         applied_revision = apply_docx_revisions(zipfile.ZipFile(bytes))
-        logging.debug(f"Converting {filename} .docx --> .txt")
-        final_bytes = await convert_file(applied_revision, f"{doc_id}", "docx", "txt")
-    elif ext == ".pdf":
-        logging.debug(f"Converting {filename} .pdf --> .txt")
-        final_bytes = await convert_file(bytes, doc_id, "pdf", "txt")
-    elif ext == ".pptx":
-        logging.debug(f"Converting {filename} .pptx --> .pdf")
-        pdf_bytes = await convert_file(bytes, doc_id, "pptx", "pdf")
-        logging.debug(f"Converting {filename} .pdf --> .txt")
-        final_bytes = await convert_file(pdf_bytes, doc_id, "pdf", "txt")
+        extracted_data = await extract_bytes(applied_revision.read(), FORMAT_MIME_TYPES[".docx"], config=KREUZBERG_CONFIG)
+        final_text = extracted_data.content
     else:
-        raise Exception(f"Unsupported file type: {ext}, filename: {filename}")
+        if ext in FORMAT_MIME_TYPES:  # file extension is supported
+            extracted_data = await extract_bytes(bytes.read(), FORMAT_MIME_TYPES[ext], config=KREUZBERG_CONFIG)
+            final_text = extracted_data.content
+        else:
+            raise Exception(
+                f"Unsupported file type: {ext}, filename: {filename}")
 
-    text_from_bytes = str(final_bytes.read(), encoding="utf-8")
     txt_data = [line.strip()
-                for line in text_from_bytes.splitlines() if line.strip()]
+                for line in final_text.splitlines() if line.strip()]
 
     return txt_data
 

app.py

@@ -21,7 +21,7 @@ logging.basicConfig(
     datefmt='%Y-%m-%d %H:%M:%S'
 )
 
-logging.info("DEBUG logging is disabled. Set `DEBUG_LOG` env var to 1 to enable debug logging.")
+logging.info(f"Set `DEBUG_LOG` env var to 1 to enable debug logging.")
 
 # Initialize global dependencies
 init_dependencies()

requirements.txt

@@ -10,4 +10,5 @@ openpyxl
 beautifulsoup4
 aiolimiter
 httpx
-Jinja2
+Jinja2
+kreuzberg