src/retrieval/corpus.py

"""Load the regulation corpus.

Each document is one article / section, tagged with its source, jurisdiction and
year so a citation can name exactly where a claim comes from. The primary format is
JSONL (one document per line) for compactness; loose ``.md`` files with simple
front-matter are also supported so the corpus is easy to extend by hand.
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import List, Optional, Union


@dataclass
class Document:
    id: str
    title: str
    source: str
    jurisdiction: str  # "MX" | "US"
    year: Optional[int]
    text: str


def _from_jsonl(path: Path) -> List[Document]:
    docs = []
    for line in path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line:
            continue
        d = json.loads(line)
        docs.append(Document(
            id=d["id"], title=d["title"], source=d["source"],
            jurisdiction=d.get("jurisdiction", ""), year=d.get("year"),
            text=d["text"],
        ))
    return docs


def _parse_front_matter(raw: str) -> Document:
    """Parse a `--- key: value ... --- body` markdown file."""
    meta, body = {}, raw
    if raw.startswith("---"):
        _, fm, body = raw.split("---", 2)
        for line in fm.strip().splitlines():
            if ":" in line:
                k, v = line.split(":", 1)
                meta[k.strip()] = v.strip()
    year = meta.get("year")
    return Document(
        id=meta.get("id", ""), title=meta.get("title", ""),
        source=meta.get("source", ""), jurisdiction=meta.get("jurisdiction", ""),
        year=int(year) if year and year.isdigit() else None,
        text=body.strip(),
    )


def load_corpus(path: Union[str, Path]) -> List[Document]:
    """Load every document under a directory (``*.jsonl`` + ``*.md``)."""
    p = Path(path)
    docs: List[Document] = []
    if p.is_file():
        return _from_jsonl(p) if p.suffix == ".jsonl" else [_parse_front_matter(p.read_text("utf-8"))]
    for f in sorted(p.glob("*.jsonl")):
        docs.extend(_from_jsonl(f))
    for f in sorted(p.glob("*.md")):
        if f.name.upper() == "NOTES.MD":
            continue
        doc = _parse_front_matter(f.read_text("utf-8"))
        if doc.text:
            docs.append(doc)
    return docs