Add PDF Mistral OCR extraction

.env.template

@@ -12,6 +12,9 @@ LOG_LEVEL=INFO
 # GET_CONTEXT_MAX_TOKENS=100000
 # MAX_FILE_SIZE=5242880 # Bytes
 # MAX_MESSAGE_SIZE=25000 # Characters
+# MISTRAL_OCR_API_KEY= # If set, PDF text extraction uses Mistral OCR; otherwise it falls back to pdfplumber
+# MISTRAL_OCR_MODEL=mistral-ocr-latest
+# MISTRAL_OCR_TIMEOUT_SECONDS=60.0
 
 # Embedding settings
 # EMBED_MESSAGES=true

config.toml.example

@@ -11,6 +11,9 @@ GET_CONTEXT_MAX_TOKENS = 100000
 MAX_FILE_SIZE = 5242880  # 5MB
 MAX_MESSAGE_SIZE = 25000 # Characters
 EMBED_MESSAGES = true
+MISTRAL_OCR_API_KEY = ""  # If set, PDF text extraction uses Mistral OCR; otherwise it falls back to pdfplumber
+MISTRAL_OCR_MODEL = "mistral-ocr-latest"
+MISTRAL_OCR_TIMEOUT_SECONDS = 60.0
 # LANGFUSE_HOST = "https://api.langfuse.com"
 # LANGFUSE_PUBLIC_KEY = "your-public-key-here"
 # COLLECT_METRICS_LOCAL = false

src/config.py

@@ -1219,6 +1219,11 @@ class AppSettings(HonchoSettings):
 
     MAX_MESSAGE_SIZE: Annotated[int, Field(default=25_000, gt=0)] = 25_000
     EMBED_MESSAGES: bool = True
+    MISTRAL_OCR_API_KEY: str | None = None
+    MISTRAL_OCR_MODEL: str = "mistral-ocr-latest"
+    MISTRAL_OCR_TIMEOUT_SECONDS: Annotated[float, Field(default=60.0, gt=0, le=300)] = (
+        60.0
+    )
     LANGFUSE_HOST: str | None = None
     LANGFUSE_PUBLIC_KEY: str | None = None
 

src/utils/files.py

@@ -1,8 +1,11 @@
+import base64
 import datetime
 import logging
 from io import BytesIO
 from typing import Any, Protocol
 
+import httpx
+import pdfplumber
 from fastapi import UploadFile
 from nanoid import generate as generate_nanoid
 from sqlalchemy import Integer, select
@@ -30,8 +33,69 @@ class PDFProcessor:
         return content_type == "application/pdf"
 
     async def extract_text(self, content: bytes) -> str:
-        import pdfplumber
+        if settings.MISTRAL_OCR_API_KEY:
+            return await self._extract_text_with_mistral_ocr(content)
+        return self._extract_text_with_pdfplumber(content)
 
+    async def _extract_text_with_mistral_ocr(self, content: bytes) -> str:
+        api_key = settings.MISTRAL_OCR_API_KEY
+        if not api_key:
+            return self._extract_text_with_pdfplumber(content)
+
+        encoded_pdf = base64.b64encode(content).decode("ascii")
+        payload = {
+            "model": settings.MISTRAL_OCR_MODEL,
+            "document": {
+                "type": "document_url",
+                "document_url": f"data:application/pdf;base64,{encoded_pdf}",
+            },
+            "include_image_base64": False,
+        }
+
+        try:
+            async with httpx.AsyncClient(
+                timeout=settings.MISTRAL_OCR_TIMEOUT_SECONDS
+            ) as client:
+                response = await client.post(
+                    "https://api.mistral.ai/v1/ocr",
+                    headers={
+                        "Authorization": f"Bearer {api_key}",
+                        "Content-Type": "application/json",
+                    },
+                    json=payload,
+                )
+                response.raise_for_status()
+                ocr_result = response.json()
+        except httpx.HTTPError as exc:
+            logger.warning(
+                "Mistral OCR request failed; falling back to pdfplumber",
+                exc_info=exc,
+            )
+            return self._extract_text_with_pdfplumber(content)
+        except ValueError as exc:
+            logger.warning(
+                "Mistral OCR response was invalid; falling back to pdfplumber",
+                exc_info=exc,
+            )
+            return self._extract_text_with_pdfplumber(content)
+
+        pages = ocr_result.get("pages") if isinstance(ocr_result, dict) else None
+        if not isinstance(pages, list):
+            logger.warning(
+                "Mistral OCR response did not include pages; falling back to pdfplumber"
+            )
+            return self._extract_text_with_pdfplumber(content)
+
+        text_parts: list[str] = []
+        for page in pages:
+            if isinstance(page, dict):
+                markdown = page.get("markdown")
+                if isinstance(markdown, str) and markdown.strip():
+                    text_parts.append(markdown.strip())
+
+        return "\n\n".join(text_parts)
+
+    def _extract_text_with_pdfplumber(self, content: bytes) -> str:
         with pdfplumber.open(BytesIO(content)) as pdf_reader:
             text_parts: list[str] = []
             for page_num, page in enumerate(pdf_reader.pages):

tests/utils/test_files.py

@@ -1,9 +1,81 @@
 import json
+from typing import Any
 
+import httpx
 import pytest
 
+from src.config import settings
 from src.exceptions import ValidationException
-from src.utils.files import JSONProcessor
+from src.utils.files import JSONProcessor, PDFProcessor
+
+
+class _FakeMistralOCRResponse:
+    def raise_for_status(self) -> None:
+        return None
+
+    def json(self) -> dict[str, Any]:
+        return {
+            "pages": [
+                {"index": 0, "markdown": "# Page 1\nHello"},
+                {"index": 1, "markdown": "Page 2 text"},
+            ],
+            "usage_info": {"pages_processed": 2},
+        }
+
+
+class _FakeAsyncClient:
+    posted_json: dict[str, Any] | None = None
+    posted_headers: dict[str, str] | None = None
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        pass
+
+    async def __aenter__(self) -> "_FakeAsyncClient":
+        return self
+
+    async def __aexit__(self, *args: Any) -> None:
+        return None
+
+    async def post(
+        self,
+        url: str,
+        *,
+        headers: dict[str, str],
+        json: dict[str, Any],
+    ) -> _FakeMistralOCRResponse:
+        assert url == "https://api.mistral.ai/v1/ocr"
+        self.__class__.posted_headers = headers
+        self.__class__.posted_json = json
+        return _FakeMistralOCRResponse()
+
+
+class _FailingAsyncClient(_FakeAsyncClient):
+    async def post(
+        self,
+        url: str,
+        *,
+        headers: dict[str, str],
+        json: dict[str, Any],
+    ) -> _FakeMistralOCRResponse:
+        raise httpx.ConnectError("Mistral unavailable")
+
+
+class _FakePDFPage:
+    def __init__(self, text: str | None) -> None:
+        self._text = text
+
+    def extract_text(self) -> str | None:
+        return self._text
+
+
+class _FakePDFReader:
+    pages = [_FakePDFPage("First page"), _FakePDFPage(None), _FakePDFPage("Second page")]
+
+    def __enter__(self) -> "_FakePDFReader":
+        return self
+
+    def __exit__(self, *args: Any) -> None:
+        return None
 
 
 @pytest.mark.asyncio
@@ -37,3 +109,57 @@ async def test_json_processor_rejects_invalid_json_content():
 
     with pytest.raises(ValidationException, match="invalid"):
         await processor.extract_text(b'{"name": }')
+
+
+@pytest.mark.asyncio
+async def test_pdf_processor_extracts_markdown_with_mistral_ocr(monkeypatch):
+    processor = PDFProcessor()
+    _FakeAsyncClient.posted_json = None
+    _FakeAsyncClient.posted_headers = None
+    monkeypatch.setattr(settings, "MISTRAL_OCR_API_KEY", "test-mistral-key")
+    monkeypatch.setattr(settings, "MISTRAL_OCR_MODEL", "mistral-ocr-test")
+    monkeypatch.setattr(settings, "MISTRAL_OCR_TIMEOUT_SECONDS", 12.5)
+    monkeypatch.setattr(httpx, "AsyncClient", _FakeAsyncClient)
+
+    result = await processor.extract_text(b"%PDF test bytes")
+
+    assert result == "# Page 1\nHello\n\nPage 2 text"
+    assert _FakeAsyncClient.posted_headers == {
+        "Authorization": "Bearer test-mistral-key",
+        "Content-Type": "application/json",
+    }
+    assert _FakeAsyncClient.posted_json == {
+        "model": "mistral-ocr-test",
+        "document": {
+            "type": "document_url",
+            "document_url": "data:application/pdf;base64,JVBERiB0ZXN0IGJ5dGVz",
+        },
+        "include_image_base64": False,
+    }
+
+
+@pytest.mark.asyncio
+async def test_pdf_processor_falls_back_to_pdfplumber_without_mistral_key(
+    monkeypatch,
+):
+    processor = PDFProcessor()
+    monkeypatch.setattr(settings, "MISTRAL_OCR_API_KEY", None)
+    monkeypatch.setattr("src.utils.files.pdfplumber.open", lambda *args: _FakePDFReader())
+
+    result = await processor.extract_text(b"%PDF test bytes")
+
+    assert result == "[Page 1]\nFirst page\n\n[Page 3]\nSecond page"
+
+
+@pytest.mark.asyncio
+async def test_pdf_processor_falls_back_to_pdfplumber_when_mistral_fails(
+    monkeypatch,
+):
+    processor = PDFProcessor()
+    monkeypatch.setattr(settings, "MISTRAL_OCR_API_KEY", "test-mistral-key")
+    monkeypatch.setattr(httpx, "AsyncClient", _FailingAsyncClient)
+    monkeypatch.setattr("src.utils.files.pdfplumber.open", lambda *args: _FakePDFReader())
+
+    result = await processor.extract_text(b"%PDF test bytes")
+
+    assert result == "[Page 1]\nFirst page\n\n[Page 3]\nSecond page"