diff --git a/projects/pgai/pgai/vectorizer/parsing.py b/projects/pgai/pgai/vectorizer/parsing.py index fb86b65d6..3f38f851f 100644 --- a/projects/pgai/pgai/vectorizer/parsing.py +++ b/projects/pgai/pgai/vectorizer/parsing.py @@ -1,5 +1,7 @@ +import asyncio import os from abc import ABC, abstractmethod +from concurrent.futures import ThreadPoolExecutor from pathlib import Path from typing import Any, Literal @@ -8,11 +10,17 @@ from pgai.vectorizer.loading import LoadedDocument +# Thread pool for CPU-intensive parsing operations +max_workers = int(os.getenv("PARSING_MAX_WORKERS", 4)) +_PARSING_EXECUTOR = ThreadPoolExecutor( + max_workers=max_workers, thread_name_prefix="parsing" +) + class ParsingNone(BaseModel): implementation: Literal["none"] - def parse(self, _1: dict[str, Any], payload: str | LoadedDocument) -> str: # noqa: ARG002 + async def parse(self, _1: dict[str, Any], payload: str | LoadedDocument) -> str: # noqa: ARG002 if isinstance(payload, LoadedDocument): raise ValueError( "Cannot chunk Document with parsing_none, " @@ -24,13 +32,15 @@ def parse(self, _1: dict[str, Any], payload: str | LoadedDocument) -> str: # no class ParsingAuto(BaseModel): implementation: Literal["auto"] - def parse(self, row: dict[str, Any], payload: str | LoadedDocument) -> str: + async def parse(self, row: dict[str, Any], payload: str | LoadedDocument) -> str: if isinstance(payload, LoadedDocument): if payload.file_type == "epub": # epub is not supported by docling, but by pymupdf - return ParsingPyMuPDF(implementation="pymupdf").parse(row, payload) + return await ParsingPyMuPDF(implementation="pymupdf").parse( + row, payload + ) - return ParsingDocling(implementation="docling").parse(row, payload) + return await ParsingDocling(implementation="docling").parse(row, payload) else: return payload @@ -40,7 +50,7 @@ class BaseDocumentParsing(BaseModel, ABC): implementation: str - def parse(self, row: dict[str, Any], payload: LoadedDocument | str) -> str: + async def parse(self, row: dict[str, Any], payload: LoadedDocument | str) -> str: """ Parse a document payload into a string representation. @@ -66,10 +76,10 @@ def parse(self, row: dict[str, Any], payload: LoadedDocument | str) -> str: if payload.file_type in ["txt", "md"]: return payload.content.getvalue().decode("utf-8") - return self.parse_doc(row, payload) + return await self.parse_doc(row, payload) @abstractmethod - def parse_doc(self, row: dict[str, Any], payload: LoadedDocument) -> str: + async def parse_doc(self, row: dict[str, Any], payload: LoadedDocument) -> str: """ Parse a binary document into a string representation, Markdown preferable. Must be implemented by subclasses. @@ -82,7 +92,15 @@ class ParsingPyMuPDF(BaseDocumentParsing): implementation: Literal["pymupdf"] # type: ignore[reportIncompatibleVariableOverride] @override - def parse_doc(self, row: dict[str, Any], payload: LoadedDocument) -> str: # noqa: ARG002 + async def parse_doc(self, row: dict[str, Any], payload: LoadedDocument) -> str: # noqa: ARG002 + # Run blocking parsing operation in thread pool + loop = asyncio.get_running_loop() + return await loop.run_in_executor( + _PARSING_EXECUTOR, self._parse_with_pymupdf, payload + ) + + def _parse_with_pymupdf(self, payload: LoadedDocument) -> str: + """Synchronous pymupdf parsing to run in thread pool.""" # Note: deferred import to avoid import overhead import pymupdf # type: ignore import pymupdf4llm # type: ignore @@ -105,7 +123,15 @@ class ParsingDocling(BaseDocumentParsing): cache_dir: Path | str = DOCLING_CACHE_DIR @override - def parse_doc(self, row: dict[str, Any], payload: LoadedDocument) -> str: # noqa: ARG002 + async def parse_doc(self, row: dict[str, Any], payload: LoadedDocument) -> str: # noqa: ARG002 + # Run blocking parsing operation in thread pool + loop = asyncio.get_running_loop() + return await loop.run_in_executor( + _PARSING_EXECUTOR, self._parse_with_docling, payload + ) + + def _parse_with_docling(self, payload: LoadedDocument) -> str: + """Synchronous docling parsing to run in thread pool.""" # Note: deferred import to avoid import overhead from docling.datamodel.base_models import ( DocumentStream, # type: ignore diff --git a/projects/pgai/pgai/vectorizer/vectorizer.py b/projects/pgai/pgai/vectorizer/vectorizer.py index a39c05335..0a797914b 100644 --- a/projects/pgai/pgai/vectorizer/vectorizer.py +++ b/projects/pgai/pgai/vectorizer/vectorizer.py @@ -1188,7 +1188,7 @@ async def _generate_embeddings( loading_errors.append((item, (LoadingError(e=e)))) continue - payload = self.vectorizer.config.parsing.parse(item, payload) + payload = await self.vectorizer.config.parsing.parse(item, payload) chunks = self.vectorizer.config.chunking.into_chunks(item, payload) for chunk_id, chunk in enumerate(chunks, 0): formatted = self.vectorizer.config.formatting.format(chunk, item)