chroma/loader.py

"""Load agenda documents from raw PDFs.

A "document" here is a plain dict::

    {
        "doc_id":   "packet:0",               # stable, unique id
        "text":     "<extracted agenda text>",
        "metadata": {file_name, source, ...},  # all scalar (str/int/float/bool)
    }

The only source is :func:`documents_from_pdfs` -- generic: accept local PDF paths or
raw bytes. (The CivicClerk OData loaders were removed when the app pivoted to
user-uploaded agenda packets; uploaded packets flow through ``webapp.backend`` which
builds documents from the cached packet pages directly.)
"""

from __future__ import annotations

import os
from typing import Iterable, Union

from pdf_utils import extract_pdf_text

PdfSource = Union[str, bytes, os.PathLike]


# --------------------------------------------------------------------------- #
# Metadata helpers
# --------------------------------------------------------------------------- #
def _scalar_metadata(meta: dict) -> dict:
    """Coerce a metadata dict to Chroma-safe scalars (str/int/float/bool), no None."""
    out: dict = {}
    for k, v in meta.items():
        if v is None:
            out[k] = ""
        elif isinstance(v, (str, int, float, bool)):
            out[k] = v
        else:
            out[k] = str(v)
    return out


# --------------------------------------------------------------------------- #
# Raw PDFs (paths or bytes)
# --------------------------------------------------------------------------- #
def load_pdf_text(source: PdfSource) -> str:
    """Extract text from a PDF given a filesystem path or raw bytes."""
    if isinstance(source, bytes):
        data = source
    else:
        with open(os.fspath(source), "rb") as fh:
            data = fh.read()
    return extract_pdf_text(data)


def documents_from_pdfs(
    sources: Iterable[PdfSource],
    *,
    metadata: dict | None = None,
) -> list[dict]:
    """Build documents from local PDF paths and/or raw PDF bytes.

    ``metadata`` (optional) is merged into every produced document's metadata.
    For path inputs the ``doc_id`` is the file's basename; for bytes it's
    ``pdf:<index>``.
    """
    base = dict(metadata or {})
    docs: list[dict] = []
    for i, src in enumerate(sources):
        text = load_pdf_text(src)
        if not text.strip():
            continue
        if isinstance(src, bytes):
            doc_id = f"pdf:{i}"
            name = f"pdf-{i}.pdf"
        else:
            path = os.fspath(src)
            name = os.path.basename(path)
            doc_id = name
        meta = {**base, "source": "pdf", "file_name": name}
        docs.append(
            {
                "doc_id": doc_id,
                "text": text,
                "metadata": _scalar_metadata(meta),
            }
        )
    return docs