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

vlmparse/clients/dotsocr.py

@@ -1,14 +1,13 @@
 import json
 import math
 from pathlib import Path
-from typing import ClassVar, Literal
+from typing import ClassVar
 
 from loguru import logger
 from PIL import Image
 from pydantic import Field
 
 from vlmparse.clients.openai_converter import (
-    LLMParams,
     OpenAIConverterClient,
     OpenAIConverterConfig,
 )
@@ -48,12 +47,13 @@ class DotsOCRDockerServerConfig(DockerServerConfig):
     )
     add_model_key_to_server: bool = True
     aliases: list[str] = Field(default_factory=lambda: ["dotsocr"])
+    default_model_name: str = DEFAULT_MODEL_NAME
 
     @property
     def client_config(self):
         return DotsOCRConverterConfig(
-            llm_params=LLMParams(
-                base_url=f"http://localhost:{self.docker_port}{self.get_base_url_suffix()}",
+            **self._create_client_kwargs(
+                f"http://localhost:{self.docker_port}{self.get_base_url_suffix()}"
             )
         )
 
@@ -65,29 +65,7 @@ class DotsOCRConverterConfig(OpenAIConverterConfig):
     model_name: str = "rednote-hilab/dots.ocr"
     preprompt: str | None = ""
     postprompt: str | None = None
-    completion_kwargs: dict | None = {
-        "temperature": 0.1,
-        "top_p": 1.0,
-        "max_completion_tokens": 16384,
-    }
-    aliases: list[str] = Field(default_factory=lambda: ["dotsocr"])
-    dpi: int = 200
-    prompt_mode: Literal["prompt_layout_all_en", "prompt_ocr"] = "prompt_ocr"
-
-    def get_client(self, **kwargs) -> "DotsOCRConverter":
-        return DotsOCRConverter(config=self, **kwargs)
-
-
-class DotsOCRConverter(OpenAIConverterClient):
-    """DotsOCR VLLM converter."""
-
-    # Constants
-    MIN_PIXELS: ClassVar[int] = 3136
-    MAX_PIXELS: ClassVar[int] = 11289600
-    IMAGE_FACTOR: ClassVar[int] = 28
-
-    # Prompts
-    PROMPTS: ClassVar[dict] = {
+    prompts: dict[str, str] = {
         "prompt_layout_all_en": """Please output the layout information from the PDF image, including each layout element's bbox, its category, and the corresponding text content within the bbox.
 
 1. Bbox format: [x1, y1, x2, y2]
@@ -108,6 +86,30 @@ class DotsOCRConverter(OpenAIConverterClient):
 """,
         "prompt_ocr": """Extract the text content from this image.""",
     }
+    prompt_mode_map: dict[str, str] = {
+        "ocr": "prompt_ocr",
+        "ocr_layout": "prompt_layout_all_en",
+        "table": "prompt_layout_all_en",
+    }
+    completion_kwargs: dict | None = {
+        "temperature": 0.1,
+        "top_p": 1.0,
+        "max_completion_tokens": 16384,
+    }
+    aliases: list[str] = Field(default_factory=lambda: ["dotsocr"])
+    dpi: int = 200
+
+    def get_client(self, **kwargs) -> "DotsOCRConverter":
+        return DotsOCRConverter(config=self, **kwargs)
+
+
+class DotsOCRConverter(OpenAIConverterClient):
+    """DotsOCR VLLM converter."""
+
+    # Constants
+    MIN_PIXELS: ClassVar[int] = 3136
+    MAX_PIXELS: ClassVar[int] = 11289600
+    IMAGE_FACTOR: ClassVar[int] = 28
 
     @staticmethod
     def round_by_factor(number: int, factor: int) -> int:
@@ -235,7 +237,7 @@ class DotsOCRConverter(OpenAIConverterClient):
         image = self.fetch_image(
             origin_image, min_pixels=self.MIN_PIXELS, max_pixels=self.MAX_PIXELS
         )
-        prompt = self.PROMPTS[prompt_mode]
+        prompt = self.config.prompts[prompt_mode]
 
         response, usage = await self._async_inference_with_vllm(image, prompt)
 
@@ -258,13 +260,15 @@ class DotsOCRConverter(OpenAIConverterClient):
     async def async_call_inside_page(self, page: Page) -> Page:
         image = page.image
 
+        prompt_key = self.get_prompt_key() or "prompt_ocr"
+
         _, response, _, usage = await self._parse_image_vllm(
-            image, prompt_mode=self.config.prompt_mode
+            image, prompt_mode=prompt_key
         )
         logger.info("Response: " + str(response))
 
         items = None
-        if self.config.prompt_mode == "prompt_layout_all_en":
+        if prompt_key == "prompt_layout_all_en":
             text = "\n\n".join([item.get("text", "") for item in response])
 
             items = []
@@ -286,5 +290,4 @@ class DotsOCRConverter(OpenAIConverterClient):
 
         page.completion_tokens = usage.completion_tokens
         page.prompt_tokens = usage.prompt_tokens
-        page.reasoning_tokens = usage.reasoning_tokens
         return page

vlmparse/clients/granite_docling.py

@@ -28,12 +28,17 @@ 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):
     """Granite Docling converter configuration."""
 
+    model_name: str = "ibm-granite/granite-docling-258M"
     preprompt: str | None = None
     postprompt: str | None = "Convert this page to docling."
     completion_kwargs: dict | None = {
@@ -69,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
-from typing import Literal
+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,25 +79,67 @@ 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
-    ) -> tuple[str, "CompletionUsage"]:  # noqa: F821
+    ) -> tuple[str, Optional["CompletionUsage"]]:  # noqa: F821
         """Helper to handle chat completion with optional streaming."""
         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,
@@ -117,10 +148,11 @@ class OpenAIConverterClient(BaseConverter):
             async for chunk in response_stream:
                 if chunk.choices and chunk.choices[0].delta.content:
                     response_parts.append(chunk.choices[0].delta.content)
-            return "".join(response_parts)
+
+            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,
             )
@@ -146,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,
@@ -175,9 +211,10 @@ class OpenAIConverterClient(BaseConverter):
 
         text = html_to_md_keep_tables(text)
         page.text = text
-        page.prompt_tokens = usage.prompt_tokens
-        page.completion_tokens = usage.completion_tokens
-        if hasattr(usage, "reasoning_tokens"):
-            page.reasoning_tokens = usage.reasoning_tokens
+        if usage is not None:
+            page.prompt_tokens = usage.prompt_tokens
+            page.completion_tokens = usage.completion_tokens
+            if hasattr(usage, "reasoning_tokens"):
+                page.reasoning_tokens = usage.reasoning_tokens
 
         return page

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/constants.py

@@ -1,2 +1,5 @@
+import os
+
 IMAGE_EXTENSIONS = [".jpg", ".jpeg", ".png", ".tiff", ".tif", ".bmp", ".gif", ".webp"]
 PDF_EXTENSION = ".pdf"
+DEFAULT_SERVER_PORT = os.getenv("VLMPARSE_DEFAULT_PORT", 8056)