rc-docparser 0.2.0__py3-none-any.whl

docparser/orchestrator.py

@@ -0,0 +1,304 @@
+"""Orchestrator: dispatch by file extension and walk a directory.
+
+The orchestrator is what most callers use: give it a directory and a
+``WorkspaceLayout`` and it parses every supported file underneath, writing
+``document.md`` + ``document.json`` per file plus a top-level ``corpus.json``
+index.
+"""
+from __future__ import annotations
+
+import time
+import traceback
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any
+
+from .common import (
+    WorkspaceLayout,
+    slugify,
+    truncate,
+    utc_now_iso,
+    write_json,
+    write_text,
+)
+from .csvtab import parse_csv
+from .docx import parse_docx
+from .text import parse_text
+from .xlsx import parse_xlsx
+
+SUPPORTED_EXTENSIONS: dict[str, str] = {
+    ".docx": "docx",
+    ".xlsx": "xlsx",
+    ".pdf": "pdf",
+    ".html": "html",
+    ".htm": "html",
+    ".pptx": "pptx",
+    ".epub": "epub",
+    ".txt": "text",
+    ".md": "text",
+    ".markdown": "text",
+    ".csv": "csv",
+    ".tsv": "csv",
+}
+
+
+def _get_parser(kind: str) -> Callable[..., dict[str, Any]]:
+    if kind == "docx":
+        return parse_docx
+    if kind == "xlsx":
+        return parse_xlsx
+    if kind == "text":
+        return parse_text
+    if kind == "csv":
+        return parse_csv
+    if kind == "pdf":
+        from .pdf import parse_pdf
+
+        return parse_pdf
+    if kind == "html":
+        from .html import parse_html
+
+        return parse_html
+    if kind == "pptx":
+        from .pptx import parse_pptx
+
+        return parse_pptx
+    if kind == "epub":
+        from .epub import parse_epub
+
+        return parse_epub
+    raise ValueError(f"unsupported kind: {kind}")
+
+
+def parse_path(
+    source: Path | str,
+    layout: WorkspaceLayout | None = None,
+    *,
+    captioner: Callable[..., dict[str, Any]] | None = None,
+    write_outputs: bool = True,
+    **kwargs: Any,
+) -> dict[str, Any]:
+    """Dispatch ``source`` to the right parser by extension.
+
+    Extra keyword arguments (e.g. ``backend``, ``ocr``, ``extract_tables``)
+    are forwarded only to the PDF parser; other parsers ignore them.
+    """
+    source = Path(source)
+    suffix = source.suffix.lower()
+    kind = SUPPORTED_EXTENSIONS.get(suffix)
+    if kind is None:
+        raise ValueError(f"Unsupported extension: {suffix}")
+    parser = _get_parser(kind)
+    if kind == "pdf" and kwargs:
+        return parser(source, layout, captioner=captioner, write_outputs=write_outputs, **kwargs)
+    return parser(source, layout, captioner=captioner, write_outputs=write_outputs)
+
+
+def make_captioner(
+    *,
+    enabled: bool,
+    layout: WorkspaceLayout,
+    max_calls: int | None = None,
+    log_path: Path | None = None,
+    provider: str | None = None,
+    model: str | None = None,
+):
+    """Build a captioner closure that wraps the image captioner.
+
+    ``provider`` selects the captioning backend: an API provider preset
+    (``openrouter`` / ``openai`` / ``gemini`` / ``local``) handled by
+    :func:`docparser.image.caption_image`, or ``transformers`` for the fully
+    local :func:`docparser.localvlm.caption_image_local`.
+
+    Enforces a per-run cap on actual VLM calls and (optionally) appends a TSV
+    log of every call.
+    """
+    if not enabled:
+        return None
+
+    use_local = (provider or "").lower() in {"transformers", "localvlm"}
+    _backend: Callable[..., Any]
+    if use_local:
+        from .localvlm import caption_image_local
+
+        _backend = caption_image_local  # lazy: needs [localvlm]
+    else:
+        from .image import caption_image
+
+        _backend = caption_image  # lazy import: needs [vlm] extra
+
+    state = {"calls": 0}
+
+    def _captioner(*, image_bytes, mime, doc_name, nearby_caption, context):
+        if max_calls is not None and state["calls"] >= max_calls:
+            return {
+                "error": f"max_calls ({max_calls}) reached for this run",
+                "caption": "",
+                "description": "",
+                "visible_text": "",
+                "tags": [],
+                "image_kind": "other",
+                "domain_relevance": "",
+                "model": "(skipped)",
+                "cached": False,
+            }
+        t0 = time.time()
+        kwargs = dict(
+            image_bytes=image_bytes,
+            mime=mime,
+            doc_name=doc_name,
+            nearby_caption=nearby_caption,
+            context=context,
+            layout=layout,
+        )
+        if not use_local:
+            kwargs["provider"] = provider
+            kwargs["model"] = model
+        result = _backend(**kwargs)
+        elapsed = time.time() - t0
+        if not result.cached:
+            state["calls"] += 1
+        if log_path is not None:
+            log_path.parent.mkdir(parents=True, exist_ok=True)
+            with log_path.open("a", encoding="utf-8") as fh:
+                fh.write(
+                    f"{utc_now_iso()}\t{doc_name}\tcached={result.cached}\t"
+                    f"elapsed={elapsed:.2f}s\tmodel={result.model}\t"
+                    f"err={result.error or ''}\tcaption={truncate(result.caption,120)}\n"
+                )
+        return result.to_dict()
+
+    return _captioner
+
+
+def build_corpus_index(
+    parsed_runs: list[dict[str, Any]], layout: WorkspaceLayout
+) -> tuple[str, dict[str, Any]]:
+    md = ["# Corpus index", ""]
+    md.append(f"_Built {utc_now_iso()} from {len(parsed_runs)} source files._")
+    md.append("")
+    md.append("| File | Kind | Sections | Images | Captioned | sha1 |")
+    md.append("| --- | --- | --- | --- | --- | --- |")
+    total_images = 0
+    total_captioned = 0
+    for r in parsed_runs:
+        src = r["source"]
+        stats = r.get("stats", {})
+        kind = src["kind"]
+        if kind == "xlsx":
+            sect = f"{stats.get('n_sheets', 0)} sheets · {stats.get('n_nonempty_cells_total', 0)} cells"
+            n_imgs = stats.get("n_images_total", 0)
+            n_capt = sum(
+                1
+                for s in r.get("sheets", [])
+                for im in s.get("images", [])
+                if (im.get("semantic") or {}).get("caption")
+            )
+        elif kind == "html":
+            sect = (
+                f"{stats.get('n_headings', 0)} hdr · "
+                f"{stats.get('n_paragraphs', 0)} para · "
+                f"{stats.get('n_tables', 0)} tbl"
+            )
+            n_imgs = stats.get("n_images", 0)
+            n_capt = 0
+        else:
+            sect = (
+                f"{stats.get('n_headings', 0)} hdr · "
+                f"{stats.get('n_paragraphs', 0)} para · "
+                f"{stats.get('n_tables', 0)} tbl"
+            )
+            n_imgs = stats.get("n_images", 0)
+            n_capt = stats.get("n_captioned_images", 0)
+        total_images += n_imgs
+        total_captioned += n_capt
+        slug = slugify(Path(src["filename"]).stem)
+        md.append(
+            f"| `{src['filename']}` | {kind} | {sect} | {n_imgs} | {n_capt} | `{src['sha1'][:10]}` |"
+        )
+        _ = slug
+    md.append("")
+    md.append(f"**Totals:** {total_images} images, {total_captioned} VLM-captioned.")
+    md.append("")
+    corpus_json = {
+        "built_at": utc_now_iso(),
+        "n_sources": len(parsed_runs),
+        "totals": {"images": total_images, "captioned_images": total_captioned},
+        "sources": [
+            {
+                "filename": r["source"]["filename"],
+                "kind": r["source"]["kind"],
+                "sha1": r["source"]["sha1"],
+                "stats": r.get("stats", {}),
+                "parsed_dir": str(layout.parsed_dir_for(Path(r["source"]["filename"]))),
+            }
+            for r in parsed_runs
+        ],
+    }
+    return "\n".join(md), corpus_json
+
+
+def run_all(
+    layout: WorkspaceLayout,
+    *,
+    use_vlm: bool = True,
+    only: str | None = None,
+    max_images: int | None = None,
+    continue_on_error: bool = False,
+    log_path: Path | None = None,
+    write_corpus: bool = True,
+    vlm_provider: str | None = None,
+    vlm_model: str | None = None,
+    pdf_backend: str = "builtin",
+    ocr: str = "off",
+    extract_tables: bool = False,
+) -> list[dict[str, Any]]:
+    """Walk ``layout.raw_dir`` and parse every supported file.
+
+    Returns the list of parsed payloads. If ``write_corpus`` is True (default),
+    a ``corpus.json`` is written under ``layout.parsed_dir`` and a top-level
+    ``CORPUS.md`` is returned via the side effect of writing it next to the
+    parsed dir's parent.
+    """
+    layout.ensure()
+    files = sorted(p for p in layout.raw_dir.iterdir() if p.is_file() or p.is_symlink())
+    files = [p for p in files if p.suffix.lower() in SUPPORTED_EXTENSIONS]
+    if only:
+        files = [p for p in files if only.lower() in p.name.lower()]
+    if not files:
+        return []
+
+    captioner = make_captioner(
+        enabled=use_vlm,
+        layout=layout,
+        max_calls=max_images,
+        log_path=log_path or (layout.cache_dir / "vlm_calls.log"),
+        provider=vlm_provider,
+        model=vlm_model,
+    )
+
+    pdf_kwargs: dict[str, Any] = {
+        "backend": pdf_backend,
+        "ocr": ocr,
+        "extract_tables": extract_tables,
+    }
+
+    parsed_runs: list[dict[str, Any]] = []
+    for p in files:
+        try:
+            t0 = time.time()
+            payload = parse_path(p, layout, captioner=captioner, **pdf_kwargs)
+            print(f"[docparser] {p.name}  ok ({time.time() - t0:.1f}s)  stats={payload.get('stats', {})}")
+            parsed_runs.append(payload)
+        except Exception as exc:
+            print(f"[docparser] {p.name}  FAILED: {exc}")
+            traceback.print_exc()
+            if not continue_on_error:
+                raise
+
+    if write_corpus:
+        md, corpus = build_corpus_index(parsed_runs, layout)
+        write_text(layout.parsed_dir.parent / "CORPUS.md", md)
+        write_json(layout.parsed_dir / "corpus.json", corpus)
+
+    return parsed_runs