stack/src/llm/links.py

"""Evidence deep links: one (url, label) per retrieved chunk (pure).

- comment → regulations.gov comment page, or the attachment on the
  downloads CDN with ``#page=N`` for PDFs (page located at index time,
  see ``llm.pages``).
- rule    → the FR paragraph anchor ``#p-N`` plus a ``:~:text=``
  scroll-to-text fragment so the browser highlights the passage;
  degrades to ``#page-N``, then the document URL.
- corpus  → the item's own URL (``#page=N`` when it is a PDF).

Labels double as the citation tokens the model is told to emit
(``[91 FR 43949 ¶4]``, ``[CMS-2026-2377-3438 p.4]``).
"""

from __future__ import annotations

from bib.frlink import text_fragment

_COMMENT = "https://www.regulations.gov/comment/{cid}"
_DOWNLOAD = "https://downloads.regulations.gov/{cid}/{name}"
_TITLE_MAX = 60
_SNIPPET_CHARS = 500


def _short(title: str) -> str:
    title = " ".join(title.split())
    if len(title) <= _TITLE_MAX:
        return title
    return title[:_TITLE_MAX].rsplit(" ", 1)[0] + "…"


def _comment(md: dict[str, str]) -> tuple[str, str]:
    cid = md.get("comment_id", "")
    if not cid:
        return "", md.get("item_key", "")
    name, page = md.get("attachment", ""), md.get("page", "")
    if name:
        url = _DOWNLOAD.format(cid=cid, name=name)
        if page and name.lower().endswith(".pdf"):
            return f"{url}#page={page}", f"{cid} p.{page}"
        return url, cid
    return _COMMENT.format(cid=cid), cid


def _rule(md: dict[str, str], snippet: str) -> tuple[str, str]:
    base = md.get("html_url", "") or md.get("url", "")
    page, p_id, ordinal, vol = (
        md.get("page", ""),
        md.get("p_id", ""),
        md.get("ordinal", ""),
        md.get("fr_volume", ""),
    )
    if p_id and page:
        url = f"{base}#p-{p_id}"
        frag = text_fragment(snippet)
        if frag:
            url = f"{url}:~:text={frag}"
        label = f"{vol} FR {page} ¶{ordinal}" if vol and ordinal else f"p-{p_id}"
        return url, label
    if page:
        label = f"{vol} FR {page}" if vol else f"page {page}"
        return f"{base}#page-{page}", label
    return base, _short(md.get("title", "")) or md.get("item_key", "")


def _corpus(md: dict[str, str]) -> tuple[str, str]:
    url, page = md.get("url", ""), md.get("page", "")
    if url and page and url.lower().endswith(".pdf"):
        url = f"{url}#page={page}"
    if md.get("doctype") == "manual":
        # IOM chapter section: "IOM 100-04 ch.18 §10.1.2", linked to the
        # chapter PDF page the section was located on
        from llm.manuals import label as manual_label

        label = manual_label(md)
        if label:
            return url, label
    title = _short(md.get("title", ""))
    year = md.get("year", "") or md.get("date", "")[:4]
    if title and year:
        label = f"{title} ({year})"
    else:
        label = title or md.get("item_key", "")
    return url, label


def for_source(md: dict[str, str], snippet: str) -> tuple[str, str]:
    """``(url, label)`` for a chunk's metadata. Never raises; degrades
    to whatever link the metadata supports."""
    kind = md.get("kind", "") or ("comment" if md.get("comment_id") else "corpus")
    if kind == "comment":
        return _comment(md)
    if kind == "rule":
        return _rule(md, snippet)
    return _corpus(md)


def iom_refs_of(text: str) -> list[dict[str, str]]:
    """``[{"pub", "chapter", "section", "label"}]`` for every CMS manual
    citation in *text* ("Pub. 100-04, Chapter 12, Section 30.6.4"),
    de-duplicated in order of appearance."""
    from llm.manuals import label as manual_label
    from pfs.guidance import iom_refs

    out: list[dict[str, str]] = []
    seen: set[tuple[str, str, str]] = set()
    for pub, chapter, section in iom_refs(text or ""):
        key = (pub, chapter, section)
        if key in seen:
            continue
        seen.add(key)
        out.append(
            {
                "pub": pub,
                "chapter": chapter,
                "section": section,
                "label": manual_label(
                    {"manual": pub, "chapter": chapter, "iom_section": section}
                ),
            }
        )
    return out


def as_source(md: dict[str, str], text: str, score: float) -> dict:
    """The source dict the chat sends to the prompt and the UI.

    ``item_key``, ``p_id`` (rules only; ``""`` otherwise), ``seq``
    (non-rule kinds only; ``""`` for rules) and ``section`` are carried
    straight from *md* so the golden evaluation can match a source back
    to its chunk anchor: ``(item_key, p_id)`` for a rule paragraph,
    ``(item_key, seq)`` for a comment/corpus chunk.
    """
    snippet = text[:_SNIPPET_CHARS].strip()
    url, label = for_source(md, snippet)
    kind = md.get("kind", "")
    return {
        "id": label,
        "label": label,
        "kind": kind,
        # proposed / final / correction for a rule (the UI badge and the
        # prompt's source line show it; the citation token stays the FR cite)
        "rule_kind": md.get("rule_kind", "") if kind == "rule" else "",
        "doctype": md.get("doctype", ""),
        # IOM chapter coordinates for manual sections (llm.manuals)
        "manual": md.get("manual", ""),
        "chapter": md.get("chapter", ""),
        "iom_section": md.get("iom_section", ""),
        # theme slugs stamped by llm.themes (comma-joined), for filters and badges
        "themes": md.get("themes", ""),
        # IOM sections a rule paragraph cites (llm.iomlink pulls them in as
        # companion sources; the search page links them)
        "iom_refs": iom_refs_of(text) if kind == "rule" else [],
        "url": url,
        "title": md.get("title", ""),
        "date": md.get("date", ""),
        "docket": md.get("docket", ""),
        "comment_id": md.get("comment_id", ""),
        "snippet": snippet,
        "score": round(score, 4),
        "item_key": md.get("item_key", ""),
        "p_id": md.get("p_id", "") if kind == "rule" else "",
        "seq": md.get("seq", "") if kind != "rule" else "",
        "section": md.get("section", ""),
        # the PDF the chunk came from and the page it was located on
        # (``llm.pages.enrich_pdf_pages``) — the UI links them to the viewer
        "attachment": md.get("attachment", ""),
        "page": md.get("page", ""),
    }