diff --git a/.env.template b/.env.template index 123af642..77bc551e 100644 --- a/.env.template +++ b/.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 diff --git a/config.toml.example b/config.toml.example index 236f3402..7a1013d0 100644 --- a/config.toml.example +++ b/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 diff --git a/src/config.py b/src/config.py index cae0c5de..6e7cd1bc 100644 --- a/src/config.py +++ b/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 diff --git a/src/utils/files.py b/src/utils/files.py index cdbac514..a31409e8 100644 --- a/src/utils/files.py +++ b/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): diff --git a/tests/utils/test_files.py b/tests/utils/test_files.py index 0b577be7..39cb67da 100644 --- a/tests/utils/test_files.py +++ b/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"