stack/notebooks/llm_search.py

import marimo

__generated_with = "0.23.13"
app = marimo.App(width="medium")


@app.cell(hide_code=True)
def _():
    import marimo as mo

    return (mo,)


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    # Semantic search over a docket

    A notebook-side look at what `LlmClient` (`llm.client`, #573) gives you: the
    same pgvector similarity search the chat UI's `GET /search` endpoint runs,
    called directly from a marimo cell over plain httpx — no `langchain`,
    `sqlalchemy` or `psycopg` required, which matters because the notebooks
    container installs the workspace without the `llm` extra (#720). Pick a
    docket and a query below; the semantic-search hits sit beside the DuckDB
    replica's own comment-analysis table for the same docket, so you can see
    both views of "what commenters said" at once.
    """)
    return


@app.cell(hide_code=True)
def _():
    # ── Setup ──
    import polars as pl

    from conf import connect
    from llm.client import LlmClient

    NOTES = {}

    def _open_replica():
        try:
            return connect.duckdb("aco", read_only=True)
        except Exception as e:  # noqa: BLE001 — degrade, never crash the page
            NOTES["replica"] = f"replica unavailable: {e}"
            return None

    def _open_bib():
        try:
            return connect.bib()
        except Exception as e:  # noqa: BLE001
            NOTES["bib"] = f"bibliography unavailable: {e}"
            return None

    con = _open_replica()
    store = _open_bib()
    client = LlmClient()

    def _service_up():
        try:
            return client.health()
        except Exception:  # noqa: BLE001 — health() never raises; belt & suspenders
            return False

    llm_up = _service_up()
    if not llm_up:
        NOTES["llm"] = f"llm service unreachable at {client.base_url}"

    return NOTES, client, con, llm_up, pl, store


@app.cell(hide_code=True)
def _(mo, store):
    # ── Docket + query pickers ──
    _DEFAULT_DOCKETS = ["CMS-2023-0121", "CMS-2025-0304"]

    def _docket_options():
        if store is not None:
            try:
                ids = sorted({d.id for d in store.dockets()})
            except Exception:  # noqa: BLE001 — fall back to the hard-coded list
                ids = []
            if ids:
                return ids
        return list(_DEFAULT_DOCKETS)

    _options = _docket_options()
    _default = "CMS-2025-0304" if "CMS-2025-0304" in _options else _options[0]
    docket_picker = mo.ui.dropdown(options=_options, value=_default, label="Docket")
    query_box = mo.ui.text(
        value="telehealth",
        label="Query",
        full_width=True,
        placeholder="a phrase to search for",
    )
    mo.vstack([docket_picker, query_box])
    return docket_picker, query_box


@app.cell(hide_code=True)
def _(NOTES, client, docket_picker, llm_up, mo, pl, query_box):
    # ── Semantic search (LlmClient.search) ──
    _docket = docket_picker.value
    _q = (query_box.value or "").strip()

    search_results = pl.DataFrame()
    if not llm_up:
        search_view = mo.md(
            f"_llm service unreachable at `{client.base_url}` — search skipped. "
            "Set `LLM_URL` or run this inside the notebooks container "
            "(`gateway` network) to reach it._"
        )
    elif not _q:
        search_view = mo.md("_Type a query above to search._")
    else:
        try:
            _hits = client.search(_q, collection="comments", docket=_docket, limit=10)
        except Exception as e:  # noqa: BLE001 — degrade, never crash the page
            NOTES["search"] = f"search failed: {e}"
            _hits = []
        if _hits:
            search_results = pl.DataFrame(
                [
                    {
                        "label": h.get("label", ""),
                        "date": h.get("date", ""),
                        "snippet": h.get("snippet", ""),
                        "url": h.get("url", ""),
                    }
                    for h in _hits
                ]
            )
            search_view = mo.ui.table(
                search_results, label=f"Semantic search — {_docket} ({len(_hits)} hits)"
            )
        else:
            search_view = mo.md(
                f"_No semantic-search hits for `{_q}` in {_docket} — either "
                "nothing matches, or the docket isn't indexed yet "
                "(`stack llm index`)._"
            )
    return search_results, search_view


@app.cell(hide_code=True)
def _(NOTES, con, docket_picker, mo, pl):
    # ── DuckDB comments view (skin_subs.rulemaking_comments) ──
    _docket = docket_picker.value

    duckdb_comments = pl.DataFrame()
    if con is None:
        duckdb_view = mo.md("_replica unavailable — DuckDB comments view skipped._")
    else:
        try:
            duckdb_comments = con.execute(
                "SELECT comment_id, posted_date, organization, position, "
                "text_length, is_form_letter FROM skin_subs.rulemaking_comments "
                "WHERE docket_id = ? ORDER BY posted_date DESC",
                [_docket],
            ).pl()
        except Exception as e:  # noqa: BLE001 — a missing table is "not built yet"
            NOTES[f"duckdb:{_docket}"] = str(e)
        if duckdb_comments.is_empty():
            duckdb_view = mo.md(
                f"_No `skin_subs.rulemaking_comments` rows for {_docket} yet — "
                "this table only covers dockets run through "
                "`uv run python dev/scripts/classify_comments.py --docket "
                f"{_docket}` (the skin-substitute-relevant subset of a "
                "docket's comments, not every comment on it)._"
            )
        else:
            duckdb_view = mo.ui.table(
                duckdb_comments,
                label=(
                    f"DuckDB skin_subs.rulemaking_comments — {_docket} "
                    f"({duckdb_comments.height} rows)"
                ),
            )
    return (duckdb_view,)


@app.cell(hide_code=True)
def _(docket_picker, duckdb_view, mo, search_view):
    # ── Side by side ──
    mo.hstack(
        [
            mo.vstack(
                [mo.md(f"### Semantic search — {docket_picker.value}"), search_view]
            ),
            mo.vstack(
                [mo.md(f"### DuckDB comments — {docket_picker.value}"), duckdb_view]
            ),
        ],
        widths="equal",
        gap=2,
    )
    return


@app.cell(hide_code=True)
def _(NOTES, mo):
    # ── Notes ──
    mo.md(
        "### Notes\n\n"
        + (
            "\n".join(f"- {k}: {v}" for k, v in NOTES.items())
            if NOTES
            else "_All sources available._"
        )
    )
    return


if __name__ == "__main__":
    app.run()