feat(llm,notebooks): httpx-only LlmClient for notebooks and an llm_search example notebook (refs #573)
This commit is contained in:
kert
215
notebooks/llm_search.py
Normal file
215
notebooks/llm_search.py
Normal file
@@ -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()
|
||||
170
src/llm/client.py
Normal file
170
src/llm/client.py
Normal file
@@ -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
|
||||
272
tests/llm/test_client.py
Normal file
272
tests/llm/test_client.py
Normal file
@@ -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"))
|
||||
@@ -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():
|
||||
|
||||
107
tests/notebooks/test_llm_search_nb.py
Normal file
107
tests/notebooks/test_llm_search_nb.py
Normal file
@@ -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
|
||||
Reference in New Issue
Block a user