parse-bench 1.0.0__py3-none-any.whl

parse_bench/inference/providers/__init__.py

@@ -0,0 +1,28 @@
+"""Provider implementations for document parsing."""
+
+# Import providers to register them
+from parse_bench.inference.providers import (
+    extract,  # noqa: F401
+    layoutdet,  # noqa: F401
+    parse,  # noqa: F401
+)
+from parse_bench.inference.providers.base import (
+    Provider,
+    ProviderConfigError,
+    ProviderError,
+    ProviderPermanentError,
+    ProviderRateLimitError,
+    ProviderTransientError,
+)
+from parse_bench.inference.providers.registry import create_provider, register_provider
+
+__all__ = [
+    "Provider",
+    "ProviderConfigError",
+    "ProviderError",
+    "ProviderPermanentError",
+    "ProviderRateLimitError",
+    "ProviderTransientError",
+    "create_provider",
+    "register_provider",
+]

parse_bench/inference/providers/base.py

@@ -0,0 +1,196 @@
+import asyncio
+import concurrent.futures
+from abc import ABC, abstractmethod
+from collections.abc import Coroutine
+from typing import Any
+
+from parse_bench.schemas.pipeline import PipelineSpec
+from parse_bench.schemas.pipeline_io import (
+    InferenceRequest,
+    InferenceResult,
+    RawInferenceResult,
+)
+
+
+class ProviderError(Exception):
+    """Base exception for provider-related failures."""
+
+    def __init__(self, message: str, *, debug_payload: dict[str, Any] | None = None):
+        super().__init__(message)
+        self.debug_payload = debug_payload
+
+
+class ProviderConfigError(ProviderError):
+    """Raised when a provider is misconfigured (missing API keys, bad endpoint, etc.)."""
+
+
+class ProviderRateLimitError(ProviderError):
+    """Raised when a provider hits rate limits or quota issues."""
+
+
+class ProviderTransientError(ProviderError):
+    """
+    Raised for transient errors that may succeed on retry
+    (e.g. network issues, 5xx responses).
+    """
+
+
+class ProviderPermanentError(ProviderError):
+    """
+    Raised for permanent errors that are not expected to succeed on retry
+    (e.g. unsupported file type, invalid request, 4xx errors).
+    """
+
+
+class Provider(ABC):
+    """Abstract base class for document parsing providers."""
+
+    def __init__(
+        self,
+        provider_name: str,
+        base_config: dict[str, Any] | None = None,
+    ):
+        """
+        Initialize a provider.
+
+        :param provider_name: Name of the provider
+        :param base_config: Optional shared configuration dictionary.
+            Can include `use_staging` (bool) to use staging environment.
+        """
+        self._provider_name = provider_name
+        self._base_config = base_config or {}
+
+    @property
+    def provider_name(self) -> str:
+        """Return the provider name."""
+        return self._provider_name
+
+    @property
+    def base_config(self) -> dict[str, Any]:
+        """Return the base configuration."""
+        return self._base_config
+
+    @property
+    def credit_rate_usd(self) -> float | None:
+        """USD cost per credit. Override in subclasses that charge credits."""
+        return None
+
+    @staticmethod
+    def run_async_from_sync(coro: Coroutine[Any, Any, Any]) -> Any:
+        """
+        Run an async coroutine from a synchronous context.
+
+        This helper handles both cases:
+        - If there's no running event loop, uses asyncio.run()
+        - If there's a running event loop, runs the coroutine in a new thread with a new event loop
+
+        :param coro: The coroutine to run
+        :return: The result of the coroutine
+        """
+        try:
+            # Try to get the current event loop
+            # If we get here, there's a running loop, so we need to run in a thread
+            asyncio.get_running_loop()
+            with concurrent.futures.ThreadPoolExecutor() as executor:
+                future = executor.submit(asyncio.run, coro)
+                return future.result()
+        except RuntimeError:
+            # No running event loop, we can use asyncio.run() directly
+            return asyncio.run(coro)
+
+    @abstractmethod
+    def run_inference(self, pipeline: PipelineSpec, request: InferenceRequest) -> RawInferenceResult:
+        """
+        Run inference for a single request and return raw results.
+
+        This method should only fetch raw data from the provider API.
+        Normalization is handled separately by the normalize() method.
+
+        :param pipeline: Pipeline specification
+        :param request: Inference request
+        :return: Raw inference result (before normalization)
+        :raises ProviderError: For any provider-related failures
+        """
+        raise NotImplementedError("Subclasses must implement this method")
+
+    @abstractmethod
+    def normalize(self, raw_result: RawInferenceResult) -> InferenceResult:
+        """
+        Normalize raw inference result to produce structured output.
+
+        This method converts the raw API response into a structured
+        format (ParseOutput or ExtractOutput) while preserving the
+        raw output for potential re-normalization.
+
+        Note: Each provider implementation is product-type specific
+        and will return either ParseOutput or ExtractOutput, not both.
+
+        :param raw_result: Raw inference result from run_inference()
+        :return: Inference result with both raw and normalized outputs
+        :raises ProviderError: For any normalization failures
+        """
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def recompute_cost(self, raw_output: dict[str, Any]) -> None:
+        """Re-derive the cost fields in ``raw_output`` from the usage it already
+        records, using this provider's current pricing.
+
+        This is the SINGLE generic seam the re-normalization path calls for every
+        provider (see ``inference/renormalize.py``), so correcting a pricing table
+        and re-running ``bench inference renormalize`` re-prices saved runs with no
+        fresh (expensive) inference. It is deliberately pipeline- and
+        provider-agnostic at the call site: the runner never names a provider.
+
+        What is irreducibly provider-specific is the rate card -- only this
+        provider knows what it charges for a cached token, a cache write, a tool
+        container, or a long-context tier -- so the base implementation is a no-op
+        and a provider that can re-price overrides this one method. A provider
+        whose cost came from an upstream credits field (not tokens we hold) simply
+        does not override it and keeps what it recorded.
+
+        Contract for overrides: mutate ``raw_output`` in place, derive everything
+        from what is already there (so it works on a saved ``.raw.json``), be
+        idempotent, and never call the API. Leave cost untouched when the recorded
+        usage is missing -- re-pricing a usage-less artifact would overwrite a real
+        number with a zero.
+        """
+        return None
+
+    def cancel(self, example_id: str) -> bool:
+        """
+        Cancel any in-flight inference work for the given ``example_id``.
+
+        Default implementation is a no-op for providers that do not spawn
+        external resources (subprocesses, remote jobs) that need explicit
+        teardown. Providers that fork subprocesses or hold network handles
+        should override this to actually terminate the work — otherwise
+        per-file timeouts in the runner will only release the calling
+        thread while the underlying work continues to run, racing the
+        retry attempt and producing duplicate / zombie processes.
+
+        :param example_id: The ``InferenceRequest.example_id`` for the
+            in-flight request that should be cancelled.
+        :return: True if a matching in-flight request was found and a
+            cancellation signal was issued; False otherwise.
+        """
+        return False
+
+    def consume_active_job_id(self, example_id: str) -> str | None:
+        """Consume a provider-owned active remote job id, if one exists."""
+
+        return None
+
+    def run_inference_normalized(self, pipeline: PipelineSpec, request: InferenceRequest) -> InferenceResult:
+        """
+        Run inference and normalize in one step (convenience method).
+
+        This is a convenience method that combines run_inference() and
+        normalize() for backward compatibility and simple use cases.
+
+        :param pipeline: Pipeline specification
+        :param request: Inference request
+        :return: Inference result with both raw and normalized outputs
+        :raises ProviderError: For any provider-related failures
+        """
+        raw_result = self.run_inference(pipeline, request)
+        return self.normalize(raw_result)

parse_bench/inference/providers/cancellation.py

@@ -0,0 +1,137 @@
+"""Reusable per-``example_id`` cancellation registry for HTTP/SDK providers.
+
+When the runner's per-file timeout fires, ``ThreadPoolExecutor.shutdown(wait=False)``
+only releases the calling thread - any in-flight HTTP request the worker
+thread spawned (httpx polling, requests session, LlamaCloud SDK call) keeps
+running, and the next retry attempt sends a duplicate request to staging.
+Closing the underlying client breaks the provider's polling loop on its
+next iteration so the worker thread unwinds with a transient error and the
+retry loop can submit a fresh request without piling on parallel duplicates.
+
+Important caveat - what closing a client does and does not abort:
+
+* It DOES break a polling loop that calls ``client.get(...)`` repeatedly
+  on a long-running job. The next call after ``close()`` raises immediately
+  (httpx: ``RuntimeError: Cannot send a request, as the client has been closed.``;
+  requests: ``ConnectionError`` on the next ``session.get``).
+
+* It does NOT interrupt a thread already blocked inside a single socket
+  read on another thread - Python threads are not OS-cancellable, and
+  closing the client object only marks it closed; the kernel ``recv`` call
+  finishes only when the server responds or the read timeout expires.
+
+For the bench bug - duplicate requests to staging during per-file timeout
+retries - the polling-loop case is the one that matters. Long-running
+parse / extract jobs are polled in tight loops; closing the client makes
+the next poll raise within milliseconds. Per-request read timeouts on the
+underlying client cap the worst-case stalled-read tail.
+
+This module provides a tiny helper that:
+  * registers a closeable handle (httpx.Client, requests.Session,
+    llama_cloud.LlamaCloud, ...) keyed by ``example_id``;
+  * exposes a ``cancel(example_id)`` that pops the handle and calls
+    ``.close()`` (best-effort - providers swallow secondary errors so the
+    cancel path can never break the runner).
+
+Each provider holds one ``CancellableClientRegistry`` instance, registers
+its client at the start of a request, and unregisters in a ``finally``.
+The registry is thread-safe so concurrent ``run_inference`` calls (one per
+``example_id``) do not collide.
+"""
+
+from __future__ import annotations
+
+import logging
+import threading
+from typing import Protocol
+
+logger = logging.getLogger(__name__)
+
+
+class _Closeable(Protocol):
+    """Anything with a no-arg ``close()``: ``httpx.Client``, ``requests.Session``,
+    ``llama_cloud.LlamaCloud`` (closes its underlying ``httpx.Client``), ..."""
+
+    def close(self) -> None: ...
+
+
+class CancellableClientRegistry:
+    """Thread-safe per-``example_id`` mapping of in-flight HTTP/SDK clients.
+
+    Providers should:
+
+        # In __init__:
+        self._inflight = CancellableClientRegistry(provider_name="...")
+
+        # At the start of run_inference (after the client is built):
+        self._inflight.register(request.example_id, client)
+        try:
+            ...
+        finally:
+            self._inflight.unregister(request.example_id, client)
+
+        # In cancel(example_id):
+        return self._inflight.cancel(example_id)
+
+    The registry never raises from ``cancel`` - a broken cancel must not
+    break the runner's retry loop.
+    """
+
+    def __init__(self, *, provider_name: str) -> None:
+        self._provider_name = provider_name
+        self._lock = threading.Lock()
+        self._inflight: dict[str, _Closeable] = {}
+
+    def register(self, example_id: str, client: _Closeable) -> None:
+        """Track ``client`` so a later ``cancel(example_id)`` can close it.
+
+        If the slot is already occupied (e.g. because the previous attempt's
+        cleanup raced the next attempt's submit), the new client wins - the
+        old one was either already cancelled or about to be unregistered.
+        """
+        with self._lock:
+            self._inflight[example_id] = client
+
+    def unregister(self, example_id: str, client: _Closeable) -> None:
+        """Remove ``client`` from the registry if it is the live entry.
+
+        Compares by identity so we never clobber a registration from a
+        concurrent retry attempt. Idempotent - safe to call from a
+        ``finally`` even when ``cancel`` already popped the entry.
+        """
+        with self._lock:
+            current = self._inflight.get(example_id)
+            if current is client:
+                self._inflight.pop(example_id, None)
+
+    def cancel(self, example_id: str) -> bool:
+        """Pop and close any registered client for ``example_id``.
+
+        :return: True if a matching client was found and ``close()`` was
+            attempted (regardless of whether close itself succeeded), False
+            if no client was registered.
+        """
+        with self._lock:
+            client = self._inflight.pop(example_id, None)
+        if client is None:
+            return False
+
+        logger.info(
+            "%s.cancel: closing in-flight client for example_id=%s",
+            self._provider_name,
+            example_id,
+        )
+        try:
+            client.close()
+        except Exception as exc:  # noqa: BLE001 - cancel must never raise
+            # Closing a client mid-request can surface various provider-
+            # specific exceptions (httpx connection state, broken pipe,
+            # SDK wrappers raising their own types). None of them should
+            # break the runner's retry loop.
+            logger.debug(
+                "%s.cancel: client.close() raised for example_id=%s: %s",
+                self._provider_name,
+                example_id,
+                exc,
+            )
+        return True

parse_bench/inference/providers/extract/__init__.py

@@ -0,0 +1,22 @@
+"""Extract providers imported for registry side effects.
+
+Only the minimal ParseBench EXTRACT integration providers are registered here.
+Imports are best-effort so base ParseBench imports do not require optional
+provider SDKs such as ``extend-ai``.
+"""
+
+import importlib
+import logging
+
+logger = logging.getLogger(__name__)
+
+_PROVIDER_MODULES = [
+    "extend",
+    "llamaextract_v2_api",
+]
+
+for _mod in _PROVIDER_MODULES:
+    try:
+        importlib.import_module(f"parse_bench.inference.providers.extract.{_mod}")
+    except ImportError:
+        logger.debug("Skipping extract provider %s (missing dependency)", _mod)