feat: optimize code + add json key extractor

__init__.py

@@ -1,4 +1,10 @@
-from core.converter import ConversionError, ConversionResult, DocumentConverter, SUPPORTED_EXTENSIONS
+"""Public re-exports for the MarkItDown conversion core."""
+from core import (
+    ConversionError,
+    ConversionResult,
+    DocumentConverter,
+    SUPPORTED_EXTENSIONS,
+)
 
 __all__ = [
     "ConversionError",

api/server.py

@@ -41,7 +41,6 @@ import os
 import time
 import uuid
 from contextlib import asynccontextmanager
-from datetime import datetime, timezone
 from pathlib import Path
 from typing import Annotated, Any, Dict, List, Optional
 from urllib.parse import urlparse
@@ -55,10 +54,6 @@ from fastapi.middleware.gzip import GZipMiddleware
 from prometheus_client import make_asgi_app
 from pydantic import BaseModel, Field, field_validator
 
-from core import ConversionError, ConversionResult, DocumentConverter, SUPPORTED_EXTENSIONS
-from extraction.generic_json_extractor import extract
-from extraction.keys_extractor import Extractor as KeysExtractor
-from logger import get_logger
 from app.api.routes import router as pdf_router
 from app.core.auth import require_api_key
 from app.core.config import settings
@@ -67,30 +62,34 @@ from app.core.logging import configure_logging
 from app.core.rate_limit import RateLimitMiddleware
 from app.services.ping import start_self_ping
 from app.utils.cleanup import cleanup_loop
+from core import ConversionError, ConversionResult, DocumentConverter, SUPPORTED_EXTENSIONS
+from extraction.generic_json_extractor import extract
+from extraction.keys_extractor import Extractor as KeysExtractor
+from logger import get_logger
 
 configure_logging()
 logger = get_logger(__name__)
 
-_START_TIME = time.time()
 APP_NAME = "reconciliation-file-processing-service"
+APP_VERSION = "2.2.0"
 
 MAX_UPLOAD_BYTES = int(os.getenv("MAX_UPLOAD_BYTES", str(100 * 1024 * 1024)))
-
 MAX_WORKERS = min(32, (os.cpu_count() or 1) + 4)
-_thread_pool = concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS)
 
+_thread_pool = concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS)
 _converter = DocumentConverter()
 
 logger.info("thread_pool_initialised", workers=MAX_WORKERS)
 
 
+
 # ---------------------------------------------------------------------------
 # Lifespan
 # ---------------------------------------------------------------------------
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    logger.info("api_starting", version="2.2.0", host="0.0.0.0:7860")
+    logger.info("api_starting", version=APP_VERSION, host="0.0.0.0:7860")
     ping_task = await start_self_ping()
     cleanup_task = asyncio.create_task(cleanup_loop())
     yield
@@ -110,13 +109,13 @@ app = FastAPI(
         "Unified API for document-to-Markdown conversion (Microsoft MarkItDown + RapidOCR) "
         "and PDF-to-image conversion with async job support."
     ),
-    version="2.2.0",
+    version=APP_VERSION,
     docs_url="/docs",
     redoc_url="/redoc",
     lifespan=lifespan,
 )
 
-# -- Middleware --
+# -- Middleware (order matters: outermost first) --
 app.add_middleware(RateLimitMiddleware)
 app.add_middleware(GZipMiddleware, minimum_size=1000)
 app.add_middleware(
@@ -149,12 +148,16 @@ async def request_context_middleware(request: Request, call_next):
 
 
 # -- Exception handlers --
+def _request_id(request: Request) -> str:
+    return getattr(request.state, "request_id", "")
+
+
 @app.exception_handler(AppException)
 async def app_exception_handler(request: Request, exc: AppException):
     logger.warning("app_exception", detail=exc.detail, status_code=exc.status_code)
     return JSONResponse(
         status_code=exc.status_code,
-        content={"error": exc.detail, "request_id": getattr(request.state, "request_id", "")},
+        content={"error": exc.detail, "request_id": _request_id(request)},
     )
 
 
@@ -163,14 +166,14 @@ async def generic_exception_handler(request: Request, exc: Exception):
     logger.exception("unhandled_exception", exc_info=exc)
     return JSONResponse(
         status_code=500,
-        content={"error": "Internal server error", "request_id": getattr(request.state, "request_id", "")},
+        content={"error": "Internal server error", "request_id": _request_id(request)},
     )
 
 
 # -- Root-level routes (only / , /ping) --
 @app.get("/", tags=["System"], summary="Root", include_in_schema=False)
 async def root():
-    return {"service": APP_NAME, "version": "2.2.0", "status": "running"}
+    return {"service": APP_NAME, "version": APP_VERSION, "status": "running"}
 
 
 @app.get("/ping", tags=["System"], summary="Ping", include_in_schema=False)
@@ -178,6 +181,7 @@ async def ping():
     return {"message": f"{APP_NAME} is running..."}
 
 
+
 # ---------------------------------------------------------------------------
 # MarkItDown router — all routes under /api/v1/markitdown
 # ---------------------------------------------------------------------------
@@ -398,10 +402,9 @@ async def convert_file(
 
     parsed_mappings: Optional[Dict[str, Any]] = None
     if mappings:
-        import json as _json
         try:
-            parsed_mappings = _json.loads(mappings)
-        except _json.JSONDecodeError:
+            parsed_mappings = json.loads(mappings)
+        except json.JSONDecodeError:
             raise HTTPException(status_code=400, detail={"success": False, "message": "Invalid JSON in mappings parameter."})
 
     logger.info("convert_file", filename=file.filename)

app/core/logging.py

@@ -2,19 +2,76 @@ from __future__ import annotations
 
 import logging
 import sys
+from typing import Any, MutableMapping
 
 import structlog
 
 from app.core.config import settings
 
+_NOISY_LOGGERS: tuple[str, ...] = (
+    "uvicorn.access",
+    "uvicorn.error",
+    "multipart",
+    "httpx",
+    "httpcore",
+    "redis",
+    "PIL",
+    "asyncio",
+)
 
-def configure_logging() -> None:
-    pre_chain = [
-        structlog.contextvars.merge_contextvars,
-        structlog.stdlib.add_log_level,
-        structlog.processors.TimeStamper(fmt="iso"),
-    ]
+_BASE_PRE_CHAIN: tuple = (
+    structlog.contextvars.merge_contextvars,
+    structlog.stdlib.add_log_level,
+    structlog.processors.TimeStamper(fmt="iso"),
+)
+
+
+def _scrub_value(value: Any) -> Any:
+    """Return a size placeholder for binary values, otherwise pass through."""
+    if isinstance(value, (bytes, bytearray, memoryview)):
+        return f"<{len(bytes(value))} bytes>"
+    return value
+
+
+def _scrub_binary(_logger: Any, _method: str, event_dict: MutableMapping[str, Any]) -> MutableMapping[str, Any]:
+    """Replace any bytes / bytearray values in event_dict with a size placeholder.
+
+    structlog will happily serialise raw bytes to JSON (escaped) or to the
+    console renderer, but a single misrouted PDF or image stream can balloon
+    a log line into megabytes. This guard catches the leak at the boundary.
+    """
+    for key, value in list(event_dict.items()):
+        event_dict[key] = _scrub_value(value)
+    return event_dict
+
+
+def _consume_positional_args(_logger: Any, _method: str, event_dict: MutableMapping[str, Any]) -> MutableMapping[str, Any]:
+    """Format `logger.info('msg %s', arg)` style calls safely.
 
+    structlog stores the trailing args under `positional_args` and leaves
+    the format string in `event` until a later renderer (e.g.
+    `PositionalArgumentsFormatter`) consumes it. We do that here, scrubbing
+    bytes so a stray PDF stream can never end up rendered as `b'\\x89PNG...'`.
+    """
+    pos_args = event_dict.pop("positional_args", None)
+    if not pos_args:
+        return event_dict
+
+    event = event_dict.get("event")
+    if not isinstance(event, str) or "%" not in event:
+        return event_dict
+
+    safe_args = tuple(_scrub_value(a) for a in pos_args)
+    try:
+        event_dict["event"] = event % safe_args
+    except (TypeError, ValueError):
+        # Fall back to the raw event if the format string is malformed.
+        pass
+    return event_dict
+
+
+def configure_logging() -> None:
+    pre_chain = list(_BASE_PRE_CHAIN)
     is_prod = settings.ENV == "production"
 
     if is_prod:
@@ -26,6 +83,8 @@ def configure_logging() -> None:
     structlog.configure(
         processors=[
             structlog.stdlib.filter_by_level,
+            _consume_positional_args,
+            _scrub_binary,
             *pre_chain,
             structlog.stdlib.ProcessorFormatter.wrap_for_formatter,
         ],
@@ -38,6 +97,8 @@ def configure_logging() -> None:
         foreign_pre_chain=pre_chain,
         processors=[
             structlog.stdlib.ProcessorFormatter.remove_processors_meta,
+            _consume_positional_args,
+            _scrub_binary,
             renderer,
         ],
     )
@@ -49,14 +110,8 @@ def configure_logging() -> None:
     root_logger.handlers = [handler]
     root_logger.setLevel(settings.LOG_LEVEL.upper())
 
-    for lib in (
-        "uvicorn.access",
-        "uvicorn.error",
-        "multipart",
-        "httpx",
-        "httpcore",
-        "redis",
-        "PIL",
-        "asyncio",
-    ):
+    for lib in _NOISY_LOGGERS:
         logging.getLogger(lib).setLevel(logging.WARNING)
+
+
+

app/core/rate_limit.py

@@ -77,18 +77,21 @@ def _is_trusted_proxy(host: str) -> bool:
 def _extract_client_ip(request: Request) -> str:
     direct_host = request.client.host if request.client else None
 
-    if direct_host and _is_trusted_proxy(direct_host):
-        forwarded_for = request.headers.get("x-forwarded-for", "")
-        if forwarded_for:
-            for candidate in reversed([ip.strip() for ip in forwarded_for.split(",")]):
-                if candidate and not _is_trusted_proxy(candidate):
-                    return candidate
-
-        real_ip = request.headers.get("x-real-ip", "").strip()
-        if real_ip and not _is_trusted_proxy(real_ip):
-            return real_ip
-
-    return direct_host or "unknown"
+    if not direct_host or not _is_trusted_proxy(direct_host):
+        return direct_host or "unknown"
+
+    forwarded_for = request.headers.get("x-forwarded-for", "")
+    if forwarded_for:
+        # Trust the rightmost non-proxy IP; that's the closest hop to us.
+        for candidate in (ip.strip() for ip in reversed(forwarded_for.split(","))):
+            if candidate and not _is_trusted_proxy(candidate):
+                return candidate
+
+    real_ip = request.headers.get("x-real-ip", "").strip()
+    if real_ip and not _is_trusted_proxy(real_ip):
+        return real_ip
+
+    return direct_host
 
 
 class RateLimitMiddleware(BaseHTTPMiddleware):

app/services/conversion.py

@@ -26,6 +26,24 @@ _render_pool = ThreadPoolExecutor(
 )
 
 
+def _format_str(fmt) -> str:
+    """Normalise an ImageFormat enum / string to its raw uppercase value."""
+    return fmt.value if hasattr(fmt, "value") else str(fmt)
+
+
+def _save_kwargs(fmt: str, quality: int) -> dict:
+    """PIL save kwargs for the requested output format."""
+    if fmt == "JPEG":
+        return {"quality": quality, "optimize": True}
+    if fmt == "PNG":
+        return {"optimize": True}
+    return {}
+
+
+def _ext_for(fmt: str) -> str:
+    return "jpg" if fmt == "JPEG" else "png"
+
+
 class ConversionService:
     def __init__(self) -> None:
         self._output_dir = Path(settings.OUTPUT_DIR)
@@ -120,11 +138,11 @@ class ConversionService:
             for page_idx in ordered_indices
         ]
 
+        raw_results = await asyncio.gather(*render_coros, return_exceptions=True)
+
         results_by_idx: dict[int, PageResult] = {}
         errors: list[str] = []
 
-        raw_results = await asyncio.gather(*render_coros, return_exceptions=True)
-
         for page_idx, outcome in zip(ordered_indices, raw_results):
             if isinstance(outcome, Exception):
                 logger.error("page_render_failed", job_id=job_id, page=page_idx + 1, error=str(outcome))
@@ -149,7 +167,7 @@ def _validate_and_get_indices(
     from pypdf import PdfReader
     from pypdf.errors import PdfReadError
 
-    fmt = params.format.value if hasattr(params.format, "value") else str(params.format)
+    fmt = _format_str(params.format)
     if fmt not in _SUPPORTED_FORMATS:
         raise InvalidParameterError(
             f"Unsupported format '{fmt}'. Only JPEG (JPG) and PNG are supported."
@@ -159,8 +177,7 @@ def _validate_and_get_indices(
         reader = PdfReader(io.BytesIO(pdf_bytes))
         if reader.is_encrypted:
             from pypdf import PasswordType
-            result = reader.decrypt("")
-            if result == PasswordType.NOT_DECRYPTED:
+            if reader.decrypt("") == PasswordType.NOT_DECRYPTED:
                 raise ConversionError("PDF is password-protected. Please provide an unlocked PDF.")
         total_pages = len(reader.pages)
     except ConversionError:
@@ -197,6 +214,28 @@ def _split_pdf_pages(pdf_bytes: bytes, page_indices: List[int]) -> dict[int, byt
     return blobs
 
 
+def _normalise_for_jpeg(img: Image.Image) -> Image.Image:
+    """Return an RGB image suitable for JPEG output, flattening alpha onto white."""
+    mode = img.mode
+    if mode == "RGB":
+        return img
+    if mode in ("RGBA", "LA", "P"):
+        if mode == "P":
+            img = img.convert("RGBA")
+        bg = Image.new("RGB", img.size, (255, 255, 255))
+        mask = img.split()[-1] if mode in ("RGBA", "LA") else None
+        bg.paste(img, mask=mask)
+        return bg
+    return img.convert("RGB")
+
+
+def _normalise_for_png(img: Image.Image) -> Image.Image:
+    """Return an image whose mode PIL accepts for PNG (RGB / RGBA / L)."""
+    if img.mode in ("RGB", "RGBA", "L"):
+        return img
+    return img.convert("RGB")
+
+
 def _render_single_page(
     page_idx: int,
     page_pdf_bytes: bytes,
@@ -207,7 +246,7 @@ def _render_single_page(
     from pdf2image import convert_from_bytes
     from pdf2image.exceptions import PDFInfoNotInstalledError, PDFPageCountError, PDFSyntaxError
 
-    fmt = params.format.value if hasattr(params.format, "value") else params.format
+    fmt = _format_str(params.format)
 
     try:
         pil_images = convert_from_bytes(
@@ -225,33 +264,17 @@ def _render_single_page(
         raise ConversionError(f"Page {page_idx + 1}: renderer returned no image")
 
     img = pil_images[0]
-
-    if fmt == "JPEG" and img.mode in ("RGBA", "LA", "P"):
-        if img.mode == "P":
-            img = img.convert("RGBA")
-        bg = Image.new("RGB", img.size, (255, 255, 255))
-        mask = img.split()[-1] if img.mode in ("RGBA", "LA") else None
-        bg.paste(img, mask=mask)
-        img = bg
-    elif fmt == "JPEG" and img.mode != "RGB":
-        img = img.convert("RGB")
-
+    if fmt == "JPEG":
+        img = _normalise_for_jpeg(img)
     if params.grayscale:
         img = img.convert("L")
 
-    ext = "jpg" if fmt == "JPEG" else "png"
+    ext = _ext_for(fmt)
     filename = f"page_{page_idx + 1:04d}.{ext}"
     out_path = job_dir / filename
 
-    save_kwargs: dict = {}
-    if fmt == "JPEG":
-        save_kwargs["quality"] = params.quality
-        save_kwargs["optimize"] = True
-    if fmt == "PNG":
-        save_kwargs["optimize"] = True
-
     try:
-        img.save(str(out_path), format=fmt, **save_kwargs)
+        img.save(str(out_path), format=fmt, **_save_kwargs(fmt, params.quality))
     except OSError as exc:
         raise ConversionError(f"Failed to write {filename}: {exc}") from exc
 
@@ -276,17 +299,17 @@ def _stitch_pages(
     Called only when split_page=False.
     Pages are placed top-to-bottom in document order with no gaps.
     """
-    fmt = params.format.value if hasattr(params.format, "value") else params.format
-    ext = "jpg" if fmt == "JPEG" else "png"
+    fmt = _format_str(params.format)
+    ext = _ext_for(fmt)
 
     images: list[Image.Image] = []
     for pr in sorted(page_results, key=lambda p: p.page_number):
         path = job_dir / Path(pr.download_url).name
         img = Image.open(str(path))
-        if fmt == "JPEG" and img.mode != "RGB":
-            img = img.convert("RGB")
-        elif fmt == "PNG" and img.mode not in ("RGB", "RGBA", "L"):
-            img = img.convert("RGB")
+        if fmt == "JPEG":
+            img = _normalise_for_jpeg(img)
+        else:
+            img = _normalise_for_png(img)
         images.append(img)
 
     if not images:
@@ -296,7 +319,8 @@ def _stitch_pages(
     total_height = sum(im.height for im in images)
 
     mode = images[0].mode
-    stitched = Image.new(mode, (total_width, total_height), color=(255, 255, 255) if mode == "RGB" else 255)
+    fill = (255, 255, 255) if mode == "RGB" else 255
+    stitched = Image.new(mode, (total_width, total_height), color=fill)
 
     y_offset = 0
     for img in images:
@@ -305,15 +329,7 @@ def _stitch_pages(
 
     out_filename = f"stitched.{ext}"
     out_path = job_dir / out_filename
-
-    save_kwargs: dict = {}
-    if fmt == "JPEG":
-        save_kwargs["quality"] = params.quality
-        save_kwargs["optimize"] = True
-    if fmt == "PNG":
-        save_kwargs["optimize"] = True
-
-    stitched.save(str(out_path), format=fmt, **save_kwargs)
+    stitched.save(str(out_path), format=fmt, **_save_kwargs(fmt, params.quality))
 
     stat = out_path.stat()
     return PageResult(
@@ -327,8 +343,7 @@ def _stitch_pages(
 
 
 def _guard_memory(page_count: int, dpi: int) -> None:
-    pixels_per_page = (dpi * 8.5) * (dpi * 11)
-    bytes_per_page = pixels_per_page * 3
+    bytes_per_page = (dpi * 8.5) * (dpi * 11) * 3
     estimated_mb = (bytes_per_page * page_count) / (1024 * 1024)
     if estimated_mb > _MAX_MEMORY_ESTIMATE_MB:
         raise InvalidParameterError(
@@ -339,3 +354,4 @@ def _guard_memory(page_count: int, dpi: int) -> None:
 
 
 conversion_service = ConversionService()
+

app/services/file_service.py

@@ -16,14 +16,16 @@ async def upload_file_to_service(
     content_type: str = "application/octet-stream",
 ) -> UploadedFileInfo:
     timeout = httpx.Timeout(settings.FILE_SERVICE_TIMEOUT, connect=settings.FILE_SERVICE_CONNECT_TIMEOUT)
+    headers = {
+        "X-API-Key": settings.FILE_SERVICE_API_KEY,
+        "Authorization": f"Bearer {settings.FILE_SERVICE_BEARER_TOKEN}",
+    }
+
     async with httpx.AsyncClient(timeout=timeout) as client:
         try:
             response = await client.post(
                 settings.FILE_SERVICE_URL,
-                headers={
-                    "X-API-Key": settings.FILE_SERVICE_API_KEY,
-                    "Authorization": f"Bearer {settings.FILE_SERVICE_BEARER_TOKEN}"
-                },
+                headers=headers,
                 files={"files": (filename, file_bytes, content_type)},
             )
             response.raise_for_status()
@@ -40,8 +42,7 @@ async def upload_file_to_service(
             logger.error("file_service_request_error", error=str(exc))
             raise FileServiceError(f"File service unreachable: {exc}") from exc
 
-    payload = response.json()
-    uploaded = _extract_uploaded_file(payload, filename, len(file_bytes))
+    uploaded = _extract_uploaded_file(response.json(), filename, len(file_bytes))
     logger.info("file_uploaded", file_id=uploaded.file_id, filename=filename)
     return uploaded
 
@@ -51,18 +52,16 @@ def _extract_uploaded_file(
     filename: str,
     size_bytes: int,
 ) -> UploadedFileInfo:
+    """Pick the first file entry out of the (variably-shaped) service response."""
     if isinstance(payload, list):
-        if not payload:
-            raise FileServiceError("File service returned an empty list")
-        first = payload[0]
-        if not isinstance(first, dict):
-            raise FileServiceError(f"Unexpected item type in file service list: {type(first)}")
-        entry = first.get("data") or first
+        if not payload or not isinstance(payload[0], dict):
+            raise FileServiceError("File service returned an empty or malformed list")
+        entry = payload[0].get("data") or payload[0]
         return _entry_to_info(entry, filename, size_bytes)
 
     if isinstance(payload, dict):
         nested = payload.get("files") or payload.get("data")
-        if isinstance(nested, list) and nested:
+        if isinstance(nested, list) and nested and isinstance(nested[0], dict):
             entry = nested[0].get("data") or nested[0]
             return _entry_to_info(entry, filename, size_bytes)
         return _entry_to_info(payload, filename, size_bytes)
@@ -71,29 +70,26 @@ def _extract_uploaded_file(
 
 
 def _entry_to_info(entry: dict, filename: str, size_bytes: int) -> UploadedFileInfo:
-    file_id = (
+    file_id = str(
         entry.get("fileid")
         or entry.get("file_id")
         or entry.get("id")
         or entry.get("_id")
         or ""
     )
-    file_url = (
+    file_url = str(
         entry.get("file_url")
         or entry.get("url")
         or entry.get("path")
         or ""
     )
-    resolved_filename = (
-        entry.get("filename")
-        or entry.get("name")
-        or filename
-    )
+    resolved_filename = str(entry.get("filename") or entry.get("name") or filename)
     resolved_size = int(entry.get("size") or entry.get("size_bytes") or size_bytes)
 
     return UploadedFileInfo(
-        file_id=str(file_id),
-        file_url=str(file_url),
-        filename=str(resolved_filename),
+        file_id=file_id,
+        file_url=file_url,
+        filename=resolved_filename,
         size_bytes=resolved_size,
     )
+

app/services/upload_orchestrator.py

@@ -34,6 +34,29 @@ from app.services.file_service import upload_file_to_service
 logger = structlog.get_logger(__name__)
 
 
+def _format_str(fmt) -> str:
+    return fmt.value if hasattr(fmt, "value") else str(fmt)
+
+
+def _resolve_local_path(download_url: str) -> Path:
+    parts = download_url.lstrip("/").split("/")
+    if len(parts) >= 5:
+        return Path(settings.OUTPUT_DIR) / parts[3] / parts[4]
+    raise ValueError(f"Cannot resolve local path from URL: {download_url}")
+
+
+def _build_filename(job_id: str, page: PageResult, ext: str) -> str:
+    """Public-facing filename for an uploaded page.
+
+    Stitched outputs are flattened to a single deterministic name; per-page
+    images use a zero-padded sequence number.
+    """
+    local = _resolve_local_path(page.download_url)
+    if local.stem == "stitched":
+        return f"{job_id}_stitched.{ext}"
+    return f"{job_id}_page_{page.page_number:04d}.{ext}"
+
+
 async def convert_and_upload(
     pdf_bytes: bytes,
     params: ConversionParams,
@@ -70,7 +93,7 @@ async def _upload_pages(
     job_id: str,
 ) -> tuple[List[PageResult], UploadSummary]:
     semaphore = asyncio.Semaphore(settings.UPLOAD_CONCURRENCY)
-    fmt = params.format.value if hasattr(params.format, "value") else str(params.format)
+    fmt = _format_str(params.format)
     ext = "jpg" if fmt == "JPEG" else "png"
     content_type = "image/jpeg" if fmt == "JPEG" else "image/png"
 
@@ -83,11 +106,9 @@ async def _upload_pages(
                 return pr
             try:
                 file_bytes = file_path.read_bytes()
-                if file_path.stem == "stitched":
-                    filename = f"{job_id}_stitched.{ext}"
-                else:
-                    filename = f"{job_id}_page_{pr.page_number:04d}.{ext}"
-                info = await upload_file_to_service(file_bytes, filename, content_type)
+                info = await upload_file_to_service(
+                    file_bytes, _build_filename(job_id, pr, ext), content_type
+                )
                 pr.file_id = info.file_id
                 pr.file_url = info.file_url
                 logger.debug("page_uploaded", job_id=job_id, page=pr.page_number, file_id=info.file_id)
@@ -109,11 +130,7 @@ async def _upload_pages(
             UploadedFileInfo(
                 file_id=p.file_id or "",
                 file_url=p.file_url or "",
-                filename=(
-                    f"{job_id}_stitched.{ext}"
-                    if Path(_resolve_local_path(p.download_url)).stem == "stitched"
-                    else f"{job_id}_page_{p.page_number:04d}.{ext}"
-                ),
+                filename=_build_filename(job_id, p, ext),
                 size_bytes=p.size_bytes,
             )
             for p in succeeded
@@ -123,9 +140,3 @@ async def _upload_pages(
     logger.info("upload_complete", job_id=job_id, uploaded=len(succeeded), failed=len(failed))
     return updated_pages, summary
 
-
-def _resolve_local_path(download_url: str) -> Path:
-    parts = download_url.lstrip("/").split("/")
-    if len(parts) >= 5:
-        return Path(settings.OUTPUT_DIR) / parts[3] / parts[4]
-    raise ValueError(f"Cannot resolve local path from URL: {download_url}")

app/utils/validators.py

@@ -15,9 +15,12 @@ from app.core.exceptions import FileTooLargeError, InvalidFileTypeError, Invalid
 logger = structlog.get_logger(__name__)
 
 _PDF_MAGIC = b"%PDF"
+_CHUNK_SIZE = 64 * 1024
 _MAX_URL_REDIRECTS = 5
-_ALLOWED_SCHEMES = {"http", "https"}
-_PRIVATE_RANGES = [
+_ALLOWED_SCHEMES = frozenset({"http", "https"})
+
+# SSRF guard: pre-collapsed private/loopback/link-local networks.
+_PRIVATE_RANGES: tuple[ipaddress.IPv4Network | ipaddress.IPv6Network, ...] = (
     ipaddress.ip_network("10.0.0.0/8"),
     ipaddress.ip_network("172.16.0.0/12"),
     ipaddress.ip_network("192.168.0.0/16"),
@@ -25,15 +28,18 @@ _PRIVATE_RANGES = [
     ipaddress.ip_network("169.254.0.0/16"),
     ipaddress.ip_network("::1/128"),
     ipaddress.ip_network("fc00::/7"),
-]
+)
+
+
+def _is_private_ip(addr: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool:
+    return any(addr in net for net in _PRIVATE_RANGES)
 
 
 async def read_and_validate_pdf(upload: UploadFile) -> bytes:
     buffer = io.BytesIO()
     total = 0
-    chunk_size = 64 * 1024
 
-    while chunk := await upload.read(chunk_size):
+    while chunk := await upload.read(_CHUNK_SIZE):
         total += len(chunk)
         if total > settings.MAX_FILE_SIZE_BYTES:
             raise FileTooLargeError(settings.MAX_FILE_SIZE_MB)
@@ -74,7 +80,7 @@ async def fetch_pdf_from_url(url: str) -> bytes:
             raise FileTooLargeError(settings.MAX_FILE_SIZE_MB)
 
         pdf_bytes = b""
-        async for chunk in response.aiter_bytes(chunk_size=64 * 1024):
+        async for chunk in response.aiter_bytes(chunk_size=_CHUNK_SIZE):
             pdf_bytes += chunk
             if len(pdf_bytes) > settings.MAX_FILE_SIZE_BYTES:
                 raise FileTooLargeError(settings.MAX_FILE_SIZE_MB)
@@ -114,8 +120,7 @@ def _validate_url_host(url: str) -> None:
     except (socket.gaierror, ValueError):
         raise InvalidParameterError(f"Could not resolve hostname: {hostname}")
 
-    for private_range in _PRIVATE_RANGES:
-        if addr in private_range:
-            raise InvalidParameterError(
-                "Requests to private or loopback IP addresses are not permitted (SSRF protection)"
-            )
+    if _is_private_ip(addr):
+        raise InvalidParameterError(
+            "Requests to private or loopback IP addresses are not permitted (SSRF protection)"
+        )

banner.py

@@ -10,14 +10,14 @@ _console = Console()
 SERVICE_NAME = os.getenv("SERVICE_NAME", "reconciliation-file-processing-service")
 API_VERSION  = os.getenv("API_VERSION", "2.2.0")
 ENVIRONMENT  = os.getenv("ENVIRONMENT", "production")
+_BANNER_ART  = pyfiglet.figlet_format("ValidOps", font="slant")
 
 
 def print_banner() -> None:
-    art = pyfiglet.figlet_format("ValidOps", font="slant")
-
-    _console.print(Text(art, style="bold cyan"))
+    _console.print(Text(_BANNER_ART, style="bold cyan"))
     _console.print(f"  [bold white]{'Service:':<14}[/bold white]  [cyan]{SERVICE_NAME}[/cyan]")
     _console.print(f"  [bold white]{'Version:':<14}[/bold white]  [cyan]v{API_VERSION}[/cyan]")
     _console.print(f"  [bold white]{'Environment:':<14}[/bold white]  [cyan]{ENVIRONMENT}[/cyan]")
     _console.print(Rule(style="dim cyan"))
     _console.print()
+

core/batch.py

@@ -15,11 +15,19 @@ from dataclasses import dataclass
 from pathlib import Path
 from typing import Callable, Optional, Sequence
 
-from .converter import ConversionError, ConversionResult, DocumentConverter, SUPPORTED_EXTENSIONS
 from logger import get_logger
 
+from .converter import (
+    ConversionError,
+    ConversionResult,
+    DocumentConverter,
+    SUPPORTED_EXTENSIONS,
+)
+
 logger = get_logger(__name__)
 
+_DEFAULT_MAX_WORKERS = min(8, (os.cpu_count() or 1) + 4)
+
 
 @dataclass
 class BatchReport:
@@ -54,7 +62,7 @@ class BatchProcessor:
     def __init__(
         self,
         converter: DocumentConverter,
-        max_workers: int = min(8, (os.cpu_count() or 1) + 4),
+        max_workers: int = _DEFAULT_MAX_WORKERS,
     ) -> None:
         self._converter = converter
         self._max_workers = max_workers
@@ -98,15 +106,18 @@ class BatchProcessor:
             "batch_processor | done | total=%d | succeeded=%d | failed=%d",
             total, len(results), len(errors),
         )
+        total_chars = sum(r.char_count for r in results)
+        total_words = sum(r.word_count for r in results)
+        total_duration = sum(r.duration_ms for r in results)
         return BatchReport(
             total=total,
             succeeded=len(results),
             failed=len(errors),
             results=results,
             errors=errors,
-            total_chars=sum(r.char_count for r in results),
-            total_words=sum(r.word_count for r in results),
-            total_duration_ms=sum(r.duration_ms for r in results),
+            total_chars=total_chars,
+            total_words=total_words,
+            total_duration_ms=total_duration,
         )
 
     def discover_files(

core/converter.py

@@ -23,46 +23,42 @@ Public classes
 from __future__ import annotations
 
 import hashlib
+import io
 import mimetypes
 import time
+from contextlib import contextmanager
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Optional
+from typing import Iterator, Optional
 from urllib.parse import urlparse
 
 from markitdown import MarkItDown
 
-from .ocr_engine import ocr_image, ocr_pdf
 from logger import get_logger
 
+from .ocr_engine import ocr_image, ocr_pdf
+
 logger = get_logger(__name__)
 
-# Extensions that route through RapidOCR rather than MarkItDown.
-IMAGE_EXTENSIONS = {
+IMAGE_EXTENSIONS: frozenset[str] = frozenset({
     ".jpg", ".jpeg", ".png", ".gif", ".bmp", ".webp", ".tiff",
-}
-
-IMAGE_MIME_PREFIXES = {"image/"}
+})
 
+IMAGE_MIME_PREFIXES: tuple[str, ...] = ("image/",)
 
-# ---------------------------------------------------------------------------
-# Supported formats
-# ---------------------------------------------------------------------------
-
-SUPPORTED_EXTENSIONS = {
+SUPPORTED_EXTENSIONS: frozenset[str] = frozenset({
     ".pdf", ".docx", ".doc", ".pptx", ".ppt",
     ".xlsx", ".xls", ".csv", ".json", ".xml",
     ".html", ".htm", ".txt", ".md", ".rst",
     ".jpg", ".jpeg", ".png", ".gif", ".bmp", ".webp", ".tiff",
     ".mp3", ".wav", ".ogg", ".flac",
     ".zip", ".epub",
-}
+})
+
 
 def _is_image(ext: str, mime: str) -> bool:
     """Return True when the input should be routed through RapidOCR."""
-    return ext.lower() in IMAGE_EXTENSIONS or any(
-        mime.startswith(p) for p in IMAGE_MIME_PREFIXES
-    )
+    return ext.lower() in IMAGE_EXTENSIONS or mime.startswith(IMAGE_MIME_PREFIXES)
 
 
 # ---------------------------------------------------------------------------
@@ -121,119 +117,87 @@ class DocumentConverter:
     def convert_file(self, path: str | Path) -> ConversionResult | ConversionError:
         """Convert a local file identified by *path*."""
         path = Path(path).resolve()
-        start = time.perf_counter()
 
         if not path.exists():
-            return ConversionError(
-                source=str(path),
-                error_type="FileNotFoundError",
-                message=f"File does not exist: {path}",
-                duration_ms=0.0,
-            )
+            return self._error(str(path), "FileNotFoundError",
+                               f"File does not exist: {path}", 0.0)
 
         file_size = path.stat().st_size
-        mime_type, _ = mimetypes.guess_type(str(path))
-        mime_type = mime_type or "application/octet-stream"
-
-        try:
-            if _is_image(path.suffix, mime_type):
-                markdown = ocr_image(str(path))
-            else:
-                markdown = self._engine.convert(str(path)).text_content
-                if not markdown.strip() and path.suffix.lower() == ".pdf":
-                    logger.info(
-                        "convert_file | no text from PDF, falling back to OCR | file=%s",
-                        path.name,
-                    )
-                    markdown = ocr_pdf(str(path))
-
-            elapsed = (time.perf_counter() - start) * 1000
-            return self._build_result(str(path), markdown, file_size, mime_type, elapsed)
-
-        except Exception as exc:
-            elapsed = (time.perf_counter() - start) * 1000
-            logger.error("convert_file | exception | file=%s | error=%s", path, exc, exc_info=True)
-            return ConversionError(
-                source=str(path),
-                error_type=type(exc).__name__,
-                message=str(exc),
-                duration_ms=elapsed,
-            )
+        mime_type = mimetypes.guess_type(str(path))[0] or "application/octet-stream"
+        ext = path.suffix
+
+        with self._timed() as elapsed:
+            try:
+                markdown = self._convert_to_markdown(ext, mime_type, str(path))
+            except Exception as exc:
+                logger.error("convert_file | exception | file=%s | error=%s",
+                             path, exc, exc_info=True)
+                return self._error(str(path), type(exc).__name__, str(exc), elapsed())
+
+        return self._build_result(str(path), markdown, file_size, mime_type, elapsed())
 
     def convert_url(self, url: str) -> ConversionResult | ConversionError:
         """Fetch and convert a public HTTP/HTTPS URL."""
         parsed = urlparse(url)
         if parsed.scheme not in {"http", "https"}:
-            return ConversionError(
-                source=url,
-                error_type="ValueError",
-                message=f"Unsupported URL scheme: {parsed.scheme!r}",
-                duration_ms=0.0,
-            )
+            return self._error(url, "ValueError",
+                               f"Unsupported URL scheme: {parsed.scheme!r}", 0.0)
 
-        start = time.perf_counter()
-        try:
-            url_ext = Path(urlparse(url).path).suffix.lower()
-            if url_ext in IMAGE_EXTENSIONS:
-                markdown = ocr_image(url)
-                mime_type = mimetypes.guess_type(url)[0] or "image/jpeg"
-            else:
-                result = self._engine.convert(url)
-                markdown = result.text_content
-                mime_type = "text/html"
-
-            elapsed = (time.perf_counter() - start) * 1000
-            return self._build_result(url, markdown, 0, mime_type, elapsed)
-
-        except Exception as exc:
-            elapsed = (time.perf_counter() - start) * 1000
-            logger.error("convert_url | exception | url=%s | error=%s", url, exc, exc_info=True)
-            return ConversionError(
-                source=url,
-                error_type=type(exc).__name__,
-                message=str(exc),
-                duration_ms=elapsed,
-            )
+        ext = Path(parsed.path).suffix.lower()
+        is_image = ext in IMAGE_EXTENSIONS
+        mime_type = mimetypes.guess_type(url)[0] or ("image/jpeg" if is_image else "text/html")
+
+        with self._timed() as elapsed:
+            try:
+                markdown = self._convert_to_markdown(ext, mime_type, url)
+            except Exception as exc:
+                logger.error("convert_url | exception | url=%s | error=%s",
+                             url, exc, exc_info=True)
+                return self._error(url, type(exc).__name__, str(exc), elapsed())
+
+        return self._build_result(url, markdown, 0, mime_type, elapsed())
 
     def convert_stream(self, data: bytes, filename: str) -> ConversionResult | ConversionError:
         """Convert raw bytes identified by *filename* (used for upload payloads)."""
-        import io
-
-        start = time.perf_counter()
-        mime_type, _ = mimetypes.guess_type(filename)
-        mime_type = mime_type or "application/octet-stream"
+        mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
         ext = Path(filename).suffix.lower()
 
-        try:
-            if _is_image(ext, mime_type):
-                markdown = ocr_image(data)
-            else:
-                result = self._engine.convert_stream(io.BytesIO(data), file_extension=ext)
-                markdown = result.text_content
-                if not markdown.strip() and ext == ".pdf":
-                    logger.info(
-                        "convert_stream | no text from PDF stream, falling back to OCR | filename=%s",
-                        filename,
-                    )
-                    markdown = ocr_pdf(data)
-
-            elapsed = (time.perf_counter() - start) * 1000
-            return self._build_result(filename, markdown, len(data), mime_type, elapsed)
-
-        except Exception as exc:
-            elapsed = (time.perf_counter() - start) * 1000
-            logger.error("convert_stream | exception | filename=%s | error=%s", filename, exc, exc_info=True)
-            return ConversionError(
-                source=filename,
-                error_type=type(exc).__name__,
-                message=str(exc),
-                duration_ms=elapsed,
-            )
+        with self._timed() as elapsed:
+            try:
+                markdown = self._convert_to_markdown(ext, mime_type, data)
+            except Exception as exc:
+                logger.error("convert_stream | exception | filename=%s | error=%s",
+                             filename, exc, exc_info=True)
+                return self._error(filename, type(exc).__name__, str(exc), elapsed())
+
+        return self._build_result(filename, markdown, len(data), mime_type, elapsed())
 
     # ------------------------------------------------------------------
     # Internal helpers
     # ------------------------------------------------------------------
 
+    def _convert_to_markdown(self, ext: str, mime_type: str, source) -> str:
+        """Dispatch to the right backend (OCR or MarkItDown) and return text.
+
+        For PDFs, if MarkItDown yields nothing, fall back to OCR.
+        `source` is a path/URL for image inputs, or raw bytes for streams.
+        """
+        if _is_image(ext, mime_type):
+            return ocr_image(source)
+
+        if isinstance(source, bytes):
+            markdown = self._engine.convert_stream(
+                io.BytesIO(source), file_extension=ext
+            ).text_content
+        else:
+            markdown = self._engine.convert(source).text_content
+
+        if not markdown.strip() and ext.lower() == ".pdf":
+            source_label = f"<{len(source)} bytes>" if isinstance(source, (bytes, bytearray)) else source
+            logger.info("convert | PDF text empty, falling back to OCR | source=%s", source_label)
+            markdown = ocr_pdf(source)
+        return markdown
+
     @staticmethod
     def _build_result(
         source: str,
@@ -242,17 +206,34 @@ class DocumentConverter:
         mime_type: str,
         elapsed: float,
     ) -> ConversionResult:
-        lines = markdown.splitlines()
-        words = markdown.split()
-        content_hash = hashlib.sha256(markdown.encode()).hexdigest()
         return ConversionResult(
             source=source,
             markdown=markdown,
             char_count=len(markdown),
-            word_count=len(words),
-            line_count=len(lines),
+            word_count=len(markdown.split()),
+            line_count=len(markdown.splitlines()),
             duration_ms=elapsed,
             file_size_bytes=file_size,
             mime_type=mime_type,
-            content_hash=content_hash,
+            content_hash=hashlib.sha256(markdown.encode()).hexdigest(),
         )
+
+    @staticmethod
+    def _error(source: str, error_type: str, message: str, duration_ms: float) -> ConversionError:
+        return ConversionError(
+            source=source,
+            error_type=error_type,
+            message=message,
+            duration_ms=duration_ms,
+        )
+
+    @staticmethod
+    @contextmanager
+    def _timed() -> Iterator:
+        """Context manager that yields a callable returning elapsed ms at any point."""
+        start = time.perf_counter()
+
+        def elapsed() -> float:
+            return (time.perf_counter() - start) * 1000.0
+
+        yield elapsed

core/ocr_engine.py

@@ -76,11 +76,14 @@ def _to_numpy(source) -> Union[np.ndarray, str]:
             img = img.convert("RGB")
         return np.array(img)
 
+    if isinstance(source, np.ndarray):
+        return source
+
     if isinstance(source, Image.Image):
         return _pil_to_array(source)
 
-    if isinstance(source, (bytes, bytearray)):
-        return _pil_to_array(Image.open(io.BytesIO(source)))
+    if isinstance(source, (bytes, bytearray, memoryview)):
+        return _pil_to_array(Image.open(io.BytesIO(bytes(source))))
 
     if isinstance(source, str):
         parsed = urlparse(source)
@@ -92,9 +95,6 @@ def _to_numpy(source) -> Union[np.ndarray, str]:
         # Local file path — RapidOCR accepts it directly.
         return source
 
-    if isinstance(source, np.ndarray):
-        return source
-
     raise TypeError(
         f"ocr_image expects bytes, str (URL or path), numpy.ndarray, or PIL.Image; "
         f"received {type(source).__name__!r}"

extraction/json_extractor.py

@@ -5,6 +5,8 @@ Supports only CSV, XLS, XLSX file types for JSON extraction.
 Returns error for unsupported file types.
 """
 
+from __future__ import annotations
+
 import io
 import warnings
 from pathlib import Path
@@ -16,15 +18,17 @@ from logger import get_logger
 
 logger = get_logger(__name__)
 
-# Supported file extensions for JSON extraction
-SUPPORTED_EXTENSIONS = {'.csv', '.xls', '.xlsx'}
+SUPPORTED_EXTENSIONS: frozenset[str] = frozenset({".csv", ".xls", ".xlsx"})
+
+MAX_CSV_ROWS = 100_000
+MAX_EXCEL_ROWS = 50_000
+MAX_CELL_COUNT = 2_000_000
 
-# Resource limits to prevent abuse
-from app.core.config import settings as _app_settings
-MAX_FILE_SIZE_BYTES = _app_settings.MAX_FILE_SIZE_BYTES
-MAX_CSV_ROWS = 100000
-MAX_EXCEL_ROWS = 50000
-MAX_MEMORY_ROWS = 100000
+try:
+    from app.core.config import settings as _app_settings
+    MAX_FILE_SIZE_BYTES: int = _app_settings.MAX_FILE_SIZE_BYTES
+except Exception:
+    MAX_FILE_SIZE_BYTES = 50 * 1024 * 1024
 
 
 def _validate_file_size(size: int) -> Optional[str]:
@@ -34,15 +38,29 @@ def _validate_file_size(size: int) -> Optional[str]:
 
 
 def _check_memory_usage(rows: int, cols: int) -> Optional[str]:
-    approx_mb = (rows * cols * 50) / (1024 * 1024)
-    if rows * cols > MAX_MEMORY_ROWS * 20:
+    cell_count = rows * cols
+    if cell_count > MAX_CELL_COUNT:
+        approx_mb = (cell_count * 50) / (1024 * 1024)
         return f"Data size too large (approx {approx_mb:.1f} MB). Too many cells: {rows}x{cols}"
     return None
 
 
+def _read_dataframe(ext: str, file_path: Union[str, Path], file_data: Optional[bytes]) -> "pd.DataFrame":
+    """Read a CSV/Excel file from either an in-memory buffer or a path."""
+    if ext == ".csv":
+        if file_data:
+            return pd.read_csv(io.BytesIO(file_data), nrows=MAX_CSV_ROWS + 1, low_memory=False)
+        return pd.read_csv(file_path, nrows=MAX_CSV_ROWS + 1, low_memory=False)
+
+    engine = "openpyxl" if ext == ".xlsx" else "xlrd"
+    if file_data:
+        return pd.read_excel(io.BytesIO(file_data), engine=engine)
+    return pd.read_excel(file_path, engine=engine)
+
+
 def extract_json_from_file(
     file_path: Union[str, Path],
-    file_data: Optional[bytes] = None
+    file_data: Optional[bytes] = None,
 ) -> Dict[str, Any]:
     """
     Extract JSON data from structured files (CSV, XLS, XLSX).
@@ -64,114 +82,66 @@ def extract_json_from_file(
     if ext not in SUPPORTED_EXTENSIONS:
         return {
             "error": f"Unsupported file type: {ext}. Supported: {', '.join(sorted(SUPPORTED_EXTENSIONS))}",
-            "file_type": ext
+            "file_type": ext,
         }
 
-    # Validate file size if we have raw data
+    # Validate file size
     if file_data is not None:
-        size_error = _validate_file_size(len(file_data))
-        if size_error:
-            return {"error": size_error, "file_type": ext}
-    elif Path(file_path).exists():
-        size_error = _validate_file_size(Path(file_path).stat().st_size)
-        if size_error:
-            return {"error": size_error, "file_type": ext}
+        source_bytes = len(file_data)
+    else:
+        path = Path(file_path)
+        source_bytes = path.stat().st_size if path.exists() else 0
+    size_error = _validate_file_size(source_bytes)
+    if size_error:
+        return {"error": size_error, "file_type": ext}
 
     try:
         with warnings.catch_warnings():
             warnings.simplefilter("ignore", UserWarning)
-
-            if ext == '.csv':
-                if file_data:
-                    stream = io.BytesIO(file_data)
-                    # Peek for encoding detection
-                    sample = stream.read(1024)
-                    stream.seek(0)
-                    df = pd.read_csv(stream, nrows=MAX_CSV_ROWS + 1, low_memory=False)
-                else:
-                    df = pd.read_csv(file_path, nrows=MAX_CSV_ROWS + 1, low_memory=False)
-            else:
-                if file_data:
-                    df = pd.read_excel(io.BytesIO(file_data), engine='openpyxl' if ext == '.xlsx' else 'xlrd')
-                else:
-                    df = pd.read_excel(file_path, engine='openpyxl' if ext == '.xlsx' else 'xlrd')
-
-        # Enforce row limits to prevent memory exhaustion
-        max_rows = MAX_EXCEL_ROWS if ext != '.csv' else MAX_CSV_ROWS
-        if len(df) > max_rows:
-            return {
-                "error": f"File contains {len(df)} rows, exceeds limit of {max_rows}",
-                "file_type": ext,
-                "row_count": len(df)
-            }
-
-        # Additional memory guard
-        mem_error = _check_memory_usage(len(df), len(df.columns))
-        if mem_error:
-            return {"error": mem_error, "file_type": ext}
-
-        # Efficient conversion to JSON-serializable format
-        result = {
-            "success": True,
-            "file_type": ext,
-            "data": {
-                "columns": list(df.columns),
-                "rows": df.where(pd.notnull(df), None).to_dict(orient='records'),
-                "shape": [len(df), len(df.columns)],
-                "dtypes": {col: str(dtype) for col, dtype in df.dtypes.items()}
-            }
-        }
-
-        logger.info("Extracted JSON from %s: %d rows, %d cols", ext, len(df), len(df.columns))
-        return result
-
+            df = _read_dataframe(ext, file_path, file_data)
     except pd.errors.EmptyDataError:
-        return {
-            "error": f"File is empty or has no data",
-            "file_type": ext
-        }
+        return {"error": "File is empty or has no data", "file_type": ext}
     except MemoryError:
-        return {
-            "error": "Out of memory processing file",
-            "file_type": ext
-        }
+        return {"error": "Out of memory processing file", "file_type": ext}
     except Exception as exc:
         logger.exception("JSON extraction failed for %s", ext)
         return {
-            "error": f"Processing failed: {str(exc)}",
+            "error": f"Processing failed: {exc}",
             "file_type": ext,
-            "exception_type": type(exc).__name__
+            "exception_type": type(exc).__name__,
         }
 
+    max_rows = MAX_EXCEL_ROWS if ext != ".csv" else MAX_CSV_ROWS
+    if len(df) > max_rows:
+        return {
+            "error": f"File contains {len(df)} rows, exceeds limit of {max_rows}",
+            "file_type": ext,
+            "row_count": len(df),
+        }
 
+    mem_error = _check_memory_usage(len(df), len(df.columns))
+    if mem_error:
+        return {"error": mem_error, "file_type": ext}
 
+    logger.info("Extracted JSON from %s: %d rows, %d cols", ext, len(df), len(df.columns))
+    return {
+        "success": True,
+        "file_type": ext,
+        "data": {
+            "columns": list(df.columns),
+            "rows": df.where(pd.notnull(df), None).to_dict(orient="records"),
+            "shape": [len(df), len(df.columns)],
+            "dtypes": {col: str(dtype) for col, dtype in df.dtypes.items()},
+        },
+    }
 
 
 def is_supported_file_type(file_path: Union[str, Path]) -> bool:
-    """
-    Check if file type is supported for JSON extraction.
-    
-    Parameters
-    ----------
-    file_path : Union[str, Path]
-        Path to the file or filename with extension
-        
-    Returns
-    -------
-    bool
-        True if supported, False otherwise
-    """
-    extension = Path(file_path).suffix.lower()
-    return extension in SUPPORTED_EXTENSIONS
+    """Check if file type is supported for JSON extraction."""
+    return Path(file_path).suffix.lower() in SUPPORTED_EXTENSIONS
 
 
-def get_supported_extensions() -> list:
-    """
-    Get list of supported file extensions for JSON extraction.
-    
-    Returns
-    -------
-    list
-        Supported file extensions
-    """
+def get_supported_extensions() -> list[str]:
+    """Get sorted list of supported file extensions for JSON extraction."""
     return sorted(SUPPORTED_EXTENSIONS)
+

extraction/spacy_extractor.py

@@ -43,15 +43,21 @@ VALID_SPACY_LABELS: Dict[str, str] = {
 }
 
 
+_WHITESPACE_RE = re.compile(r"\s+")
+_CURRENCY_RE = re.compile(r"[$€£¥₹]")
+_NON_NUMERIC_RE = re.compile(r"[^\d.]")
+_DATE_SEP_RE = re.compile(r"[/.]")
+
 # ---------------------------------------------------------------------------
 # spaCy singleton
 # ---------------------------------------------------------------------------
 
 _nlp_lock = threading.Lock()
-_nlp      = None
+_nlp: Any = None
 
 
-def _get_nlp():
+def _get_nlp() -> Any:
+    """Return the shared spaCy pipeline, initialising it on first call."""
     global _nlp
     if _nlp is not None:
         return _nlp
@@ -66,11 +72,20 @@ def _get_nlp():
     return _nlp
 
 
+def clean_text(text: str) -> str:
+    """Normalise whitespace on raw text before handing it to spaCy.
+
+    Defined as a module-level function so it is not shadowed by local
+    variables named `cleaned` in the public extract_* functions.
+    """
+    return _WHITESPACE_RE.sub(" ", text).strip()
+
+
 # ---------------------------------------------------------------------------
 # Normalizer registry
 # ---------------------------------------------------------------------------
 
-_norm_lock  = threading.Lock()
+_normalizer_lock = threading.Lock()
 _NORMALIZERS: Dict[str, Callable[[str], str]] = {
     "strip":               lambda s: s.strip(),
     "upper":               lambda s: s.upper(),
@@ -78,16 +93,16 @@ _NORMALIZERS: Dict[str, Callable[[str], str]] = {
     "remove_commas":       lambda s: s.replace(",", ""),
     "remove_spaces":       lambda s: s.replace(" ", ""),
     "remove_newlines":     lambda s: s.replace("\n", " ").replace("\r", ""),
-    "collapse_whitespace": lambda s: re.sub(r"\s+", " ", s).strip(),
-    "remove_currency":     lambda s: re.sub(r"[$€£¥₹]", "", s),
-    "remove_non_numeric":  lambda s: re.sub(r"[^\d.]", "", s),
-    "normalize_date_sep":  lambda s: re.sub(r"[/.]", "-", s),
+    "collapse_whitespace": lambda s: _WHITESPACE_RE.sub(" ", s).strip(),
+    "remove_currency":     lambda s: _CURRENCY_RE.sub("", s),
+    "remove_non_numeric":  lambda s: _NON_NUMERIC_RE.sub("", s),
+    "normalize_date_sep":  lambda s: _DATE_SEP_RE.sub("-", s),
 }
 
 
 def register_normalizer(name: str, fn: Callable[[str], str]) -> None:
     """Register a custom normalizer. Thread-safe, overwrites silently."""
-    with _norm_lock:
+    with _normalizer_lock:
         _NORMALIZERS[name] = fn
 
 
@@ -99,8 +114,7 @@ def _apply_normalizers(value: Optional[str], normalize: Any) -> Optional[str]:
     if isinstance(normalize, str):
         normalize = [normalize]
     for key in normalize:
-        with _norm_lock:
-            fn = _NORMALIZERS.get(key)
+        fn = _NORMALIZERS.get(key)
         if fn is None:
             logger.warning("Unknown normalizer %r — skipped", key)
             continue
@@ -132,8 +146,8 @@ def register_resolver(
 # Regex resolver
 # ---------------------------------------------------------------------------
 
-def _build_flags(rule: Dict[str, Any]) -> re.RegexFlag:
-    flags = re.RegexFlag(0)
+def _build_flags(rule: Dict[str, Any]) -> int:
+    flags = 0
     for name in rule.get("flags", []):
         obj = getattr(re, name.upper(), None)
         if obj is None:
@@ -167,7 +181,7 @@ def _resolve_regex(rule: Dict[str, Any], text: str) -> Optional[str]:
     strip_chars   = rule.get("strip_chars", "")
     fallbacks     = rule.get("fallback_patterns", [])
 
-    for pat in [primary, *fallbacks]:
+    for pat in (primary, *fallbacks):
         try:
             matches = list(re.finditer(pat, text, flags))
         except re.error as exc:
@@ -177,22 +191,15 @@ def _resolve_regex(rule: Dict[str, Any], text: str) -> Optional[str]:
         if not matches:
             continue
 
-        try:
-            target_matches = [matches[match_index]]
-        except IndexError:
-            target_matches = [matches[-1]]
-
-        for m in target_matches:
-            exists, result = _try_group(m, capture_group)
-            if not exists:
-                break
-            if result is None:
-                break
-            result = _apply_normalizers(result, normalize)
-            if result is None:
-                break
-            result = result.strip(strip_chars) if strip_chars else result.strip()
-            return result or None
+        target = matches[match_index] if match_index < len(matches) else matches[-1]
+        exists, result = _try_group(target, capture_group)
+        if not exists or result is None:
+            return None
+        result = _apply_normalizers(result, normalize)
+        if result is None:
+            return None
+        result = result.strip(strip_chars) if strip_chars else result.strip()
+        return result or None
 
     return None
 
@@ -254,7 +261,7 @@ def _resolve_entity(rule: Dict[str, Any], doc: Any) -> Optional[str]:
     labels = rule.get("label")
     if isinstance(labels, str):
         labels = [labels]
-    labels = set(labels or [])
+    label_set = set(labels or [])
 
     match_index   = rule.get("match_index", 0)
     min_length    = rule.get("min_length", 1)
@@ -264,7 +271,7 @@ def _resolve_entity(rule: Dict[str, Any], doc: Any) -> Optional[str]:
 
     candidates = [
         ent.text for ent in doc.ents
-        if ent.label_ in labels
+        if ent.label_ in label_set
         and len(ent.text) >= min_length
         and not (exclude_pat and re.search(exclude_pat, ent.text, exclude_flags))
     ]
@@ -272,11 +279,7 @@ def _resolve_entity(rule: Dict[str, Any], doc: Any) -> Optional[str]:
     if not candidates:
         return None
 
-    try:
-        result = candidates[match_index]
-    except IndexError:
-        result = candidates[-1]
-
+    result = candidates[match_index] if match_index < len(candidates) else candidates[-1]
     return _apply_normalizers(result, normalize)
 
 
@@ -299,7 +302,7 @@ def _resolve_entity_all(rule: Dict[str, Any], doc: Any) -> List[Optional[str]]:
     labels = rule.get("label")
     if isinstance(labels, str):
         labels = [labels]
-    labels = set(labels or [])
+    label_set = set(labels or [])
 
     min_length    = rule.get("min_length", 1)
     exclude_pat   = rule.get("exclude_pattern", "")
@@ -312,7 +315,7 @@ def _resolve_entity_all(rule: Dict[str, Any], doc: Any) -> List[Optional[str]]:
     seen: set = set()
 
     for ent in doc.ents:
-        if ent.label_ not in labels:
+        if ent.label_ not in label_set:
             continue
         if len(ent.text) < min_length:
             continue
@@ -351,11 +354,7 @@ def _resolve_token_attr(rule: Dict[str, Any], doc: Any) -> Optional[str]:
     if not candidates:
         return None
 
-    try:
-        result = candidates[match_index]
-    except IndexError:
-        result = candidates[-1]
-
+    result = candidates[match_index] if match_index < len(candidates) else candidates[-1]
     return _apply_normalizers(result, normalize)
 
 
@@ -381,8 +380,7 @@ def _resolve_scalar_field(
     text: str,
 ) -> Optional[str]:
     src = rule.get("source_type")
-    with _resolver_lock:
-        fn = _RESOLVERS.get(src)
+    fn = _RESOLVERS.get(src)
     if fn is None:
         logger.warning("Unknown source_type %r — no resolver registered", src)
         return None
@@ -501,7 +499,6 @@ def _resolve_array_node(
     """
     item_schema: SchemaNode  = node.get("items", {})
     split_pat: Optional[str] = node.get("split_pattern")
-    split_flags_rule          = {"flags": node.get("split_flags", [])}
     max_items: Optional[int]  = node.get("max_items")
 
     results: List[ResultNode] = []
@@ -509,7 +506,7 @@ def _resolve_array_node(
     if split_pat:
         # ── MODE A: segment-per-item ────────────────────────────────────
         try:
-            flags    = _build_flags(split_flags_rule)
+            flags    = _build_flags({"flags": node.get("split_flags", [])})
             segments = re.split(split_pat, text, flags=flags)
         except re.error as exc:
             logger.error("Invalid split_pattern %r: %s", split_pat, exc)
@@ -577,6 +574,15 @@ def _safe_resolve_node(
         return None
 
 
+def _doc_for_text(nlp: Any, text: str) -> Any:
+    """Run *text* through spaCy, returning None on failure instead of raising."""
+    try:
+        return next(iter(nlp.pipe([text])))
+    except Exception as exc:
+        logger.error("spaCy pipe failed: %s", exc, exc_info=True)
+        return None
+
+
 # ---------------------------------------------------------------------------
 # Public API
 # ---------------------------------------------------------------------------
@@ -595,17 +601,10 @@ def extract_fields(
 
     Backward-compatible: callers that pass flat scalar rules unchanged still work.
     """
-    nlp     = _get_nlp()
-    cleaned = text
-
-    try:
-        doc = next(iter(nlp.pipe([cleaned])))
-    except Exception as exc:
-        logger.error("spaCy pipe failed: %s", exc, exc_info=True)
-        doc = None
-
+    nlp = _get_nlp()
+    doc = _doc_for_text(nlp, text)
     return {
-        field: _safe_resolve_node(field, node, doc, cleaned)
+        field: _safe_resolve_node(field, node, doc, text)
         for field, node in fields.items()
     }
 
@@ -637,16 +636,9 @@ def extract_schema(
     result = extract_schema(invoice_text, schema)
     # → [{"date": "Jan 2024", "amount": "$1,200"}, ...]
     """
-    nlp     = _get_nlp()
-    cleaned = text
-
-    try:
-        doc = next(iter(nlp.pipe([cleaned])))
-    except Exception as exc:
-        logger.error("spaCy pipe failed: %s", exc, exc_info=True)
-        doc = None
-
-    return _safe_resolve_node("<root>", schema, doc, cleaned)
+    nlp = _get_nlp()
+    doc = _doc_for_text(nlp, text)
+    return _safe_resolve_node("<root>", schema, doc, text)
 
 
 def extract_fields_batch(
@@ -661,7 +653,7 @@ def extract_fields_batch(
     Note: array nodes with split_pattern trigger their own inner pipe call per
     text; the outer pass still processes the top-level texts efficiently.
     """
-    nlp     = _get_nlp()
+    nlp = _get_nlp()
     cleaned = [clean_text(t) for t in texts]
 
     try:
@@ -688,7 +680,7 @@ def extract_schema_batch(
 
     Returns one ResultNode per input text, in the same order.
     """
-    nlp     = _get_nlp()
+    nlp = _get_nlp()
     cleaned = [clean_text(t) for t in texts]
 
     try:
@@ -700,4 +692,4 @@ def extract_schema_batch(
     return [
         _safe_resolve_node(f"<root>[{i}]", schema, doc, text)
         for i, (doc, text) in enumerate(zip(docs, cleaned))
-    ]
+    ]