fastdocparse 0.1.0__py3-none-any.whl

docextract/json_repair.py

@@ -0,0 +1,79 @@
+"""Recover a JSON object from raw LLM output that isn't guaranteed to be clean JSON."""
+
+import json
+import re
+from typing import Any, Dict
+
+
+def parse_json_from_llm(text: str) -> Dict[str, Any]:
+    """Safely parse JSON from LLM output, handling markdown blocks and `<think>` tags."""
+    text = text.strip()
+    text = re.sub(r"<think\b[^>]*>.*?</think>", "", text, flags=re.DOTALL | re.IGNORECASE)
+
+    match = re.search(r"```json\s*(.*?)\s*```", text, re.DOTALL | re.IGNORECASE)
+    if match:
+        try:
+            return json.loads(match.group(1).strip())
+        except Exception:
+            pass
+
+    brace_positions = [i for i, c in enumerate(text) if c == "{"]
+    for pos in reversed(brace_positions):
+        end = text.rfind("}", pos)
+        if end > pos:
+            try:
+                obj = json.loads(text[pos:end + 1])
+                if isinstance(obj, dict):
+                    return obj
+            except Exception:
+                continue
+
+    # Nothing above found valid JSON — likely the response was cut off mid-object
+    # (hit max_tokens). Try to close whatever was left open, so the fields that were
+    # fully generated before the cutoff still come back instead of the whole response
+    # being discarded.
+    return _repair_truncated_json(text)
+
+
+def _repair_truncated_json(text: str) -> Dict[str, Any]:
+    start = text.find("{")
+    if start == -1:
+        return {}
+    snippet = text[start:]
+
+    stack = []
+    in_string = False
+    escape = False
+
+    for ch in snippet:
+        if in_string:
+            if escape:
+                escape = False
+            elif ch == "\\":
+                escape = True
+            elif ch == '"':
+                in_string = False
+        else:
+            if ch == '"':
+                in_string = True
+            elif ch in "{[":
+                stack.append(ch)
+            elif ch in "}]" and stack:
+                stack.pop()
+
+    repaired = snippet
+    if in_string:
+        repaired += '"'
+    # Drop a dangling comma or an incomplete "key": fragment with no value at all —
+    # neither can be closed into valid JSON by just appending brackets.
+    repaired = re.sub(r',\s*"[^"]*"?\s*:?\s*$', "", repaired)
+    repaired = re.sub(r',\s*$', "", repaired)
+
+    for opener in reversed(stack):
+        repaired += "}" if opener == "{" else "]"
+
+    try:
+        obj = json.loads(repaired)
+        return obj if isinstance(obj, dict) else {}
+    except Exception:
+        return {}

docextract/llm_client.py

@@ -0,0 +1,82 @@
+"""OpenAI-compatible LLM client for document extraction."""
+
+import time
+from typing import List, Optional
+
+from openai import (
+    APIConnectionError,
+    APITimeoutError,
+    AuthenticationError,
+    OpenAI,
+    RateLimitError,
+)
+
+
+class LLMClientError(Exception):
+    """Raised when a call to the LLM endpoint fails (bad auth, unreachable, or exhausted retries)."""
+
+
+class LLMClient:
+    """A thin adapter over any OpenAI-compatible endpoint."""
+
+    def __init__(self, base_url: Optional[str] = None, api_key: Optional[str] = None, model: Optional[str] = None):
+        """Initialize the LLM client.
+
+        Args:
+            base_url: The base URL of the OpenAI-compatible API. Defaults to OpenAI if None.
+            api_key: The API key. Can be dummy for local endpoints like Ollama.
+            model: The model string to use for completions.
+        """
+        # If not provided, it will fallback to standard environment variables if used.
+        self.base_url = base_url
+        self.api_key = api_key or "dummy-key-for-local"
+        self.model = model or "gpt-4o-mini"
+
+        self.client = OpenAI(
+            base_url=self.base_url,
+            api_key=self.api_key,
+            timeout=60.0
+        )
+
+    def _call(self, messages: List[dict], temperature: float, max_tokens: int, retries: int = 2, backoff: float = 1.0) -> str:
+        """Run a chat completion, retrying transient failures and raising LLMClientError on exhaustion."""
+        last_error: Optional[Exception] = None
+        for attempt in range(retries + 1):
+            try:
+                response = self.client.chat.completions.create(
+                    model=self.model,
+                    messages=messages,
+                    temperature=temperature,
+                    max_tokens=max_tokens,
+                )
+                return response.choices[0].message.content or ""
+            except AuthenticationError as e:
+                raise LLMClientError(f"Authentication failed for model '{self.model}'. Check your API key.") from e
+            except (APIConnectionError, APITimeoutError, RateLimitError) as e:
+                last_error = e
+                if attempt < retries:
+                    time.sleep(backoff * (attempt + 1))
+                    continue
+            except Exception as e:
+                raise LLMClientError(f"LLM call to '{self.model}' failed: {e}") from e
+
+        raise LLMClientError(
+            f"Could not reach the LLM endpoint at {self.base_url or 'https://api.openai.com'} "
+            f"after {retries + 1} attempts: {last_error}"
+        )
+
+    def extract(self, prompt: str, document_text: str, temperature: float = 0.0, max_tokens: int = 4096) -> str:
+        """Run the prompt against the LLM to get the structured extraction."""
+        final_prompt = prompt.replace("{document_text}", document_text)
+        return self._call(
+            [
+                {"role": "system", "content": "You are a ultra-fast document extractor. Output strict JSON."},
+                {"role": "user", "content": final_prompt},
+            ],
+            temperature=temperature,
+            max_tokens=max_tokens,
+        )
+
+    def complete(self, prompt: str, temperature: float = 0.0, max_tokens: int = 2048) -> str:
+        """Run a single free-form prompt against the LLM (no document-extraction framing)."""
+        return self._call([{"role": "user", "content": prompt}], temperature=temperature, max_tokens=max_tokens)

docextract/ocr_engine.py

@@ -0,0 +1,105 @@
+"""Phase 1 OCR Engine: Fast, layout-aware OCR for images and scanned PDFs."""
+
+from __future__ import annotations
+
+import importlib.util
+import io
+import logging
+import threading
+from typing import List, Tuple
+from PIL import Image
+
+logger = logging.getLogger(__name__)
+
+# Whether the package is installed at all — checked via find_spec, which locates the
+# module without executing it, so this costs nothing at import time. The actual
+# `import rapidocr_onnxruntime` (which pulls in onnxruntime, opencv, numpy — measured
+# at ~1.2s) and RapidOCR() construction only happen lazily, on first real OCR call, via
+# _get_rapid_ocr() below. Without this, every `import docextract` would eagerly pay
+# that cost even for code that never touches OCR (e.g. just building a Schema).
+HAS_RAPID_OCR = importlib.util.find_spec("rapidocr_onnxruntime") is not None
+
+_rapid_ocr = None
+_init_lock = threading.Lock()
+_init_failed = False
+
+
+def _get_rapid_ocr():
+    global _rapid_ocr, _init_failed
+    if _rapid_ocr is not None or _init_failed:
+        return _rapid_ocr
+    with _init_lock:
+        if _rapid_ocr is not None or _init_failed:
+            return _rapid_ocr
+        try:
+            from rapidocr_onnxruntime import RapidOCR
+            _rapid_ocr = RapidOCR()
+        except Exception:
+            logger.warning("RapidOCR failed to initialize; OCR extraction will return empty text.", exc_info=True)
+            _init_failed = True
+    return _rapid_ocr
+
+
+_ocr_lock = threading.Lock()
+
+
+def extract_text_from_image_ocr(image_bytes: bytes, structured_mode: bool = False, min_confidence: float = 0.3) -> str:
+    """Extract layout-preserved text from an image (PNG/JPG) using local OCR engine.
+
+    Returns a clean, line-grouped text representation preserving horizontal spacing.
+    """
+    if not HAS_RAPID_OCR:
+        return ""
+
+    rapid_ocr = _get_rapid_ocr()
+    if rapid_ocr is None:
+        return ""
+
+    try:
+        img = Image.open(io.BytesIO(image_bytes))
+        if img.mode != "RGB":
+            img = img.convert("RGB")
+
+        buf = io.BytesIO()
+        img.save(buf, format="JPEG")
+
+        with _ocr_lock:
+            result, _ = rapid_ocr(buf.getvalue())
+        if not result:
+            return ""
+
+        lines: List[Tuple[float, float, str]] = []
+        for box, text, conf in result:
+            conf_val = float(conf) if conf is not None else 0.0
+            if conf_val > min_confidence and text and str(text).strip():
+                y0 = min(pt[1] for pt in box)
+                x0 = min(pt[0] for pt in box)
+                lines.append((y0, x0, str(text).strip()))
+
+        lines.sort(key=lambda item: (round(item[0] / 15), item[1]))
+
+        grouped_lines: List[str] = []
+        current_y_group = -1
+        current_line_parts: List[str] = []
+
+        for y0, x0, text in lines:
+            group = round(y0 / 15)
+            if structured_mode:
+                text_formatted = f"[X:{int(x0)}] {text}"
+            else:
+                text_formatted = text
+                
+            if group != current_y_group:
+                if current_line_parts:
+                    grouped_lines.append("    ".join(current_line_parts))
+                current_y_group = group
+                current_line_parts = [text_formatted]
+            else:
+                current_line_parts.append(text_formatted)
+
+        if current_line_parts:
+            grouped_lines.append("    ".join(current_line_parts))
+
+        return "\n".join(grouped_lines)
+    except Exception:
+        return ""

docextract/parser.py

@@ -0,0 +1,313 @@
+"""Parser module orchestrating document extraction."""
+
+import asyncio
+import functools
+import logging
+import re
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, Callable, Dict, List, Optional
+
+import pymupdf
+
+from .cache import Cache, make_cache_key
+from .config import ExtractionConfig
+from .grounding import Issue, _is_present, check_substring, cross_check, validate_field_constraints
+from .json_repair import parse_json_from_llm as _parse_json_from_llm
+from .llm_client import LLMClient
+from .ocr_engine import extract_text_from_image_ocr
+from .pdf_utils import chunk_document_text, extract_text_from_pdf, pdf_to_page_images
+from .prompt_compiler import compile_prompt
+from .schema import Schema
+
+logger = logging.getLogger(__name__)
+
+
+class EmptyDocumentError(ValueError):
+    """Raised when no text could be extracted from a document by any route."""
+
+
+_LAYOUT_TAG_RE = re.compile(r"\[X:\d+\]\s*")
+
+
+def _strip_layout_tags(text: str) -> str:
+    """Remove the "[X:nnn]" column-position markers structured_mode adds for the LLM's
+    benefit before running any grounding check against the text. Those digits are pixel
+    coordinates, not document content — left in, they can splice into the middle of a
+    real value's character stream (e.g. an address split across OCR lines) and break a
+    fuzzy match that would otherwise succeed, flagging a correct value as ungrounded.
+    """
+    return _LAYOUT_TAG_RE.sub("", text)
+
+
+class UnknownIngestionKindError(ValueError):
+    """Raised when extract()/aextract() is given a kind with no registered handler.
+
+    Kept distinct from a bare ValueError so callers (the CLI included) can tell "you
+    passed an unregistered --kind" apart from any other ValueError that might surface
+    from deeper in the pipeline (e.g. from inside a custom ingestion handler).
+    """
+
+
+def _merge_extracted_data(results: List[Dict[str, Any]], chunks: List[str], schema: Schema) -> Dict[str, Any]:
+    """Merge per-chunk extractions into one result.
+
+    List fields concatenate across chunks (unchanged). For scalar fields, when more than
+    one chunk produced a non-null value, prefer whichever one is grounded in the chunk
+    it actually came from — we have that chunk's source text right here, so use it
+    instead of blindly taking the first chunk's answer regardless of whether it's real.
+    Falls back to first-non-null when nothing grounds (matches the old behavior exactly,
+    so a single-chunk document — the common case — is unaffected).
+    """
+    if not results:
+        return {}
+
+    merged: Dict[str, Any] = {}
+    for f in schema.fields:
+        merged[f.name] = [] if f.type == "list" else None
+
+    for f in schema.fields:
+        if f.type == "list":
+            for res in results:
+                val = res.get(f.name)
+                if val and isinstance(val, list):
+                    merged[f.name].extend(val)
+            continue
+
+        first_non_null = None
+        grounded_value = None
+        for res, chunk_text in zip(results, chunks):
+            val = res.get(f.name)
+            if not _is_present(val):
+                continue
+            if first_non_null is None:
+                first_non_null = val
+            if grounded_value is None and check_substring(
+                val, _strip_layout_tags(chunk_text), numeric=f.is_numeric, date=f.is_date
+            ):
+                grounded_value = val
+                break
+        merged[f.name] = grounded_value if grounded_value is not None else first_non_null
+
+    return merged
+
+
+# Ingestion handlers keyed by document kind. To support a new input format
+# (e.g. DOCX), add a handler here and register it — no changes needed to
+# DocumentParser itself.
+def _ingest_pdf(document_bytes: bytes, structured_mode: bool, config: ExtractionConfig) -> str:
+    doc_text = extract_text_from_pdf(document_bytes, max_pages=config.max_pages, structured_mode=structured_mode)
+    if len(doc_text.strip()) < 30:
+        # Scanned PDF: run local OCR on first page image
+        logger.info("Digital text layer too short (%d chars); falling back to OCR on page 1.", len(doc_text.strip()))
+        pages = pdf_to_page_images(
+            document_bytes, max_pages=1, dpi=config.pdf_render_dpi, max_dim=config.max_image_dim
+        )
+        if pages:
+            doc_text = extract_text_from_image_ocr(
+                pages[0].png_bytes, structured_mode=structured_mode, min_confidence=config.ocr_min_confidence
+            )
+    return doc_text
+
+
+def _ingest_image(document_bytes: bytes, structured_mode: bool, config: ExtractionConfig) -> str:
+    return extract_text_from_image_ocr(
+        document_bytes, structured_mode=structured_mode, min_confidence=config.ocr_min_confidence
+    )
+
+
+INGESTION_HANDLERS: Dict[str, Callable[[bytes, bool, ExtractionConfig], str]] = {
+    "pdf": _ingest_pdf,
+    "image": _ingest_image,
+}
+
+
+def register_default_ingestion_handler(kind: str, handler: Callable[[bytes, bool, ExtractionConfig], str]) -> None:
+    """Register a new default ingestion route, e.g. for DOCX/XLSX support, process-wide.
+
+    Named distinctly from DocumentParser.register_ingestion_handler() on purpose — that
+    one scopes a handler to a single instance; this one changes what *new* instances get
+    by default. They used to share a name, which made it easy to call the process-wide
+    one by habit when a scoped registration was actually intended.
+
+    handler receives (document_bytes, structured_mode, config) and returns the extracted
+    text. This affects DocumentParser instances created *after* this call — each instance
+    copies the default registry at construction time, so it can't be silently clobbered by
+    another part of the program (or another test) registering a different handler under
+    the same kind later.
+    """
+    INGESTION_HANDLERS[kind] = handler
+
+
+class DocumentParser:
+    """Orchestrates document extraction."""
+
+    def __init__(
+        self,
+        client: LLMClient,
+        config: Optional[ExtractionConfig] = None,
+        cache: Optional[Cache] = None,
+        ingestion_handlers: Optional[Dict[str, Callable[[bytes, bool, ExtractionConfig], str]]] = None,
+    ):
+        self.client = client
+        self.config = config or ExtractionConfig()
+        self.cache = cache
+        # Copied, not referenced, so registering a handler on one instance (or globally,
+        # after this instance already exists) never affects an instance already in use.
+        self._ingestion_handlers = dict(ingestion_handlers) if ingestion_handlers is not None else dict(INGESTION_HANDLERS)
+
+    def register_ingestion_handler(self, kind: str, handler: Callable[[bytes, bool, ExtractionConfig], str]) -> None:
+        """Register an ingestion route scoped to this DocumentParser instance only."""
+        self._ingestion_handlers[kind] = handler
+
+    def extract(
+        self,
+        document_bytes: bytes,
+        schema: Schema,
+        is_image: bool = False,
+        rules: list = None,
+        kind: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """Extract information from a document matching the schema.
+
+        kind overrides routing when set (must be registered via register_ingestion_handler
+        first, e.g. kind="docx"). Otherwise routing falls back to is_image, as before.
+        """
+        structured_mode = any(f.type == "list" for f in schema.fields)
+        resolved_kind = kind or ("image" if is_image else "pdf")
+        handler = self._resolve_handler(resolved_kind)
+        logger.info("Extracting schema=%r via kind=%r (structured_mode=%s)", schema.name, resolved_kind, structured_mode)
+
+        cache_key = None
+        if self.cache is not None and not rules:
+            # Keyed on the handler itself, not just the kind name — two DocumentParser
+            # instances can register different handlers under the same kind, and a
+            # shared cache must not conflate them.
+            cache_key = make_cache_key(document_bytes, schema, resolved_kind, self.config, handler)
+            cached = self.cache.get(cache_key)
+            if cached is not None:
+                logger.info("Cache hit for schema=%r.", schema.name)
+                return cached
+
+        doc_text = handler(document_bytes, structured_mode, self.config)
+        is_truncated, truncation_reason = self._check_truncation(document_bytes, resolved_kind)
+        if is_truncated:
+            logger.warning(truncation_reason)
+
+        if not doc_text.strip():
+            raise EmptyDocumentError("Could not extract any text from the document.")
+
+        chunks = chunk_document_text(doc_text, max_tokens=self.config.chunk_max_tokens)
+        logger.info("Document split into %d chunk(s) for extraction.", len(chunks))
+        merged_data = self._run_extraction(chunks, schema)
+
+        result = self._build_result(schema, merged_data, doc_text, rules, is_truncated, truncation_reason)
+
+        if cache_key is not None:
+            self.cache.set(cache_key, result)
+
+        return result
+
+    async def aextract(
+        self,
+        document_bytes: bytes,
+        schema: Schema,
+        is_image: bool = False,
+        rules: list = None,
+        kind: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """Async wrapper around extract(). The work itself is still synchronous (OCR,
+        PDF parsing, and the LLM SDK calls are all blocking) — this runs it in a
+        background thread so an asyncio event loop (e.g. inside a FastAPI route) isn't
+        blocked while it happens, not because the underlying pipeline became non-blocking.
+        """
+        loop = asyncio.get_running_loop()
+        call = functools.partial(self.extract, document_bytes, schema, is_image=is_image, rules=rules, kind=kind)
+        return await loop.run_in_executor(None, call)
+
+    def _resolve_handler(self, kind: str) -> Callable[[bytes, bool, ExtractionConfig], str]:
+        handler = self._ingestion_handlers.get(kind)
+        if handler is None:
+            raise UnknownIngestionKindError(f"No ingestion handler registered for document kind {kind!r}")
+        return handler
+
+    def _check_truncation(self, document_bytes: bytes, kind: str) -> tuple[bool, Optional[str]]:
+        if kind != "pdf":
+            return False, None
+        try:
+            doc = pymupdf.open(stream=document_bytes, filetype="pdf")
+            try:
+                if len(doc) > self.config.max_pages:
+                    return True, f"Document is {len(doc)} pages long, truncated to {self.config.max_pages} pages."
+                return False, None
+            finally:
+                doc.close()
+        except Exception:
+            return False, None
+
+    def _run_extraction(self, chunks: List[str], schema: Schema) -> Dict[str, Any]:
+        prompt_template = compile_prompt(schema)
+        max_workers = self.config.max_concurrent_chunks
+
+        if max_workers <= 1 or len(chunks) <= 1:
+            # Default path: identical to the original sequential loop, no thread pool
+            # involved — keeps behavior (and mock call ordering in tests) unchanged.
+            parsed_chunks = [_parse_json_from_llm(self.client.extract(prompt_template, chunk)) for chunk in chunks]
+        else:
+            def process(chunk: str) -> Dict[str, Any]:
+                return _parse_json_from_llm(self.client.extract(prompt_template, chunk))
+
+            with ThreadPoolExecutor(max_workers=max_workers) as executor:
+                parsed_chunks = list(executor.map(process, chunks))
+
+        return _merge_extracted_data(parsed_chunks, chunks, schema)
+
+    def _build_result(
+        self,
+        schema: Schema,
+        merged_data: Dict[str, Any],
+        doc_text: str,
+        rules: Optional[list],
+        is_truncated: bool,
+        truncation_reason: Optional[str],
+    ) -> Dict[str, Any]:
+        issues = cross_check(schema, merged_data, rules) + validate_field_constraints(schema, merged_data)
+        issues_by_field: Dict[str, List[Issue]] = {}
+        for issue in issues:
+            issues_by_field.setdefault(issue.field, []).append(issue)
+        issue_kind_to_flag = {
+            "cross_check": "failed_check",
+            "missing_required": "missing_required",
+            "invalid_format": "invalid_format",
+        }
+
+        final_result: Dict[str, Any] = {
+            "_meta": {
+                "truncated": is_truncated,
+                "truncation_reason": truncation_reason,
+            }
+        }
+        grounding_text = _strip_layout_tags(doc_text)
+        for f in schema.fields:
+            value = merged_data.get(f.name, None)
+            flags = []
+            confidence = "low"
+
+            if _is_present(value):
+                if check_substring(value, grounding_text, numeric=f.is_numeric, date=f.is_date):
+                    flags.append("grounded")
+                    confidence = "high"
+                else:
+                    flags.append("ungrounded")
+
+            for issue in issues_by_field.get(f.name, []):
+                flag = issue_kind_to_flag.get(issue.kind, "failed_check")
+                if flag not in flags:
+                    flags.append(flag)
+
+            final_result[f.name] = {
+                "value": value,
+                "confidence": confidence,
+                "flags": flags,
+            }
+
+        return final_result