section-anything 0.1.0__py3-none-any.whl

section_anything/__init__.py

@@ -0,0 +1,90 @@
+"""Section Anything: find the sections of a PDF, with a cost/accuracy dial.
+
+    >>> import section_anything
+    >>> tree = section_anything.section("paper.pdf", mode="free")
+    >>> print(tree.to_markdown())
+
+Pre-release: the API may change before 1.0. Names load lazily, so importing the package is cheap.
+"""
+
+from __future__ import annotations
+
+from importlib import import_module
+from typing import TYPE_CHECKING, Any
+
+__version__ = "0.1.0"
+
+_EXPORTS = {
+    "section": "section_anything.api",
+    "asection": "section_anything.api",
+    "SectionTree": "section_anything.api",
+    "resolve_mode": "section_anything.api",
+    "Settings": "section_anything.config",
+    "Price": "section_anything.config",
+    "DEFAULT_PRICES": "section_anything.config",
+    "extract_pdf": "section_anything.isolate",
+    "check_pdf": "section_anything.isolate",
+    "Extraction": "section_anything.extract",
+    "SectionAnythingError": "section_anything.errors",
+    "ConfigError": "section_anything.errors",
+    "UnpricedModelError": "section_anything.errors",
+    "CostCapExceeded": "section_anything.errors",
+    "LlmCallFailed": "section_anything.errors",
+    "PdfError": "section_anything.errors",
+    "PdfTooLarge": "section_anything.errors",
+    "ExtractionTimeout": "section_anything.errors",
+    "ExtractionFailed": "section_anything.errors",
+}
+
+__all__ = [
+    "DEFAULT_PRICES",
+    "ConfigError",
+    "CostCapExceeded",
+    "Extraction",
+    "ExtractionFailed",
+    "ExtractionTimeout",
+    "LlmCallFailed",
+    "PdfError",
+    "PdfTooLarge",
+    "Price",
+    "SectionAnythingError",
+    "SectionTree",
+    "Settings",
+    "UnpricedModelError",
+    "__version__",
+    "asection",
+    "check_pdf",
+    "extract_pdf",
+    "resolve_mode",
+    "section",
+]
+
+if TYPE_CHECKING:
+    from section_anything.api import SectionTree, asection, resolve_mode, section
+    from section_anything.config import DEFAULT_PRICES, Price, Settings
+    from section_anything.errors import (
+        ConfigError,
+        CostCapExceeded,
+        ExtractionFailed,
+        ExtractionTimeout,
+        LlmCallFailed,
+        PdfError,
+        PdfTooLarge,
+        SectionAnythingError,
+        UnpricedModelError,
+    )
+    from section_anything.extract import Extraction
+    from section_anything.isolate import check_pdf, extract_pdf
+
+
+def __getattr__(name: str) -> Any:
+    module = _EXPORTS.get(name)
+    if module is None:
+        raise AttributeError(f"module 'section_anything' has no attribute {name!r}")
+    value = getattr(import_module(module), name)
+    globals()[name] = value
+    return value
+
+
+def __dir__() -> list[str]:
+    return sorted(__all__)

section_anything/__main__.py

@@ -0,0 +1,5 @@
+import sys
+
+from section_anything.cli import main
+
+sys.exit(main())

section_anything/_worker.py

@@ -0,0 +1,66 @@
+"""Extraction worker: ``python -m section_anything._worker PDF OUT MAX_PAGES MAX_CHARS MEMORY_MB``.
+
+Started by `section_anything.isolate`. Writes the extraction to OUT as JSON and reports progress
+and the outcome as JSON lines on its original stdout. Error messages carry only text written by
+this package or an exception's type name, never document text.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import sys
+from pathlib import Path
+
+
+def _limit_memory(mb: int) -> None:
+    """Cap the address space on Linux. macOS does not enforce RLIMIT_AS, so it is skipped there."""
+    if mb <= 0 or not sys.platform.startswith("linux"):
+        return
+    try:
+        import resource
+
+        limit = mb * 1024 * 1024
+        resource.setrlimit(resource.RLIMIT_AS, (limit, limit))
+    except (ImportError, ValueError, OSError):
+        pass
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = sys.argv[1:] if argv is None else argv
+    pdf, out, max_pages, max_chars, memory_mb = args[0], args[1], int(args[2]), int(args[3]), int(args[4])
+    # Keep the protocol on the original stdout; anything else that prints goes to stderr.
+    proto = os.fdopen(os.dup(1), "w", buffering=1, encoding="utf-8")
+    os.dup2(2, 1)
+    sys.stdout = sys.stderr
+
+    def emit(msg: dict) -> None:
+        proto.write(json.dumps(msg) + "\n")
+        proto.flush()
+
+    _limit_memory(memory_mb)
+    from section_anything.errors import PdfError, PdfTooLarge
+    from section_anything.extract import extract
+
+    try:
+        ext = extract(Path(pdf), lambda n, total: emit({"page": n, "total": total}),
+                      max_pages=max_pages, max_chars=max_chars)
+        ext.save(Path(out))
+    except PdfTooLarge as exc:
+        emit({"error": "too_large", "message": str(exc)})
+        return 2
+    except PdfError as exc:
+        emit({"error": "pdf", "message": str(exc)})
+        return 2
+    except MemoryError:
+        emit({"error": "memory", "message": f"out of memory (extract_memory_mb={memory_mb})"})
+        return 3
+    except Exception as exc:  # noqa: BLE001 — reported by type only: messages may quote the PDF
+        emit({"error": "internal", "message": type(exc).__name__})
+        return 1
+    emit({"done": True})
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

section_anything/api.py

@@ -0,0 +1,267 @@
+"""Public API: `section(pdf, mode)` -> `SectionTree`.
+
+Modes, cheapest first:
+
+- ``free``: no model. PDF bookmarks or the printed table of contents as the outline, then
+  typography and numbering rules. Nothing leaves the machine.
+- ``quick``: the rules propose candidate heading lines; a small model keeps the real ones and
+  sets their levels. Sends only the candidate lines (plus a short peek at the next line).
+- ``standard``: a model reads every line, with font hints, and marks where sections start and
+  end. Sends the document text.
+- ``local:<model>``: quick, answered by a local Ollama model. Free; nothing leaves the machine.
+
+Costs: a cloud model with no known price is refused before anything is read. Before a paid mode
+runs, its estimated cost is checked against the per-document cap (`max_usd`); a run over the cap
+is downgraded (standard -> quick -> free) or refused (``over_budget="fail"``). While it runs, the
+cap is re-checked before every call and each call's output is limited to what the remaining budget
+can pay for.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import concurrent.futures
+import json
+import os
+import time
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+from section_anything.config import LOCAL_PREFIX, Price, Settings
+from section_anything.errors import ConfigError, CostCapExceeded, PdfError, UnpricedModelError
+from section_anything.export import heading_text, markdown_hashes, section_bodies
+from section_anything.extract import Extraction
+from section_anything.isolate import extract_pdf
+from section_anything.llm import Usage
+
+MODES = ("free", "quick", "standard", "local")
+STRATEGY = {"free": "rules", "quick": "llm_review", "standard": "llm_text", "local": "llm_review"}
+DOWNGRADE = {"standard": "quick", "quick": "free", "local": None, "free": None}
+SECTION_KEYS = ("number", "title", "level", "page_start", "page_end", "start", "end", "flags")
+
+Progress = Callable[[float, str], None]
+
+
+def resolve_mode(mode: str, settings: Settings, model: str | None = None) -> tuple[str, str | None]:
+    """'free' | 'quick' | 'standard' | 'local:<ollama model>' -> (mode, model id or None).
+
+    `model` replaces the configured model of a quick or standard run, or names the local model."""
+    m = (mode or "free").strip()
+    if m == "local" or m.startswith("local:"):
+        tag = m.partition(":")[2].strip() or (model or "").strip() or settings.local_model.strip()
+        if not tag:
+            raise ConfigError("local mode needs a model: mode='local:<ollama model>' or SECTIONANY_LOCAL_MODEL")
+        return "local", LOCAL_PREFIX + tag
+    if m == "free":
+        return "free", None
+    if m == "quick":
+        return "quick", model or settings.quick_model
+    if m == "standard":
+        return "standard", model or settings.standard_model
+    raise ConfigError(f"unknown mode {m!r}: use free, quick, standard or local:<ollama model>")
+
+
+def _require_price(mode: str, model: str | None, settings: Settings) -> Price | None:
+    if model is None:
+        return None
+    price = settings.price_for(model)
+    if price is None:
+        raise UnpricedModelError(
+            f"no price for model {model!r}: pass prices={{{model!r}: {{'input': ..., 'output': ...}}}} "
+            "(USD per 1M tokens) or set SECTIONANY_PRICES, so the per-document cost cap can be enforced"
+        )
+    if mode in ("quick", "standard") and settings.api_key is None:
+        raise ConfigError(f"{mode} mode needs an API key: set SECTIONANY_API_KEY or pass api_key=")
+    return price
+
+
+@dataclass
+class SectionTree:
+    """The sections of one document, in reading order.
+
+    `sections` holds one dict per section with number, title, level (1 = top), page_start,
+    page_end, start and end (character offsets into `text`) and flags (for example
+    ``doc_title``, ``front_matter``, ``start_marker_fuzzy``)."""
+
+    sections: list[dict]
+    mode: str
+    model: str | None
+    pages: int
+    usage: dict
+    seconds: float
+    downgraded: list[str] = field(default_factory=list)
+    extraction: Extraction | None = field(default=None, repr=False)
+    resolved: list[dict] = field(default_factory=list, repr=False)  # full pipeline output
+
+    @property
+    def text(self) -> str:
+        """The document text the character offsets refer to."""
+        return self.extraction.text if self.extraction is not None else ""
+
+    @property
+    def cost_usd(self) -> float:
+        return float(self.usage.get("cost_usd") or 0.0)
+
+    def summary(self) -> dict:
+        """Mode, model, size, calls, tokens and cost of the run."""
+        return {
+            "mode": self.mode,
+            "model": self.model,
+            "pages": self.pages,
+            "sections": len(self.sections),
+            "flagged": sum(bool(s["flags"]) for s in self.sections),
+            "calls": self.usage.get("calls", 0),
+            "input_tokens": self.usage.get("input_tokens", 0),
+            "output_tokens": self.usage.get("output_tokens", 0),
+            "cost_usd": self.cost_usd,
+            "seconds": self.seconds,
+            "downgraded": list(self.downgraded),
+        }
+
+    def to_markdown(self, body: bool = False) -> str:
+        """Headings as Markdown: the document title (flag ``doc_title``) as '#', other headings
+        as '#' x (level + 1), capped at six. With `body`, each heading is followed by its section's
+        text reflowed into paragraphs; otherwise front matter is left out."""
+        out: list[str] = []
+        if body and self.extraction is not None and self.resolved:
+            from section_anything.structure import furniture_of
+
+            bodies = section_bodies(self.extraction, self.resolved, furniture_of(self.extraction))
+            for sec, paras in zip(self.resolved, bodies, strict=True):
+                if "front_matter" not in sec["flags"]:
+                    out.append(f"{markdown_hashes(sec)} {heading_text(sec)}")
+                out.extend(paras)
+            return "\n\n".join(out) + "\n"
+        for sec in self.sections:
+            if "front_matter" in sec["flags"]:
+                continue
+            out.append(f"{markdown_hashes(sec)} {heading_text(sec)}")
+        return "\n".join(out) + ("\n" if out else "")
+
+    def to_dict(self) -> dict:
+        from section_anything import __version__
+
+        return {"section_anything": __version__, **self.summary(), "usage": self.usage,
+                "sections": self.sections}
+
+    def to_json(self, indent: int | None = 2) -> str:
+        return json.dumps(self.to_dict(), ensure_ascii=False, indent=indent)
+
+
+def _public(sec: dict) -> dict:
+    out = {k: (list(sec[k]) if k == "flags" else sec[k]) for k in SECTION_KEYS}
+    if sec.get("title_from_text") is False and "front_matter" not in out["flags"]:
+        out["flags"].append("generated_title")
+    return out
+
+
+def _settings_for(settings: Settings | None, **overrides: Any) -> Settings:
+    base = settings if settings is not None else Settings()
+    prices = overrides.pop("prices", None)
+    if prices is not None:
+        overrides["prices"] = {m: Price.model_validate(p) for m, p in dict(prices).items()}
+    return base.with_overrides(**overrides)
+
+
+async def asection(
+    pdf_path: str | os.PathLike[str],
+    mode: str = "free",
+    *,
+    settings: Settings | None = None,
+    api_key: str | None = None,
+    base_url: str | None = None,
+    model: str | None = None,
+    max_usd: float | None = None,
+    prices: Mapping[str, Mapping[str, float] | Price] | None = None,
+    over_budget: str | None = None,
+    max_pages: int | None = None,
+    extract_timeout: float | None = None,
+    progress: Progress | None = None,
+) -> SectionTree:
+    """Async form of `section`."""
+    from section_anything import pipeline
+
+    s = _settings_for(settings, api_key=api_key, base_url=base_url, max_usd_per_doc=max_usd, prices=prices,
+                      over_budget=over_budget, max_pages=max_pages, extract_timeout_seconds=extract_timeout)
+    mode, model = resolve_mode(mode, s, model)
+    _require_price(mode, model, s)  # before the PDF is read: a refused run costs nothing
+    report = progress or (lambda _f, _m: None)
+
+    t0 = time.monotonic()
+    ext = await asyncio.to_thread(
+        extract_pdf, Path(pdf_path), settings=s,
+        on_page=lambda n, total: report(0.5 * n / total, f"Reading page {n} of {total}"),
+    )
+    if not ext.text.strip():
+        raise PdfError("no text layer: the PDF needs OCR first")
+
+    cap = s.max_usd_per_doc
+    usage = Usage()  # one tally per document: a downgrade continues the same budget
+    downgraded: list[str] = []
+    while True:
+        strategy = pipeline.BY_ID[STRATEGY[mode]]
+        price = s.price_for(model) if model else None
+        if model and price is not None and strategy.uses_llm:
+            est = strategy.estimate(ext)
+            usd = price.cost(est.input_tokens, est.output_tokens)
+            if usage.cost_usd + usd > cap:
+                note = f"{mode}: estimated ${usd:.2f}, over the ${cap:.2f} per-document cap"
+                if s.over_budget == "fail" or DOWNGRADE[mode] is None:
+                    raise CostCapExceeded(f"{note} (max_usd, SECTIONANY_MAX_USD_PER_DOC)")
+                downgraded.append(note)
+                mode, model = _next_mode(mode, s, downgraded)
+                continue
+
+        def stage(f: float, msg: str) -> None:
+            report(0.5 + 0.5 * f, msg)
+
+        try:
+            _raw, resolved, _metrics = await pipeline.run(
+                ext, strategy.id, model or "", cap, stage, usage=usage, settings=s)
+        except CostCapExceeded as exc:
+            if s.over_budget == "fail" or DOWNGRADE[mode] is None:
+                raise
+            downgraded.append(f"{mode}: {exc}")
+            mode, model = _next_mode(mode, s, downgraded)
+            continue
+        return SectionTree(
+            sections=[_public(x) for x in resolved], mode=mode, model=model, pages=len(ext.pages),
+            usage=usage.summary(), seconds=round(time.monotonic() - t0, 2), downgraded=downgraded,
+            extraction=ext, resolved=resolved,
+        )
+
+
+def _next_mode(mode: str, s: Settings, notes: list[str]) -> tuple[str, str | None]:
+    nxt = DOWNGRADE[mode]
+    while nxt is not None:
+        nxt_mode, nxt_model = resolve_mode(nxt, s)
+        if nxt_model is None or s.price_for(nxt_model) is not None:
+            return nxt_mode, nxt_model
+        notes.append(f"{nxt}: skipped, no price for model {nxt_model!r}")
+        nxt = DOWNGRADE[nxt]
+    raise ConfigError("no mode left to downgrade to")  # unreachable: free is always available
+
+
+def section(
+    pdf_path: str | os.PathLike[str],
+    mode: str = "free",
+    **opts: Any,
+) -> SectionTree:
+    """Find the sections of a PDF.
+
+    `mode` is "free", "quick", "standard" or "local:<ollama model>". Options (each overrides the
+    matching ``SECTIONANY_*`` environment variable): `api_key`, `base_url`, `model`, `max_usd`
+    (per-document cap), `prices` ({model: {"input", "output"}} USD per 1M tokens, merged with the
+    built-in list), `over_budget` ("downgrade" or "fail"), `max_pages`, `extract_timeout`
+    (seconds), `progress` (callback(fraction, message)) and `settings` (a `Settings` object).
+    """
+    coro = asection(pdf_path, mode, **opts)
+    try:
+        asyncio.get_running_loop()
+    except RuntimeError:
+        return asyncio.run(coro)
+    # Called from inside an event loop (a notebook, an async app): run on a separate thread.
+    with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
+        return pool.submit(asyncio.run, coro).result()

section_anything/bridge.py

@@ -0,0 +1,70 @@
+"""Browser bridge: answer a run's model calls with a model on someone else's machine.
+
+A run whose model is "browser/<name>" never calls a provider. Each model call is queued here; a
+client (for example a web page that can reach the user's own Ollama or LM Studio) long-polls
+`next()`, sends the prompt to its model and posts the reply back with `resolve()`. The reply is
+validated against the call's schema like any other model output.
+
+The queue is in memory and per process: the host application must route a run's polling
+requests to the process that runs it.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import uuid
+from typing import Any
+
+from section_anything.config import BROWSER_PREFIX
+from section_anything.errors import SectionAnythingError
+
+__all__ = ["BROWSER_PREFIX", "BridgeError", "BridgeTimeout", "BrowserBridge"]
+
+
+class BridgeError(SectionAnythingError, RuntimeError):
+    """The remote model reported an error."""
+
+
+class BridgeTimeout(BridgeError, TimeoutError):
+    """No reply arrived in time."""
+
+
+class BrowserBridge:
+    def __init__(self, wait_seconds: float = 300.0) -> None:
+        self.wait_seconds = wait_seconds
+        self._queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
+        self._pending: dict[str, asyncio.Future[str]] = {}
+
+    async def call(self, payload: dict[str, Any], wait_seconds: float | None = None) -> str:
+        """Queue one model call and wait for the reply text."""
+        rid = uuid.uuid4().hex
+        fut: asyncio.Future[str] = asyncio.get_running_loop().create_future()
+        self._pending[rid] = fut
+        await self._queue.put({"id": rid, **payload})
+        try:
+            async with asyncio.timeout(self.wait_seconds if wait_seconds is None else wait_seconds):
+                return await fut
+        except TimeoutError as exc:
+            raise BridgeTimeout(
+                "no reply from your local model: keep the page open and check that the model is running"
+            ) from exc
+        finally:
+            self._pending.pop(rid, None)
+
+    async def next(self, wait: float) -> dict[str, Any] | None:
+        """The next queued call, or None after `wait` seconds."""
+        try:
+            async with asyncio.timeout(wait):
+                return await self._queue.get()
+        except TimeoutError:
+            return None
+
+    def resolve(self, rid: str, content: str | None, error: str | None) -> bool:
+        fut = self._pending.get(rid)
+        if fut is None or fut.done():
+            return False
+        if error is not None:
+            fut.set_exception(BridgeError(f"local model error: {error[:300]}"))
+        else:
+            fut.set_result(content or "")
+        return True

section_anything/cli.py

@@ -0,0 +1,77 @@
+"""Command line: ``section-anything FILE.pdf --mode quick --format md``.
+
+The result goes to stdout (or ``--output``); a one-line cost summary goes to stderr.
+Exit codes: 0 done, 2 bad arguments or configuration, 3 refused by a cost cap, 4 the PDF could not
+be read, 1 anything else.
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+from section_anything.errors import ConfigError, CostCapExceeded, ExtractionFailed, ExtractionTimeout, PdfError
+
+
+def build_parser() -> argparse.ArgumentParser:
+    from section_anything import __version__
+
+    p = argparse.ArgumentParser(
+        prog="section-anything",
+        description="Find the sections (headings, levels, page and character spans) of a PDF.",
+    )
+    p.add_argument("pdf", type=Path, help="PDF file")
+    p.add_argument("--mode", default="free",
+                   help="free (default), quick, standard, or local:<ollama model>")
+    p.add_argument("--format", choices=("md", "json"), default="md", help="output format (default md)")
+    p.add_argument("--body", action="store_true", help="md: include each section's text, not only headings")
+    p.add_argument("--max-usd", type=float, default=None, help="per-document cost cap in USD (default 0.50)")
+    p.add_argument("--over-budget", choices=("downgrade", "fail"), default=None,
+                   help="over the cap: downgrade standard -> quick -> free (default), or fail")
+    p.add_argument("--model", default=None, help="model for quick/standard, overriding the configured one")
+    p.add_argument("--base-url", default=None, help="OpenAI-compatible endpoint (default OpenAI)")
+    p.add_argument("--max-pages", type=int, default=None, help="refuse PDFs with more pages (default 500)")
+    p.add_argument("--timeout", type=float, default=None, help="extraction time limit in seconds (default 120)")
+    p.add_argument("-o", "--output", type=Path, default=None, help="write the result to this file")
+    p.add_argument("-q", "--quiet", action="store_true", help="no summary on stderr")
+    p.add_argument("--version", action="version", version=f"%(prog)s {__version__}")
+    return p
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = build_parser().parse_args(argv)
+    from section_anything.api import section
+
+    try:
+        tree = section(
+            args.pdf, args.mode, model=args.model, base_url=args.base_url, max_usd=args.max_usd,
+            over_budget=args.over_budget, max_pages=args.max_pages, extract_timeout=args.timeout,
+        )
+    except (PdfError, ExtractionTimeout, ExtractionFailed) as exc:
+        print(f"section-anything: {exc}", file=sys.stderr)
+        return 4
+    except CostCapExceeded as exc:
+        print(f"section-anything: {exc}", file=sys.stderr)
+        return 3
+    except (ConfigError, ValueError) as exc:  # includes invalid settings values
+        print(f"section-anything: {exc}", file=sys.stderr)
+        return 2
+
+    out = tree.to_json() + "\n" if args.format == "json" else tree.to_markdown(body=args.body)
+    if args.output:
+        args.output.write_text(out, encoding="utf-8")
+    else:
+        sys.stdout.write(out)
+    if not args.quiet:
+        s = tree.summary()
+        cost = "unpriced" if tree.usage.get("cost_usd") is None else f"${s['cost_usd']:.4f}"
+        note = f"; downgraded: {' | '.join(s['downgraded'])}" if s["downgraded"] else ""
+        print(f"{s['sections']} sections, {s['pages']} pages, mode {s['mode']}"
+              f"{' (' + s['model'] + ')' if s['model'] else ''}, {s['calls']} model calls, {cost}, "
+              f"{s['seconds']:.1f} s{note}", file=sys.stderr)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())