diff --git a/notebooks/llm_search.py b/notebooks/llm_search.py new file mode 100644 index 0000000..23077c7 --- /dev/null +++ b/notebooks/llm_search.py @@ -0,0 +1,215 @@ +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() diff --git a/src/llm/client.py b/src/llm/client.py new file mode 100644 index 0000000..8294e36 --- /dev/null +++ b/src/llm/client.py @@ -0,0 +1,170 @@ +"""``LlmClient``: an httpx-only client for the llm chat/search service (#573). + +Notebook-side counterpart to the retrieval and chat endpoints exposed by +``llm.api`` (``GET /search``, ``GET /similar/{key}``, ``POST /chat``, +``GET /health``) — the same shape used by the chat UI, but callable from a +marimo cell without pulling in the ``llm`` extra (langchain/sqlalchemy/ +psycopg). The notebooks container installs the workspace *without* that +extra (#720); this module only needs ``httpx``, which it always has. + +Base URL resolution, in order: + 1. the *base_url* constructor argument + 2. the ``LLM_URL`` environment variable + 3. ``http://llm:8000`` — the compose service name, reachable because + the notebooks container sits on the same ``gateway`` network as + ``llm`` (compose.yml) + 4. ``http://localhost:8000`` — outside compose (llm publishes no host + port, so this only works when something else forwards it) + +Usage:: + + from llm.client import LlmClient + + client = LlmClient() + if client.health(): + hits = client.search("telehealth", docket="CMS-2025-0304") +""" + +from __future__ import annotations + +import json +import os +import socket +from typing import Any, Iterator + +import httpx + +_DEFAULT_TIMEOUT = 30.0 +_COMPOSE_URL = "http://llm:8000" +_LOCAL_URL = "http://localhost:8000" + + +def _on_gateway_network() -> bool: + """Whether the ``llm`` compose service name resolves from here — true + inside a container on the ``gateway`` network (e.g. ``notebooks``), + false on the host, where nothing publishes that name.""" + try: + socket.gethostbyname("llm") + return True + except OSError: + return False + + +def _resolve_base_url(base_url: str | None) -> str: + if base_url: + return base_url.rstrip("/") + env = os.environ.get("LLM_URL", "").strip() + if env: + return env.rstrip("/") + return _COMPOSE_URL if _on_gateway_network() else _LOCAL_URL + + +class LlmClient: + """Thin httpx client for the llm service's retrieval and chat endpoints. + + ``search``/``similar``/``chat`` raise ``httpx.HTTPError`` (with the + request context in the message) on a network failure or non-2xx + response. ``health`` never raises — it is meant for a notebook guard + cell that wants a plain yes/no before doing anything else. + """ + + def __init__( + self, + base_url: str | None = None, + *, + timeout: float = _DEFAULT_TIMEOUT, + ) -> None: + self.base_url = _resolve_base_url(base_url) + self.timeout = timeout + + def _client(self) -> httpx.Client: + return httpx.Client(base_url=self.base_url, timeout=self.timeout) + + def _get(self, path: str, *, params: dict[str, Any]) -> dict: + try: + with self._client() as http: + resp = http.get(path, params=params) + resp.raise_for_status() + return resp.json() + except httpx.HTTPError as exc: + raise httpx.HTTPError(f"GET {self.base_url}{path} failed: {exc}") from exc + + def health(self) -> bool: + """``True`` iff ``GET /health`` responds ok. Never raises.""" + try: + with self._client() as http: + resp = http.get("/health") + resp.raise_for_status() + return bool(resp.json().get("status") == "ok") + except (httpx.HTTPError, ValueError): + return False + + def search( + self, + q: str, + *, + collection: str = "all", + docket: str | None = None, + item_key: str | None = None, + year: str | int | None = None, + kind: str | None = None, + limit: int = 10, + offset: int = 0, + ) -> list[dict]: + """``GET /search`` — metadata-filtered similarity search. + + Mirrors ``llm.api.search_endpoint``'s query parameters; returns + the ``results`` list (each a source dict: ``label``, ``date``, + ``url``, ``snippet``, ``docket``, ``item_key``, ``score``, ...). + """ + params: dict[str, Any] = { + "q": q, + "collection": collection, + "limit": limit, + "offset": offset, + } + if docket: + params["docket"] = docket + if item_key: + params["item_key"] = item_key + if year: + params["year"] = str(year) + if kind: + params["kind"] = kind + return self._get("/search", params=params)["results"] + + def similar( + self, key: str, *, collection: str = "all", limit: int = 10 + ) -> list[dict]: + """``GET /similar/{key}`` — nearest neighbours of an already-indexed + bib item's own chunk(s). Returns the ``results`` list.""" + params = {"collection": collection, "limit": limit} + return self._get(f"/similar/{key}", params=params)["results"] + + def chat( + self, question: str, *, mode: str = "auto", since: str | None = None + ) -> Iterator[dict]: + """``POST /chat`` — stream the SSE events the chat UI consumes. + + Yields each ``data:`` line's JSON payload (``{"type": ..., ...}``) + as the server sends it; the caller decides what to do with + ``token``/``source``/``error``/etc. event types. + """ + payload: dict[str, Any] = {"question": question, "mode": mode} + if since: + payload["since"] = since + try: + with ( + self._client() as http, + http.stream("POST", "/chat", json=payload) as resp, + ): + resp.raise_for_status() + for line in resp.iter_lines(): + if not line.startswith("data:"): + continue + data = line[len("data:") :].strip() + if not data: + continue + yield json.loads(data) + except httpx.HTTPError as exc: + raise httpx.HTTPError(f"POST {self.base_url}/chat failed: {exc}") from exc diff --git a/tests/llm/test_client.py b/tests/llm/test_client.py new file mode 100644 index 0000000..df77ad8 --- /dev/null +++ b/tests/llm/test_client.py @@ -0,0 +1,272 @@ +"""``llm.client.LlmClient`` — httpx-only client for the llm service (#573). + +All requests go through ``httpx.MockTransport``; no real network I/O. +""" + +from __future__ import annotations + +import json + +import httpx +import pytest + +from llm.client import LlmClient, _resolve_base_url + + +def _client(handler, **kw) -> LlmClient: + """An LlmClient whose internal httpx.Client is wired to *handler*.""" + c = LlmClient(base_url="http://llm-test:8000", **kw) + transport = httpx.MockTransport(handler) + c._client = lambda: httpx.Client( # noqa: SLF001 — test seam + base_url=c.base_url, timeout=c.timeout, transport=transport + ) + return c + + +# ── base URL resolution ────────────────────────────────────────────── + + +def test_explicit_base_url_wins(): + assert _resolve_base_url("http://example:9000/") == "http://example:9000" + + +def test_env_var_used_when_no_explicit_base_url(monkeypatch): + monkeypatch.setenv("LLM_URL", "http://from-env:1234/") + assert _resolve_base_url(None) == "http://from-env:1234" + + +def test_falls_back_to_compose_name_when_it_resolves(monkeypatch): + monkeypatch.delenv("LLM_URL", raising=False) + monkeypatch.setattr("llm.client.socket.gethostbyname", lambda host: "192.168.5.9") + assert _resolve_base_url(None) == "http://llm:8000" + + +def test_falls_back_to_localhost_when_compose_name_does_not_resolve(monkeypatch): + import socket + + monkeypatch.delenv("LLM_URL", raising=False) + + def _raise(host): + raise socket.gaierror("not found") + + monkeypatch.setattr("llm.client.socket.gethostbyname", _raise) + assert _resolve_base_url(None) == "http://localhost:8000" + + +def test_default_constructor_resolves_base_url(monkeypatch): + monkeypatch.delenv("LLM_URL", raising=False) + monkeypatch.setattr( + "llm.client.socket.gethostbyname", + lambda host: (_ for _ in ()).throw(OSError("no")), + ) + client = LlmClient() + assert client.base_url == "http://localhost:8000" + assert client.timeout == 30.0 + + +# ── health ──────────────────────────────────────────────────────────── + + +def test_health_true_on_ok(): + def handler(request: httpx.Request) -> httpx.Response: + assert request.url.path == "/health" + return httpx.Response(200, json={"status": "ok"}) + + assert _client(handler).health() is True + + +def test_health_false_on_bad_status(): + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(503, json={"status": "down"}) + + assert _client(handler).health() is False + + +def test_health_false_on_connect_error(): + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("refused", request=request) + + assert _client(handler).health() is False + + +def test_health_false_on_bad_json(): + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, text="not json") + + assert _client(handler).health() is False + + +# ── search ──────────────────────────────────────────────────────────── + + +def test_search_sends_expected_query_params_and_returns_results(): + seen = {} + + def handler(request: httpx.Request) -> httpx.Response: + seen["path"] = request.url.path + seen["params"] = dict(request.url.params) + return httpx.Response( + 200, + json={ + "query": "telehealth", + "filters": {}, + "total": 1, + "results": [ + { + "label": "CMS-2025-0304-1", + "url": "https://x", + "date": "2025-01-01", + "snippet": "hi", + } + ], + }, + ) + + results = _client(handler).search( + "telehealth", + collection="comments", + docket="CMS-2025-0304", + item_key="ABC123", + year=2025, + kind="comment", + limit=5, + offset=2, + ) + assert seen["path"] == "/search" + assert seen["params"] == { + "q": "telehealth", + "collection": "comments", + "limit": "5", + "offset": "2", + "docket": "CMS-2025-0304", + "item_key": "ABC123", + "year": "2025", + "kind": "comment", + } + assert results == [ + { + "label": "CMS-2025-0304-1", + "url": "https://x", + "date": "2025-01-01", + "snippet": "hi", + } + ] + + +def test_search_omits_unset_optional_filters(): + seen = {} + + def handler(request: httpx.Request) -> httpx.Response: + seen["params"] = dict(request.url.params) + return httpx.Response(200, json={"results": []}) + + _client(handler).search("q") + assert seen["params"] == { + "q": "q", + "collection": "all", + "limit": "10", + "offset": "0", + } + + +def test_search_raises_httperror_on_bad_status(): + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(400, json={"detail": "unknown collection"}) + + with pytest.raises(httpx.HTTPError): + _client(handler).search("q", collection="nope") + + +def test_search_raises_httperror_on_connect_error(): + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("refused", request=request) + + with pytest.raises(httpx.HTTPError): + _client(handler).search("q") + + +# ── similar ─────────────────────────────────────────────────────────── + + +def test_similar_hits_expected_path_and_returns_results(): + seen = {} + + def handler(request: httpx.Request) -> httpx.Response: + seen["path"] = request.url.path + seen["params"] = dict(request.url.params) + return httpx.Response( + 200, json={"key": "ABC123", "total": 1, "results": [{"label": "x"}]} + ) + + results = _client(handler).similar("ABC123", collection="rules", limit=3) + assert seen["path"] == "/similar/ABC123" + assert seen["params"] == {"collection": "rules", "limit": "3"} + assert results == [{"label": "x"}] + + +def test_similar_raises_httperror_on_404(): + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(404, json={"detail": "no indexed chunks"}) + + with pytest.raises(httpx.HTTPError): + _client(handler).similar("nope") + + +# ── chat (SSE) ──────────────────────────────────────────────────────── + + +def _sse_body(*events: dict) -> bytes: + return "".join(f"data: {json.dumps(e)}\n\n" for e in events).encode() + + +def test_chat_parses_sse_events_in_order(): + events = [ + {"type": "token", "text": "Hel"}, + {"type": "token", "text": "lo"}, + {"type": "done"}, + ] + + def handler(request: httpx.Request) -> httpx.Response: + assert request.method == "POST" + assert json.loads(request.content) == {"question": "hi?", "mode": "auto"} + return httpx.Response(200, content=_sse_body(*events)) + + got = list(_client(handler).chat("hi?")) + assert got == events + + +def test_chat_includes_since_and_mode_when_given(): + seen = {} + + def handler(request: httpx.Request) -> httpx.Response: + seen["body"] = json.loads(request.content) + return httpx.Response(200, content=_sse_body({"type": "done"})) + + list(_client(handler).chat("q", mode="timeline", since="2025-01-01")) + assert seen["body"] == {"question": "q", "mode": "timeline", "since": "2025-01-01"} + + +def test_chat_ignores_blank_and_non_data_lines(): + raw = b': comment line\n\ndata: {"type": "done"}\n\n\n' + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, content=raw) + + got = list(_client(handler).chat("q")) + assert got == [{"type": "done"}] + + +def test_chat_raises_httperror_on_bad_status(): + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(400, json={"detail": "empty question"}) + + with pytest.raises(httpx.HTTPError): + list(_client(handler).chat("")) + + +def test_chat_raises_httperror_on_connect_error(): + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("refused", request=request) + + with pytest.raises(httpx.HTTPError): + list(_client(handler).chat("q")) diff --git a/tests/llm/test_init_lazy.py b/tests/llm/test_init_lazy.py index 0e39c9e..6d706a3 100644 --- a/tests/llm/test_init_lazy.py +++ b/tests/llm/test_init_lazy.py @@ -2,7 +2,11 @@ ``llm`` extra (no langchain, no sqlalchemy). ``llm.config``, ``llm.links`` and ``llm.lineage`` must import there anyway — the pool exports are resolved lazily, and ``llm.evidence`` only touches -sqlalchemy when a pgvector query actually runs.""" +sqlalchemy when a pgvector query actually runs. + +#573: ``llm.client`` (the notebook-side ``LlmClient``) is httpx-only and +must import there too — it's the one module notebooks are expected to +import directly.""" from __future__ import annotations @@ -21,8 +25,9 @@ def _fake(name, *a, **k): return _real(name, *a, **k) builtins.__import__ = _fake import llm.config, llm.links, llm.lineage, llm.evidence # must not need langchain/sqlalchemy +import llm.client # #573: notebook-side LlmClient — httpx only import llm -print('ok', llm.LlmConfig.__name__) +print('ok', llm.LlmConfig.__name__, llm.client.LlmClient.__name__) """ @@ -31,7 +36,7 @@ def test_config_links_lineage_evidence_import_without_langchain_or_sqlalchemy(): [sys.executable, "-c", _PROBE], capture_output=True, text=True, check=False ) assert r.returncode == 0, r.stderr[-800:] - assert "ok LlmConfig" in r.stdout + assert "ok LlmConfig LlmClient" in r.stdout def test_pool_exports_resolve_lazily(): diff --git a/tests/notebooks/test_llm_search_nb.py b/tests/notebooks/test_llm_search_nb.py new file mode 100644 index 0000000..e8999a9 --- /dev/null +++ b/tests/notebooks/test_llm_search_nb.py @@ -0,0 +1,107 @@ +"""notebooks/llm_search.py — structure and headless degradation (#573).""" + +from __future__ import annotations + +import ast +import importlib.util +from pathlib import Path + +import pytest + +NB = Path(__file__).resolve().parents[2] / "notebooks" / "llm_search.py" + + +def _load(): + spec = importlib.util.spec_from_file_location("llm_search_nb", NB) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +def test_notebook_is_a_marimo_app(): + mod = _load() + assert mod.app.__class__.__name__ == "App" + + +def test_cells_are_anonymous(): + src = NB.read_text() + tree = ast.parse(src) + names = [n.name for n in tree.body if isinstance(n, ast.FunctionDef)] + assert names and set(names) == {"_"} + + +def test_headless_run_degrades_without_service_or_data(monkeypatch): + """With the llm service unreachable and no replica/bib, every guard cell + should render a note instead of raising.""" + mod = _load() + + import conf.connect as cc + from llm.client import LlmClient + + monkeypatch.setattr(LlmClient, "health", lambda self: False) + monkeypatch.setattr( + cc, + "duckdb", + lambda *a, **k: (_ for _ in ()).throw(FileNotFoundError("no replica")), + ) + monkeypatch.setattr( + cc, "bib", lambda *a, **k: (_ for _ in ()).throw(FileNotFoundError("no bib")) + ) + + outputs, _defs = mod.app.run() + assert outputs is not None + rendered = "\n".join(o._repr_html_() for o in outputs if hasattr(o, "_repr_html_")) + assert "Traceback" not in rendered + assert "llm service unreachable" in rendered + assert "replica unavailable" in rendered + assert "bibliography unavailable" in rendered + + +def test_headless_run_still_renders_with_service_up_but_no_data(monkeypatch): + """Health true but no replica/bib and no real search results (network + still stubbed out) — the search cell should degrade to "no hits", not + raise, and the duckdb cell should still show "replica unavailable".""" + mod = _load() + + import conf.connect as cc + from llm.client import LlmClient + + monkeypatch.setattr(LlmClient, "health", lambda self: True) + monkeypatch.setattr( + LlmClient, + "search", + lambda self, *a, **k: (_ for _ in ()).throw( + __import__("httpx").ConnectError("refused") + ), + ) + monkeypatch.setattr( + cc, + "duckdb", + lambda *a, **k: (_ for _ in ()).throw(FileNotFoundError("no replica")), + ) + monkeypatch.setattr( + cc, "bib", lambda *a, **k: (_ for _ in ()).throw(FileNotFoundError("no bib")) + ) + + outputs, _defs = mod.app.run() + rendered = "\n".join(o._repr_html_() for o in outputs if hasattr(o, "_repr_html_")) + assert "Traceback" not in rendered + assert "replica unavailable" in rendered + + +_ROOT = Path(__file__).resolve().parents[2] +_HAS_DATA = (_ROOT / "data" / "replica" / "aco.ro.duckdb").exists() and ( + _ROOT / "data" / "bib.sqlite" +).exists() + + +@pytest.mark.skipif(not _HAS_DATA, reason="needs the live replica and bib") +def test_docket_picker_offers_known_dockets_against_real_bib(monkeypatch): + """Live smoke test: with the real bib store, the docket dropdown's + options come from ``store.dockets()`` and are non-empty.""" + from llm.client import LlmClient + + monkeypatch.setattr(LlmClient, "health", lambda self: False) + mod = _load() + outputs, _defs = mod.app.run() + assert outputs is not None