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

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."""

vlmparse/converter_with_server.py

@@ -5,10 +5,61 @@ from typing import Literal
 
 from loguru import logger
 
+from vlmparse.constants import DEFAULT_SERVER_PORT
 from vlmparse.servers.utils import get_model_from_uri
 from vlmparse.utils import get_file_paths
 
 
+def start_server(
+    model: str,
+    gpus: str,
+    port: None | int = None,
+    with_vllm_server: bool = True,
+    vllm_args: list[str] = {},
+    forget_predefined_vllm_args: bool = False,
+    auto_stop: bool = False,
+):
+    from vlmparse.registries import docker_config_registry
+
+    base_url = ""
+    container = None
+    docker_config = docker_config_registry.get(model, default=with_vllm_server)
+
+    if port is None:
+        port = DEFAULT_SERVER_PORT
+
+    if docker_config is None:
+        logger.warning(
+            f"No Docker configuration found for model: {model}, using default configuration"
+        )
+        return "", container, None, docker_config
+
+    gpu_device_ids = None
+    if gpus is not None:
+        gpu_device_ids = [g.strip() for g in str(gpus).split(",")]
+
+    if docker_config is not None:
+        if port is not None:
+            docker_config.docker_port = port
+        docker_config.gpu_device_ids = gpu_device_ids
+        docker_config.update_command_args(
+            vllm_args,
+            forget_predefined_vllm_args=forget_predefined_vllm_args,
+        )
+
+        logger.info(
+            f"Deploying VLLM server for {docker_config.model_name} on port {port}..."
+        )
+        server = docker_config.get_server(auto_stop=auto_stop)
+        if server is None:
+            logger.error(f"Model server not found for model: {model}")
+            return "", container, None, docker_config
+
+        base_url, container = server.start()
+
+    return base_url, container, server, docker_config
+
+
 class ConverterWithServer:
     def __init__(
         self,
@@ -18,58 +69,60 @@ class ConverterWithServer:
         port: int | None = None,
         with_vllm_server: bool = False,
         concurrency: int = 10,
-        vllm_kwargs: dict | None = None,
-        forget_predefined_vllm_kwargs: bool = False,
+        vllm_args: dict | None = None,
+        forget_predefined_vllm_args: bool = False,
+        return_documents: bool = False,
     ):
+        if model is None and uri is None:
+            raise ValueError("Either 'model' or 'uri' must be provided")
+
+        if concurrency < 1:
+            raise ValueError("concurrency must be at least 1")
+
         self.model = model
         self.uri = uri
         self.port = port
         self.gpus = gpus
         self.with_vllm_server = with_vllm_server
         self.concurrency = concurrency
-        self.vllm_kwargs = vllm_kwargs
-        self.forget_predefined_vllm_kwargs = forget_predefined_vllm_kwargs
+        self.vllm_args = vllm_args
+        self.forget_predefined_vllm_args = forget_predefined_vllm_args
+        self.return_documents = return_documents
         self.server = None
         self.client = None
 
-        if self.uri is not None and self.model is None:
+        if self.uri is not None:
             self.model = get_model_from_uri(self.uri)
 
     def start_server_and_client(self):
-        from vlmparse.registries import (
-            converter_config_registry,
-            docker_config_registry,
-        )
-
-        gpu_device_ids = None
-        if self.gpus is not None:
-            gpu_device_ids = [g.strip() for g in self.gpus.split(",")]
+        from vlmparse.registries import converter_config_registry
 
         if self.uri is None:
-            docker_config = docker_config_registry.get(
-                self.model, default=self.with_vllm_server
+            _, _, self.server, docker_config = start_server(
+                model=self.model,
+                gpus=self.gpus,
+                port=self.port,
+                with_vllm_server=self.with_vllm_server,
+                vllm_args=self.vllm_args,
+                forget_predefined_vllm_args=self.forget_predefined_vllm_args,
+                auto_stop=True,
             )
 
             if docker_config is not None:
-                if self.port is not None:
-                    docker_config.docker_port = self.port
-                docker_config.gpu_device_ids = gpu_device_ids
-                docker_config.update_command_args(
-                    self.vllm_kwargs,
-                    forget_predefined_vllm_kwargs=self.forget_predefined_vllm_kwargs,
+                self.client = docker_config.get_client(
+                    return_documents_in_batch_mode=self.return_documents
                 )
-                self.server = docker_config.get_server(auto_stop=True)
-
-                self.server.start()
-
-                self.client = docker_config.get_client()
             else:
-                self.client = converter_config_registry.get(self.model).get_client()
+                self.client = converter_config_registry.get(self.model).get_client(
+                    return_documents_in_batch_mode=self.return_documents
+                )
 
         else:
             client_config = converter_config_registry.get(self.model, uri=self.uri)
 
-            self.client = client_config.get_client()
+            self.client = client_config.get_client(
+                return_documents_in_batch_mode=self.return_documents
+            )
 
     def stop_server(self):
         if self.server is not None and self.server.auto_stop:
@@ -80,16 +133,30 @@ class ConverterWithServer:
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
-        self.stop_server()
+        try:
+            self.stop_server()
+        except Exception as e:
+            logger.warning(f"Error stopping server during cleanup: {e}")
+        return False  # Don't suppress exceptions
 
     def parse(
         self,
         inputs: str | list[str],
         out_folder: str = ".",
         mode: Literal["document", "md", "md_page"] = "document",
+        conversion_mode: Literal[
+            "ocr",
+            "ocr_layout",
+            "table",
+            "image_description",
+            "formula",
+            "chart",
+        ]
+        | None = None,
         dpi: int | None = None,
         debug: bool = False,
         retrylast: bool = False,
+        completion_kwargs: dict | None = None,
     ):
         assert (
             self.client is not None
@@ -126,6 +193,14 @@ class ConverterWithServer:
         if dpi is not None:
             self.client.config.dpi = int(dpi)
 
+        if conversion_mode is not None:
+            self.client.config.conversion_mode = conversion_mode
+
+        if completion_kwargs is not None and hasattr(
+            self.client.config, "completion_kwargs"
+        ):
+            self.client.config.completion_kwargs |= completion_kwargs
+
         if debug:
             self.client.debug = debug
 

vlmparse/registries.py

@@ -1,37 +1,24 @@
 import os
 from collections.abc import Callable
 
-from vlmparse.clients.chandra import ChandraConverterConfig, ChandraDockerServerConfig
-from vlmparse.clients.deepseekocr import (
-    DeepSeekOCRConverterConfig,
-    DeepSeekOCRDockerServerConfig,
-)
-from vlmparse.clients.docling import DoclingConverterConfig, DoclingDockerServerConfig
-from vlmparse.clients.dotsocr import DotsOCRConverterConfig, DotsOCRDockerServerConfig
-from vlmparse.clients.granite_docling import (
-    GraniteDoclingConverterConfig,
-    GraniteDoclingDockerServerConfig,
-)
-from vlmparse.clients.hunyuanocr import (
-    HunyuanOCRConverterConfig,
-    HunyuanOCRDockerServerConfig,
-)
+from vlmparse.clients.chandra import ChandraDockerServerConfig
+from vlmparse.clients.deepseekocr import DeepSeekOCRDockerServerConfig
+from vlmparse.clients.docling import DoclingDockerServerConfig
+from vlmparse.clients.dotsocr import DotsOCRDockerServerConfig
+from vlmparse.clients.granite_docling import GraniteDoclingDockerServerConfig
+from vlmparse.clients.hunyuanocr import HunyuanOCRDockerServerConfig
 from vlmparse.clients.lightonocr import (
-    LightOnOCRConverterConfig,
+    LightonOCR21BServerConfig,
     LightOnOCRDockerServerConfig,
 )
-from vlmparse.clients.mineru import MinerUConverterConfig, MinerUDockerServerConfig
-from vlmparse.clients.nanonetocr import (
-    NanonetOCR2ConverterConfig,
-    NanonetOCR2DockerServerConfig,
-)
-from vlmparse.clients.olmocr import OlmOCRConverterConfig, OlmOCRDockerServerConfig
-from vlmparse.clients.openai_converter import LLMParams, OpenAIConverterConfig
-from vlmparse.clients.paddleocrvl import (
-    PaddleOCRVLConverterConfig,
-    PaddleOCRVLDockerServerConfig,
-)
-from vlmparse.servers.docker_server import DEFAULT_MODEL_NAME, docker_config_registry
+from vlmparse.clients.mineru import MinerUDockerServerConfig
+from vlmparse.clients.mistral_converter import MistralOCRConverterConfig
+from vlmparse.clients.nanonetocr import NanonetOCR2DockerServerConfig
+from vlmparse.clients.olmocr import OlmOCRDockerServerConfig
+from vlmparse.clients.openai_converter import OpenAIConverterConfig
+from vlmparse.clients.paddleocrvl import PaddleOCRVLDockerServerConfig
+from vlmparse.converter import ConverterConfig
+from vlmparse.servers.docker_server import DockerServerConfig, docker_config_registry
 
 
 def get_default(cls, field_name):
@@ -43,7 +30,8 @@ def get_default(cls, field_name):
     return field_info.default
 
 
-for server_config_cls in [
+# All server configs - single source of truth
+SERVER_CONFIGS: list[type[DockerServerConfig]] = [
     ChandraDockerServerConfig,
     LightOnOCRDockerServerConfig,
     DotsOCRDockerServerConfig,
@@ -55,7 +43,11 @@ for server_config_cls in [
     MinerUDockerServerConfig,
     DeepSeekOCRDockerServerConfig,
     GraniteDoclingDockerServerConfig,
-]:
+    LightonOCR21BServerConfig,
+]
+
+# Register docker server configs
+for server_config_cls in SERVER_CONFIGS:
     aliases = get_default(server_config_cls, "aliases") or []
     model_name = get_default(server_config_cls, "model_name")
     names = [n for n in aliases + [model_name] if isinstance(n, str)]
@@ -64,37 +56,81 @@ for server_config_cls in [
 
 
 class ConverterConfigRegistry:
-    """Registry for mapping model names to their Docker configurations."""
+    """Registry for mapping model names to their converter configurations.
+
+    Thread-safe registry that maps model names to their converter configuration factories.
+    """
 
     def __init__(self):
-        self._registry = dict()
+        import threading
+
+        self._registry: dict[str, Callable[[str | None], ConverterConfig]] = {}
+        self._lock = threading.RLock()
 
     def register(
         self,
         model_name: str,
-        config_factory: Callable[[str], OpenAIConverterConfig | None],
+        config_factory: Callable[[str | None], ConverterConfig],
     ):
-        """Register a config factory for a model name."""
-        self._registry[model_name] = config_factory
+        """Register a config factory for a model name (thread-safe)."""
+        with self._lock:
+            self._registry[model_name] = config_factory
 
-    def get(self, model_name: str, uri: str | None = None) -> OpenAIConverterConfig:
-        """Get config for a model name. Returns default if not registered."""
-        if model_name in self._registry:
-            return self._registry[model_name](uri=uri)
+    def register_from_server(
+        self,
+        server_config_cls: type[DockerServerConfig],
+    ):
+        """Register converter config derived from a server config class.
+
+        This ensures model_name and default_model_name are consistently
+        passed from server to client config via _create_client_kwargs.
+        """
+        aliases = get_default(server_config_cls, "aliases") or []
+        model_name = get_default(server_config_cls, "model_name")
+        names = [n for n in aliases + [model_name] if isinstance(n, str)]
+        # Also register short name (after last /)
+        if model_name and "/" in model_name:
+            names.append(model_name.split("/")[-1])
+
+        def factory(uri: str | None, cls=server_config_cls) -> ConverterConfig:
+            server = cls()
+            client_config = server.client_config
+            # Override base_url if provided
+            if uri is not None:
+                client_config = client_config.model_copy(update={"base_url": uri})
+            return client_config
+
+        with self._lock:
+            for name in names:
+                self._registry[name] = factory
+
+    def get(self, model_name: str, uri: str | None = None) -> ConverterConfig:
+        """Get config for a model name (thread-safe). Returns default if not registered."""
+        with self._lock:
+            factory = self._registry.get(model_name)
+
+        if factory is not None:
+            return factory(uri)
         # Fallback to OpenAIConverterConfig for unregistered models
         if uri is not None:
-            return OpenAIConverterConfig(
-                llm_params=LLMParams(base_url=uri, model_name=model_name)
-            )
-        return OpenAIConverterConfig(llm_params=LLMParams(model_name=model_name))
+            return OpenAIConverterConfig(base_url=uri)
+        return OpenAIConverterConfig(model_name=model_name)
 
     def list_models(self) -> list[str]:
-        """List all registered model names."""
-        return list(self._registry.keys())
+        """List all registered model names (thread-safe)."""
+        with self._lock:
+            return list(self._registry.keys())
 
 
 # Global registry instance
 converter_config_registry = ConverterConfigRegistry()
+
+# Register all server-backed converters through the server config
+# This ensures model_name and default_model_name are consistently passed
+for server_config_cls in SERVER_CONFIGS:
+    converter_config_registry.register_from_server(server_config_cls)
+
+# External API configs (no server config - these are cloud APIs)
 GOOGLE_API_BASE_URL = (
     os.getenv("GOOGLE_API_BASE_URL")
     or "https://generativelanguage.googleapis.com/v1beta/openai/"
@@ -111,11 +147,10 @@ for gemini_model in [
     converter_config_registry.register(
         gemini_model,
         lambda uri=None, model=gemini_model: OpenAIConverterConfig(
-            llm_params=LLMParams(
-                model_name=model,
-                base_url=GOOGLE_API_BASE_URL if uri is None else uri,
-                api_key=os.getenv("GOOGLE_API_KEY"),
-            )
+            model_name=model,
+            base_url=GOOGLE_API_BASE_URL if uri is None else uri,
+            api_key=os.getenv("GOOGLE_API_KEY"),
+            default_model_name=model,
         ),
     )
 for openai_model in [
@@ -126,45 +161,18 @@ for openai_model in [
     converter_config_registry.register(
         openai_model,
         lambda uri=None, model=openai_model: OpenAIConverterConfig(
-            llm_params=LLMParams(
-                model_name=model,
-                base_url=None,
-                api_key=os.getenv("OPENAI_API_KEY"),
-            )
+            model_name=model,
+            base_url=None,
+            api_key=os.getenv("OPENAI_API_KEY"),
+            default_model_name=model,
         ),
     )
 
-for converter_config_cls in [
-    ChandraConverterConfig,
-    LightOnOCRConverterConfig,
-    DotsOCRConverterConfig,
-    PaddleOCRVLConverterConfig,
-    NanonetOCR2ConverterConfig,
-    HunyuanOCRConverterConfig,
-    DeepSeekOCRConverterConfig,
-    GraniteDoclingConverterConfig,
-    OlmOCRConverterConfig,
-]:
-    aliases = get_default(converter_config_cls, "aliases") or []
-    model_name = get_default(converter_config_cls, "model_name")
-    names = [n for n in aliases + [model_name] if isinstance(n, str)]
-    for name in names:
-        converter_config_registry.register(
-            name,
-            lambda uri, cls=converter_config_cls: cls(
-                llm_params=LLMParams(
-                    base_url=uri,
-                    model_name=DEFAULT_MODEL_NAME,
-                    api_key="",
-                )
-            ),
-        )
-for converter_config_cls in [MinerUConverterConfig, DoclingConverterConfig]:
-    aliases = get_default(converter_config_cls, "aliases") or []
-    model_name = get_default(converter_config_cls, "model_name")
-    names = [n for n in aliases + [model_name] if isinstance(n, str)]
-    for name in names:
-        converter_config_registry.register(
-            name,
-            lambda uri, cls=converter_config_cls: cls(base_url=uri),
-        )
+for mistral_model in ["mistral-ocr-latest", "mistral-ocr"]:
+    converter_config_registry.register(
+        mistral_model,
+        lambda uri=None, model=mistral_model: MistralOCRConverterConfig(
+            base_url="https://api.mistral.ai/v1" if uri is None else uri,
+            api_key=os.getenv("MISTRAL_API_KEY"),
+        ),
+    )