vlmparse 0.1.7__py3-none-any.whl → 0.1.8__py3-none-any.whl

vlmparse/clients/granite_docling.py

@@ -28,7 +28,11 @@ class GraniteDoclingDockerServerConfig(VLLMDockerServerConfig):
 
     @property
     def client_config(self):
-        return GraniteDoclingConverterConfig(llm_params=self.llm_params)
+        return GraniteDoclingConverterConfig(
+            **self._create_client_kwargs(
+                f"http://localhost:{self.docker_port}{self.get_base_url_suffix()}"
+            )
+        )
 
 
 class GraniteDoclingConverterConfig(OpenAIConverterConfig):
@@ -70,49 +74,18 @@ class GraniteDoclingConverter(OpenAIConverterClient):
             }
         ]
 
-        doctags = await self._get_chat_completion_adaptive(
+        doctags, usage = await self._get_chat_completion(
             messages, completion_kwargs=self.config.completion_kwargs
         )
         doctags = clean_response(doctags)
 
         page.raw_response = doctags
         page.text = _doctags_to_markdown(doctags, image)
+        if usage is not None:
+            page.prompt_tokens = usage.prompt_tokens
+            page.completion_tokens = usage.completion_tokens
         return page
 
-    async def _get_chat_completion_adaptive(
-        self, messages: list[dict], completion_kwargs: dict | None
-    ) -> str:
-        """
-        vLLM enforces input+output <= model context length. If `max_tokens` is too
-        high (especially for multimodal prompts), retry with progressively smaller
-        `max_tokens`.
-        """
-        kwargs = (completion_kwargs or {}).copy()
-        max_tokens = kwargs.get("max_tokens") or kwargs.get("max_completion_tokens")
-
-        for _ in range(6):
-            try:
-                return await self._get_chat_completion(
-                    messages, completion_kwargs=kwargs
-                )
-            except Exception as e:
-                msg = str(e)
-                too_large = (
-                    "max_tokens" in msg
-                    and "maximum context length" in msg
-                    and "is too large" in msg
-                )
-                if not too_large or not isinstance(max_tokens, int):
-                    raise
-
-                max_tokens = max(256, int(max_tokens * 0.75))
-                if "max_tokens" in kwargs:
-                    kwargs["max_tokens"] = max_tokens
-                if "max_completion_tokens" in kwargs:
-                    kwargs["max_completion_tokens"] = max_tokens
-
-        return await self._get_chat_completion(messages, completion_kwargs=kwargs)
-
 
 def _doctags_to_markdown(doctags: str, image):
     try:

vlmparse/clients/hunyuanocr.py

@@ -25,7 +25,11 @@ class HunyuanOCRDockerServerConfig(VLLMDockerServerConfig):
 
     @property
     def client_config(self):
-        return HunyuanOCRConverterConfig(llm_params=self.llm_params)
+        return HunyuanOCRConverterConfig(
+            **self._create_client_kwargs(
+                f"http://localhost:{self.docker_port}{self.get_base_url_suffix()}"
+            )
+        )
 
 
 class HunyuanOCRConverterConfig(OpenAIConverterConfig):

vlmparse/clients/lightonocr.py

@@ -25,7 +25,11 @@ class LightOnOCRDockerServerConfig(VLLMDockerServerConfig):
 
     @property
     def client_config(self):
-        return LightOnOCRConverterConfig(llm_params=self.llm_params)
+        return LightOnOCRConverterConfig(
+            **self._create_client_kwargs(
+                f"http://localhost:{self.docker_port}{self.get_base_url_suffix()}"
+            )
+        )
 
 
 class LightOnOCRConverterConfig(OpenAIConverterConfig):
@@ -41,3 +45,21 @@ class LightOnOCRConverterConfig(OpenAIConverterConfig):
     }
     dpi: int = 200
     aliases: list[str] = Field(default_factory=lambda: ["lightonocr"])
+
+
+class LightonOCR21BServerConfig(LightOnOCRDockerServerConfig):
+    model_name: str = "lightonai/LightOnOCR-2-1B"
+    aliases: list[str] = Field(default_factory=lambda: ["lightonocr2"])
+
+    @property
+    def client_config(self):
+        return LightonOCR21BConverterConfig(
+            **self._create_client_kwargs(
+                f"http://localhost:{self.docker_port}{self.get_base_url_suffix()}"
+            )
+        )
+
+
+class LightonOCR21BConverterConfig(LightOnOCRConverterConfig):
+    model_name: str = "lightonai/LightOnOCR-2-1B"
+    aliases: list[str] = Field(default_factory=lambda: ["lightonocr2"])

vlmparse/clients/mineru.py

@@ -31,7 +31,6 @@ class MinerUDockerServerConfig(DockerServerConfig):
 class MinerUConverterConfig(ConverterConfig):
     """Configuration for MinerU API converter."""
 
-    base_url: str
     model_name: str = "opendatalab/MinerU2.5-2509-1.2B"
     aliases: list[str] = Field(default_factory=lambda: ["mineru25"])
     timeout: int = 600

vlmparse/clients/mistral_converter.py

@@ -0,0 +1,85 @@
+import os
+
+import httpx
+import orjson
+from loguru import logger
+from pydantic import Field
+
+from vlmparse.clients.pipe_utils.html_to_md_conversion import html_to_md_keep_tables
+from vlmparse.clients.pipe_utils.utils import clean_response
+from vlmparse.converter import BaseConverter, ConverterConfig
+from vlmparse.data_model.document import Page
+from vlmparse.utils import to_base64
+
+
+class MistralOCRConverterConfig(ConverterConfig):
+    """Configuration for Mistral OCR converter."""
+
+    base_url: str = "https://api.mistral.ai/v1"
+    model_name: str = "mistral-ocr-latest"
+    api_key: str | None = None
+    timeout: int = 300
+    aliases: list[str] = Field(
+        default_factory=lambda: ["mistral-ocr-latest", "mistral-ocr"]
+    )
+
+    def get_client(self, **kwargs) -> "MistralOCRConverter":
+        return MistralOCRConverter(config=self, **kwargs)
+
+
+class MistralOCRConverter(BaseConverter):
+    """Client for Mistral OCR API."""
+
+    config: MistralOCRConverterConfig
+
+    def __init__(self, config: MistralOCRConverterConfig, **kwargs):
+        super().__init__(config=config, **kwargs)
+        if not self.config.api_key:
+            self.config.api_key = os.getenv("MISTRAL_API_KEY")
+        if not self.config.api_key:
+            raise ValueError("MISTRAL_API_KEY environment variable not set")
+        self._base_url = self.config.base_url.rstrip("/")
+
+    async def _async_ocr(self, image) -> httpx.Response:
+        payload = {
+            "model": self.config.model_name,
+            "document": {
+                "type": "image_url",
+                "image_url": f"data:image/png;base64,{to_base64(image)}",
+            },
+        }
+        headers = {"Authorization": f"Bearer {self.config.api_key}"}
+
+        async with httpx.AsyncClient(timeout=self.config.timeout) as client:
+            response = await client.post(
+                f"{self._base_url}/ocr",
+                json=payload,
+                headers=headers,
+            )
+        response.raise_for_status()
+        return response
+
+    async def async_call_inside_page(self, page: Page) -> Page:
+        response = await self._async_ocr(page.image)
+        page.raw_response = response.text
+
+        try:
+            data = response.json()
+        except ValueError:
+            logger.warning("Mistral OCR returned non-JSON response")
+            page.text = clean_response(response.text)
+            return page
+
+        pages = data.get("pages") or []
+        if pages:
+            page_data = pages[0]
+            text = page_data.get("markdown") or page_data.get("text") or ""
+        else:
+            text = (
+                data.get("markdown") or data.get("text") or orjson.dumps(data).decode()
+            )
+
+        text = clean_response(text)
+        text = html_to_md_keep_tables(text)
+        page.text = text
+        return page

vlmparse/clients/nanonetocr.py

@@ -12,7 +12,11 @@ class NanonetOCR2DockerServerConfig(VLLMDockerServerConfig):
 
     @property
     def client_config(self):
-        return NanonetOCR2ConverterConfig(llm_params=self.llm_params)
+        return NanonetOCR2ConverterConfig(
+            **self._create_client_kwargs(
+                f"http://localhost:{self.docker_port}{self.get_base_url_suffix()}"
+            )
+        )
 
 
 class NanonetOCR2ConverterConfig(OpenAIConverterConfig):

vlmparse/clients/olmocr.py

@@ -23,7 +23,11 @@ class OlmOCRDockerServerConfig(VLLMDockerServerConfig):
 
     @property
     def client_config(self):
-        return OlmOCRConverterConfig(llm_params=self.llm_params)
+        return OlmOCRConverterConfig(
+            **self._create_client_kwargs(
+                f"http://localhost:{self.docker_port}{self.get_base_url_suffix()}"
+            )
+        )
 
 
 class OlmOCRConverterConfig(OpenAIConverterConfig):
@@ -37,7 +41,7 @@ class OlmOCRConverterConfig(OpenAIConverterConfig):
         "Return your output as markdown, with a front matter section on top specifying values for the primary_language, is_rotation_valid, rotation_correction, is_table, and is_diagram parameters."
     )
     postprompt: str | None = None
-    completion_kwargs: dict | None = {
+    completion_kwargs: dict = {
         "temperature": 0.1,
         "max_tokens": 8000,
     }

vlmparse/clients/openai_converter.py

@@ -1,15 +1,13 @@
-import os
+import asyncio
 from typing import Literal, Optional
 
 from loguru import logger
 from pydantic import Field
 
-from vlmparse.base_model import VLMParseBaseModel
 from vlmparse.clients.pipe_utils.html_to_md_conversion import html_to_md_keep_tables
 from vlmparse.clients.pipe_utils.utils import clean_response
 from vlmparse.converter import BaseConverter, ConverterConfig
 from vlmparse.data_model.document import Page
-from vlmparse.servers.docker_server import DEFAULT_MODEL_NAME
 from vlmparse.utils import to_base64
 
 from .prompts import PDF2MD_PROMPT
@@ -17,50 +15,14 @@ from .prompts import PDF2MD_PROMPT
 GOOGLE_API_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/openai/"
 
 
-class LLMParams(VLMParseBaseModel):
+class OpenAIConverterConfig(ConverterConfig):
     api_key: str = ""
-    base_url: str | None = None
-    model_name: str = DEFAULT_MODEL_NAME
     timeout: int | None = 500
     max_retries: int = 1
-
-
-def get_llm_params(model_name: str, uri: str | None = None):
-    if uri is not None:
-        return LLMParams(base_url=uri, model_name="vllm-model", api_key="")
-    if model_name in [
-        "gpt-4o",
-        "gpt-4o-mini",
-        "gpt-4.1",
-        "gpt-4.1-mini",
-        "gpt-4.1-nano",
-        "gpt-5",
-        "gpt-5-mini",
-        "gpt-5-nano",
-    ]:
-        base_url = None
-        api_key = os.getenv("OPENAI_API_KEY")
-        if api_key is None:
-            raise ValueError("OPENAI_API_KEY environment variable not set")
-    else:
-        if model_name in [
-            "gemini-2.5-flash-lite",
-            "gemini-2.5-flash",
-            "gemini-2.5-pro",
-        ]:
-            base_url = GOOGLE_API_BASE_URL
-            api_key = os.getenv("GOOGLE_API_KEY")
-            if api_key is None:
-                raise ValueError("GOOGLE_API_KEY environment variable not set")
-        else:
-            return None
-    return LLMParams(base_url=base_url, model_name=model_name, api_key=api_key)
-
-
-class OpenAIConverterConfig(ConverterConfig):
-    llm_params: LLMParams
     preprompt: str | None = None
-    postprompt: str | None = PDF2MD_PROMPT
+    postprompt: str | dict[str, str] | None = PDF2MD_PROMPT
+    prompts: dict[str, str] = Field(default_factory=dict)
+    prompt_mode_map: dict[str, str] = Field(default_factory=dict)
     completion_kwargs: dict = Field(default_factory=dict)
     stream: bool = False
 
@@ -71,6 +33,33 @@ class OpenAIConverterConfig(ConverterConfig):
 class OpenAIConverterClient(BaseConverter):
     """Client for OpenAI-compatible API servers."""
 
+    def get_prompt_key(self) -> str | None:
+        """Resolve a prompt key from conversion_mode using class mappings."""
+        mode = getattr(self.config, "conversion_mode", None) or "ocr"
+        prompts = self._get_prompts()
+        if mode in prompts:
+            return mode
+        mapped = self._get_prompt_mode_map().get(mode)
+        if mapped in prompts:
+            return mapped
+        return None
+
+    def get_prompt_for_mode(self) -> str | None:
+        key = self.get_prompt_key()
+        if key is None:
+            return None
+        return self._get_prompts().get(key)
+
+    def _get_prompts(self) -> dict[str, str]:
+        if self.config.prompts:
+            return self.config.prompts
+        if isinstance(self.config.postprompt, dict):
+            return self.config.postprompt
+        return {}
+
+    def _get_prompt_mode_map(self) -> dict[str, str]:
+        return self.config.prompt_mode_map or {}
+
     def __init__(
         self,
         config: OpenAIConverterConfig,
@@ -90,14 +79,54 @@ class OpenAIConverterClient(BaseConverter):
             debug=debug,
             return_documents_in_batch_mode=return_documents_in_batch_mode,
         )
-        from openai import AsyncOpenAI
-
-        self.model = AsyncOpenAI(
-            base_url=self.config.llm_params.base_url,
-            api_key=self.config.llm_params.api_key,
-            timeout=self.config.llm_params.timeout,
-            max_retries=self.config.llm_params.max_retries,
-        )
+        self._model = None
+        self._model_loop = None
+
+    async def _get_async_model(self):
+        loop = asyncio.get_running_loop()
+        if self._model is None or self._model_loop is not loop:
+            await self._close_model()
+            from openai import AsyncOpenAI
+
+            self._model = AsyncOpenAI(
+                base_url=self.config.base_url,
+                api_key=self.config.api_key,
+                timeout=self.config.timeout,
+                max_retries=self.config.max_retries,
+            )
+            self._model_loop = loop
+        return self._model
+
+    async def _close_model(self):
+        """Close the async OpenAI client if it exists."""
+        if self._model is not None:
+            try:
+                await self._model.close()
+            except RuntimeError:
+                # Event loop may already be closed
+                pass
+            finally:
+                self._model = None
+                self._model_loop = None
+
+    async def aclose(self):
+        """Close the converter and release resources."""
+        await self._close_model()
+
+    def close(self):
+        """Synchronously close the converter if possible."""
+        if self._model is not None:
+            try:
+                loop = asyncio.get_running_loop()
+                loop.create_task(self._close_model())
+            except RuntimeError:
+                # No running loop, try to close synchronously
+                try:
+                    asyncio.run(self._close_model())
+                except RuntimeError:
+                    # Event loop already closed, force cleanup
+                    self._model = None
+                    self._model_loop = None
 
     async def _get_chat_completion(
         self, messages: list[dict], completion_kwargs: dict | None = None
@@ -106,9 +135,11 @@ class OpenAIConverterClient(BaseConverter):
         if completion_kwargs is None:
             completion_kwargs = self.config.completion_kwargs
 
+        model = await self._get_async_model()
+
         if self.config.stream:
-            response_stream = await self.model.chat.completions.create(
-                model=self.config.llm_params.model_name,
+            response_stream = await model.chat.completions.create(
+                model=self.config.default_model_name,
                 messages=messages,
                 stream=True,
                 **completion_kwargs,
@@ -120,8 +151,8 @@ class OpenAIConverterClient(BaseConverter):
 
             return "".join(response_parts), None
         else:
-            response_obj = await self.model.chat.completions.create(
-                model=self.config.llm_params.model_name,
+            response_obj = await model.chat.completions.create(
+                model=self.config.default_model_name,
                 messages=messages,
                 **completion_kwargs,
             )
@@ -147,11 +178,15 @@ class OpenAIConverterClient(BaseConverter):
         else:
             preprompt = []
 
-        postprompt = (
-            [{"type": "text", "text": self.config.postprompt}]
-            if self.config.postprompt
-            else []
-        )
+        selected_prompt = self.get_prompt_for_mode()
+        if selected_prompt is not None:
+            postprompt = [{"type": "text", "text": selected_prompt}]
+        else:
+            postprompt = (
+                [{"type": "text", "text": self.config.postprompt}]
+                if isinstance(self.config.postprompt, str) and self.config.postprompt
+                else []
+            )
 
         messages = [
             *preprompt,

vlmparse/clients/paddleocrvl.py

@@ -22,7 +22,11 @@ class PaddleOCRVLDockerServerConfig(VLLMDockerServerConfig):
 
     @property
     def client_config(self):
-        return PaddleOCRVLConverterConfig(llm_params=self.llm_params)
+        return PaddleOCRVLConverterConfig(
+            **self._create_client_kwargs(
+                f"http://localhost:{self.docker_port}{self.get_base_url_suffix()}"
+            )
+        )
 
 
 # Task-specific base prompts
@@ -39,7 +43,10 @@ class PaddleOCRVLConverterConfig(OpenAIConverterConfig):
 
     model_name: str = "PaddlePaddle/PaddleOCR-VL"
     preprompt: str | None = None
-    postprompt: str | None = TASKS["ocr"]
+    postprompt: dict[str, str] = TASKS
+    prompt_mode_map: dict[str, str] = {
+        "ocr_layout": "ocr",
+    }
     completion_kwargs: dict | None = {
         "temperature": 0.0,
         "max_completion_tokens": 16384,

vlmparse/converter.py

@@ -1,7 +1,6 @@
 import asyncio
 import threading
 import time
-import traceback
 from pathlib import Path
 from typing import Literal
 
@@ -9,6 +8,8 @@ from loguru import logger
 from PIL import Image
 from pydantic import Field
 
+from vlmparse.servers.docker_server import DEFAULT_MODEL_NAME
+
 from .base_model import VLMParseBaseModel
 from .build_doc import convert_specific_page_to_image, get_page_count, resize_image
 from .constants import IMAGE_EXTENSIONS, PDF_EXTENSION
@@ -19,9 +20,20 @@ PDFIUM_LOCK = threading.Lock()
 
 
 class ConverterConfig(VLMParseBaseModel):
+    model_name: str
     aliases: list[str] = Field(default_factory=list)
-    dpi: int = 175
-    max_image_size: int | None = 4000
+    dpi: int = Field(default=175, ge=30, le=600)
+    max_image_size: int | None = Field(default=4000, ge=50)
+    base_url: str | None = None
+    default_model_name: str = DEFAULT_MODEL_NAME
+    conversion_mode: Literal[
+        "ocr",
+        "ocr_layout",
+        "table",
+        "image_description",
+        "formula",
+        "chart",
+    ] = "ocr"
 
     def get_client(self, **kwargs) -> "BaseConverter":
         return BaseConverter(config=self, **kwargs)
@@ -94,14 +106,22 @@ class BaseConverter:
                         page = await self.async_call_inside_page(page)
                         toc = time.perf_counter()
                         page.latency = toc - tic
-                        logger.debug(f"Time taken: {page.latency} seconds")
+                        logger.debug(
+                            "Page {page_idx} processed in {latency:.2f}s",
+                            page_idx=page_idx,
+                            latency=page.latency,
+                        )
                     except KeyboardInterrupt:
                         raise
                     except Exception:
                         if self.debug:
                             raise
                         else:
-                            logger.exception(traceback.format_exc())
+                            logger.opt(exception=True).error(
+                                "Error processing page {page_idx} of {file_path}",
+                                page_idx=page_idx,
+                                file_path=str(file_path),
+                            )
                             page.error = ProcessingError.from_class(self)
                     if not self.save_page_images:
                         page.buffer_image = dict(
@@ -122,12 +142,19 @@ class BaseConverter:
             if self.debug:
                 raise
             else:
-                logger.exception(traceback.format_exc())
+                logger.opt(exception=True).error(
+                    "Error processing document {file_path}",
+                    file_path=str(file_path),
+                )
                 document.error = ProcessingError.from_class(self)
                 return document
         toc = time.perf_counter()
         document.latency = toc - tic
-        logger.debug(f"Time taken to process the document: {document.latency} seconds")
+        logger.debug(
+            "Document {file_path} processed in {latency:.2f}s",
+            file_path=str(file_path),
+            latency=document.latency,
+        )
         if self.save_folder is not None:
             self._save_document(document)
 
@@ -169,8 +196,16 @@ class BaseConverter:
         else:
             logger.warning(f"Unknown save_mode: {self.save_mode}, skipping save")
 
+    async def _async_call_with_cleanup(self, file_path: str | Path):
+        """Call async_call and ensure cleanup."""
+        try:
+            return await self.async_call(file_path)
+        finally:
+            if hasattr(self, "aclose"):
+                await self.aclose()
+
     def __call__(self, file_path: str | Path):
-        return asyncio.run(self.async_call(file_path))
+        return asyncio.run(self._async_call_with_cleanup(file_path))
 
     async def async_batch(self, file_paths: list[str | Path]) -> list[Document] | None:
         """Process multiple files concurrently with semaphore limit."""
@@ -184,9 +219,14 @@ class BaseConverter:
                     await self.async_call(file_path)
 
         tasks = [asyncio.create_task(worker(file_path)) for file_path in file_paths]
-        documents = await asyncio.gather(*tasks)
-        if self.return_documents_in_batch_mode:
-            return documents
+        try:
+            documents = await asyncio.gather(*tasks)
+            if self.return_documents_in_batch_mode:
+                return documents
+        finally:
+            # Close async resources before the event loop ends
+            if hasattr(self, "aclose"):
+                await self.aclose()
 
     def batch(self, file_paths: list[str | Path]) -> list[Document] | None:
         """Synchronous wrapper for async_batch."""