homelab/stack
1
0
Fork 0
Code Issues 114 Pull Requests 1 Actions Packages Projects 1 Releases Wiki Activity

feat(llm): whole-library retrieval with recency blend, deep-linked sources, largest-GPU generation (refs #571 #654)

Browse Source
This commit is contained in:
kert
2026-09-03 10:49:16 -04:00
parent d3ae6e9300
commit c63208f5fe
2 changed files with 365 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files

145
src/llm/rag.py
View File

@@ -1,83 +1,150 @@
"""RAG chain for the chat UI: retrieve comments, stream a grounded answer.
"""RAG chain for the chat UI: retrieve across the library, stream a
grounded answer from the largest live GPU.
Single-shot and stateless — each question is retrieved and answered on its
own. Retrieval hits pgvector's ``comments`` collection via the embedding
pool; generation streams the configured instruct model from an Ollama host.
Single-shot and stateless — each question is retrieved and answered on
its own. Retrieval embeds the question once and searches the
``comments``, ``rules`` and ``corpus`` pgvector collections; hits are
merged and re-ranked by similarity × recency (``llm.rerank``), then each
source gets a deep link (``llm.links``). Generation streams from the
largest live Ollama host (``HostPool.acquire_generation``) with the
model tier that host can hold (``pick_model``).
"""
from __future__ import annotations
import json
from datetime import date
from typing import Iterator
import httpx
from llm.config import LlmConfig
from llm.pool import HostPool
from llm.links import for_source
from llm.pool import HostPool, PoolEmbeddings, pick_model
from llm.rerank import Hit, blend, filter_since
_TIMEOUT = httpx.Timeout(300.0, connect=5.0)
_COLLECTIONS = {"comment": "comments", "rule": "rules", "corpus": "corpus"}
_OVERFETCH = 3
_SNIPPET_CHARS = 500
_SYSTEM = (
"You answer questions about public comments submitted to CMS rulemaking "
"dockets. Use ONLY the comment excerpts provided below. When you use an "
"excerpt, cite its id in square brackets, e.g. [CMS-2017-0092-1306]. If "
"the excerpts do not contain the answer, say you don't have information "
"on that in the indexed comments — do not invent facts."
"You answer questions about CMS rulemaking using ONLY the excerpts "
"provided below. Excerpts come from three kinds of sources: public "
"comments submitted to regulations.gov dockets, Federal Register rules "
"(proposed and final), and a reference library (journal articles, CMS "
"manuals, regulations, agency documents). Each excerpt is prefixed with "
"its citation label in square brackets and its kind and date. When you "
"use an excerpt, cite its label exactly, e.g. [CMS-2026-2377-3438] or "
"[91 FR 43949 ¶4]. Prefer the most recent comments when excerpts "
"conflict or describe a changing position, and say what year a "
"statement comes from when it matters. If the excerpts do not contain "
"the answer, say you don't have information on that in the indexed "
"library — do not invent facts."
)
def retrieve(
question: str, *, cfg: LlmConfig, pool: HostPool, k: int = 6
) -> list[dict]:
"""Top-k comment chunks for ``question`` as source dicts."""
def _hits(question_vec: list[float], *, cfg: LlmConfig, pool: HostPool) -> list[Hit]:
from llm.index import vectorstore
store = vectorstore("comments", cfg, pool)
hits = store.similarity_search_with_score(question, k=k)
sources = []
for doc, score in hits:
md = doc.metadata or {}
sources.append(
{
"comment_id": md.get("comment_id") or md.get("item_key", ""),
"docket": md.get("docket", ""),
"snippet": doc.page_content[:500].strip(),
"score": round(float(score), 4),
}
)
return sources
hits: list[Hit] = []
for kind, collection in _COLLECTIONS.items():
k = int(cfg.k_per_kind.get(kind, 0))
if k <= 0:
continue
store = vectorstore(collection, cfg, pool)
for doc, distance in store.similarity_search_with_score_by_vector(
question_vec, k=k * _OVERFETCH
):
md = {key: str(v) for key, v in (doc.metadata or {}).items()}
md.setdefault("kind", kind)
hits.append(
Hit(text=doc.page_content, metadata=md, distance=float(distance))
)
return hits
def _source(hit: Hit) -> dict:
md = hit.metadata
snippet = hit.text[:_SNIPPET_CHARS].strip()
url, label = for_source(md, snippet)
return {
"id": label,
"label": label,
"kind": md.get("kind", ""),
"url": url,
"title": md.get("title", ""),
"date": md.get("date", ""),
"docket": md.get("docket", ""),
"comment_id": md.get("comment_id", ""),
"snippet": snippet,
"score": round(hit.score, 4),
}
def retrieve(
question: str,
*,
cfg: LlmConfig,
pool: HostPool,
since: str = "",
now: date | None = None,
) -> list[dict]:
"""Top sources for ``question`` across all collections, recency-blended.
``since`` (ISO date) hard-filters to material dated on/after it.
"""
vec = PoolEmbeddings(pool, cfg.embed_model).embed_query(question)
hits = filter_since(_hits(vec, cfg=cfg, pool=pool), since)
ranked = blend(
hits,
weight=cfg.recency_weight,
half_life_days=cfg.recency_half_life_days,
now=now or date.today(),
top_n=cfg.top_n,
)
return [_source(h) for h in ranked]
def build_messages(question: str, sources: list[dict]) -> list[dict]:
"""Grounded chat messages: system rules + question with excerpts."""
if sources:
context = "\n\n".join(f"[{s['comment_id']}] {s['snippet']}" for s in sources)
context = "\n\n".join(
f"[{s['label']}] ({s.get('kind', '')}, {s.get('date', '') or 'undated'}) "
f"{s['snippet']}"
for s in sources
)
else:
context = "(no relevant comments found)"
user = f"Comment excerpts:\n\n{context}\n\nQuestion: {question}"
context = "(no relevant excerpts found)"
user = f"Excerpts:\n\n{context}\n\nQuestion: {question}"
return [
{"role": "system", "content": _SYSTEM},
{"role": "user", "content": user},
]
def stream_answer(question: str, *, cfg: LlmConfig, pool: HostPool) -> Iterator[dict]:
"""Retrieve, then stream a grounded answer.
def stream_answer(
question: str, *, cfg: LlmConfig, pool: HostPool, since: str = ""
) -> Iterator[dict]:
"""Retrieve, then stream a grounded answer from the largest live host.
Yields ``{"type":"token","text":…}`` events as the model generates, then
one ``{"type":"sources",…}`` and a final ``{"type":"done"}``.
Yields ``{"type":"token","text":…}`` events as the model generates,
then one ``{"type":"sources", "sources": […], "model": …, "host": …}``
and a final ``{"type":"done"}``.
"""
sources = retrieve(question, cfg=cfg, pool=pool)
sources = retrieve(question, cfg=cfg, pool=pool, since=since)
pool.check(cfg.instruct_model)
messages = build_messages(question, sources)
with pool.acquire() as host, httpx.Client(timeout=_TIMEOUT) as client:
with pool.acquire_generation() as host, httpx.Client(timeout=_TIMEOUT) as client:
model = pick_model(cfg, pool, host)
with client.stream(
"POST",
f"{host}/api/chat",
json={
"model": cfg.instruct_model,
"model": model,
"messages": messages,
"stream": True,
"options": {"num_ctx": cfg.chat_num_ctx},
},
) as resp:
resp.raise_for_status()
@@ -90,5 +157,5 @@ def stream_answer(question: str, *, cfg: LlmConfig, pool: HostPool) -> Iterator[
yield {"type": "token", "text": chunk}
if data.get("done"):
break
yield {"type": "sources", "sources": sources}
yield {"type": "sources", "sources": sources, "model": model, "host": host}
yield {"type": "done"}

315
tests/llm/test_rag.py
View File

@@ -1,7 +1,9 @@
"""llm.rag — retrieval + grounded streaming answer."""
"""llm.rag — multi-collection retrieval + grounded streaming answer."""
from datetime import date
from unittest.mock import MagicMock, patch
import pytest
from langchain_core.documents import Document
from llm.config import LlmConfig
@@ -9,117 +11,318 @@ from llm.rag import build_messages, retrieve, stream_answer
CFG = LlmConfig(
ollama_hosts=("http://h1:11434",),
host_vram={"http://h1:11434": 24},
embed_model="embed",
instruct_model="chat",
instruct_model_large="big",
embed_dim=768,
pg_host="x",
pg_port=5432,
pg_db="llm",
pg_user="llm",
build_ann_index=True,
k_per_kind={"comment": 2, "rule": 1, "corpus": 1},
top_n=3,
)
NOW = date(2026, 9, 3)
def _doc(text, **md):
return Document(page_content=text, metadata=md)
class TestRetrieve:
@patch("llm.index.vectorstore")
def test_maps_hits_to_sources(self, mock_vs):
store = mock_vs.return_value
store.similarity_search_with_score.return_value = [
(
_doc(
"Telehealth comment.",
comment_id="CMS-2017-0092-1306",
docket="CMS-2017-0092",
item_key="K1",
),
0.21,
),
]
pool = MagicMock()
out = retrieve("telehealth", cfg=CFG, pool=pool, k=3)
assert out == [
{
"comment_id": "CMS-2017-0092-1306",
"docket": "CMS-2017-0092",
"snippet": "Telehealth comment.",
"score": 0.21,
}
]
store.similarity_search_with_score.assert_called_once_with("telehealth", k=3)
def _stores(by_collection):
"""vectorstore(collection, cfg, pool) → a store whose
similarity_search_with_score_by_vector returns by_collection[name]."""
def factory(collection, cfg, pool):
s = MagicMock()
s.similarity_search_with_score_by_vector.return_value = by_collection.get(
collection, []
)
return s
return factory
class TestRetrieve:
@patch("llm.rag.PoolEmbeddings")
@patch("llm.index.vectorstore")
def test_falls_back_to_item_key_when_no_comment_id(self, mock_vs):
store = mock_vs.return_value
store.similarity_search_with_score.return_value = [
(_doc("x", item_key="K9", docket="D"), 0.5)
]
out = retrieve("q", cfg=CFG, pool=MagicMock())
assert out[0]["comment_id"] == "K9"
def test_merges_kinds_and_builds_links(self, mock_vs, MockEmb):
MockEmb.return_value.embed_query.return_value = [0.1] * 3
mock_vs.side_effect = _stores(
{
"comments": [
(
_doc(
"Telehealth comment.",
kind="comment",
comment_id="CMS-2026-2377-3438",
docket="CMS-2026-2377",
item_key="K1",
date="2026-08-19",
title="Anand M.",
),
0.25,
)
],
"rules": [
(
_doc(
"Under this proposal, the new G codes apply.",
kind="rule",
item_key="R1",
html_url="https://fr.test/doc",
p_id="935",
page="43949",
ordinal="1",
fr_volume="91",
date="2026-07-16",
title="CY2027 PFS NPRM",
),
0.20,
)
],
"corpus": [],
}
)
out = retrieve("telehealth", cfg=CFG, pool=MagicMock(), now=NOW)
assert [s["kind"] for s in out] == ["rule", "comment"]
rule, comment = out
assert rule["label"] == "91 FR 43949 ¶1"
assert rule["id"] == rule["label"]
assert rule["url"].startswith("https://fr.test/doc#p-935:~:text=Under%20this")
assert (
comment["url"] == "https://www.regulations.gov/comment/CMS-2026-2377-3438"
)
assert comment["comment_id"] == "CMS-2026-2377-3438"
assert comment["docket"] == "CMS-2026-2377"
assert comment["date"] == "2026-08-19"
assert comment["snippet"] == "Telehealth comment."
assert 0 < comment["score"] <= 1
@patch("llm.rag.PoolEmbeddings")
@patch("llm.index.vectorstore")
def test_over_fetches_three_per_kind_and_embeds_once(self, mock_vs, MockEmb):
MockEmb.return_value.embed_query.return_value = [0.0]
stores = {}
def factory(collection, cfg, pool):
s = MagicMock()
s.similarity_search_with_score_by_vector.return_value = []
stores[collection] = s
return s
mock_vs.side_effect = factory
retrieve("q", cfg=CFG, pool=MagicMock(), now=NOW)
assert MockEmb.return_value.embed_query.call_count == 1
comments = stores["comments"].similarity_search_with_score_by_vector
rules = stores["rules"].similarity_search_with_score_by_vector
assert comments.call_args.kwargs["k"] == 6
assert rules.call_args.kwargs["k"] == 3
@patch("llm.rag.PoolEmbeddings")
@patch("llm.index.vectorstore")
def test_since_filters_old_hits(self, mock_vs, MockEmb):
MockEmb.return_value.embed_query.return_value = [0.0]
mock_vs.side_effect = _stores(
{
"comments": [
(
_doc(
"old",
kind="comment",
comment_id="C-1",
item_key="A",
date="2019-01-01",
),
0.1,
),
(
_doc(
"new",
kind="comment",
comment_id="C-2",
item_key="B",
date="2026-01-01",
),
0.3,
),
]
}
)
out = retrieve("q", cfg=CFG, pool=MagicMock(), since="2025-01-01", now=NOW)
assert [s["comment_id"] for s in out] == ["C-2"]
@patch("llm.rag.PoolEmbeddings")
@patch("llm.index.vectorstore")
def test_recent_comment_outranks_slightly_closer_old_one(self, mock_vs, MockEmb):
MockEmb.return_value.embed_query.return_value = [0.0]
mock_vs.side_effect = _stores(
{
"comments": [
(
_doc(
"old",
kind="comment",
comment_id="C-1",
item_key="A",
date="2019-01-01",
),
0.20,
),
(
_doc(
"new",
kind="comment",
comment_id="C-2",
item_key="B",
date="2026-08-19",
),
0.25,
),
]
}
)
out = retrieve("q", cfg=CFG, pool=MagicMock(), now=NOW)
assert [s["comment_id"] for s in out] == ["C-2", "C-1"]
@patch("llm.rag.PoolEmbeddings")
@patch("llm.index.vectorstore")
def test_legacy_chunks_without_kind_are_treated_as_comments(self, mock_vs, MockEmb):
MockEmb.return_value.embed_query.return_value = [0.0]
mock_vs.side_effect = _stores(
{
"comments": [
(_doc("x", comment_id="C-9", item_key="K9", docket="D"), 0.5)
]
}
)
(s,) = retrieve("q", cfg=CFG, pool=MagicMock(), now=NOW)
assert s["kind"] == "comment" and s["label"] == "C-9"
class TestBuildMessages:
def test_includes_ids_and_abstention_rule(self):
def test_includes_labels_kinds_dates_and_rules(self):
sources = [
{
"comment_id": "CMS-2017-0092-1",
"docket": "D",
"id": "91 FR 43949 ¶1",
"label": "91 FR 43949 ¶1",
"kind": "rule",
"date": "2026-07-16",
"snippet": "Under this proposal",
"url": "u",
"title": "t",
"docket": "",
"comment_id": "",
"score": 0.9,
},
{
"id": "CMS-2026-2377-1",
"label": "CMS-2026-2377-1",
"kind": "comment",
"date": "2026-08-19",
"snippet": "reduce documentation",
"score": 0.1,
}
"url": "u",
"title": "t",
"docket": "CMS-2026-2377",
"comment_id": "CMS-2026-2377-1",
"score": 0.8,
},
]
msgs = build_messages("why?", sources)
sys_msg = msgs[0]["content"].lower()
assert msgs[0]["role"] == "system"
assert "only" in msgs[0]["content"].lower()
assert "don't have information" in msgs[0]["content"].lower()
assert "[CMS-2017-0092-1]" in msgs[1]["content"]
assert "only" in sys_msg
assert "don't have information" in sys_msg
assert "most recent" in sys_msg
assert (
"[91 FR 43949 ¶1] (rule, 2026-07-16) Under this proposal"
in msgs[1]["content"]
)
assert "[CMS-2026-2377-1] (comment, 2026-08-19)" in msgs[1]["content"]
assert "why?" in msgs[1]["content"]
def test_no_sources_marks_empty_context(self):
msgs = build_messages("q", [])
assert "no relevant comments" in msgs[1]["content"].lower()
assert "no relevant excerpts" in msgs[1]["content"].lower()
class TestStreamAnswer:
def _pool(self, vram=24.0, serves_big=True):
pool = MagicMock()
pool.acquire_generation.return_value.__enter__.return_value = "http://h1:11434"
pool.vram.return_value = vram
pool.serves.return_value = serves_big
return pool
@patch("llm.rag.httpx.Client")
@patch("llm.rag.retrieve")
def test_yields_tokens_then_sources_then_done(self, mock_retrieve, MockClient):
mock_retrieve.return_value = [
{"comment_id": "C1", "docket": "D", "snippet": "s", "score": 0.1}
]
src = {
"id": "C1",
"label": "C1",
"kind": "comment",
"snippet": "s",
"score": 0.1,
}
mock_retrieve.return_value = [src]
lines = [
'{"message":{"content":"Doc"},"done":false}',
"", # keep-alive blank line — must be skipped, not parsed
'{"message":{"content":"tors"},"done":false}',
'{"message":{"content":""},"done":true}',
]
stream_cm = MockClient.return_value.__enter__.return_value.stream.return_value
resp = stream_cm.__enter__.return_value
client = MockClient.return_value.__enter__.return_value
resp = client.stream.return_value.__enter__.return_value
resp.iter_lines.return_value = iter(lines)
pool = MagicMock()
pool = self._pool()
events = list(stream_answer("q", cfg=CFG, pool=pool))
pool.check.assert_called_once_with("chat")
mock_retrieve.assert_called_once_with("q", cfg=CFG, pool=pool, since="")
assert events[0] == {"type": "token", "text": "Doc"}
assert events[1] == {"type": "token", "text": "tors"}
assert events[-2] == {
"type": "sources",
"sources": [
{"comment_id": "C1", "docket": "D", "snippet": "s", "score": 0.1}
],
"sources": [src],
"model": "big",
"host": "http://h1:11434",
}
assert events[-1] == {"type": "done"}
body = client.stream.call_args.kwargs["json"]
assert body["model"] == "big"
assert body["options"] == {"num_ctx": 8192}
assert client.stream.call_args.args[1] == "http://h1:11434/api/chat"
@patch("llm.rag.httpx.Client")
@patch("llm.rag.retrieve")
def test_small_host_uses_baseline_model(self, mock_retrieve, MockClient):
mock_retrieve.return_value = []
client = MockClient.return_value.__enter__.return_value
resp = client.stream.return_value.__enter__.return_value
resp.iter_lines.return_value = iter(['{"message":{"content":""},"done":true}'])
events = list(stream_answer("q", cfg=CFG, pool=self._pool(vram=12.0)))
assert client.stream.call_args.kwargs["json"]["model"] == "chat"
assert events[-2]["model"] == "chat"
@patch("llm.rag.httpx.Client")
@patch("llm.rag.retrieve")
def test_since_forwarded(self, mock_retrieve, MockClient):
mock_retrieve.return_value = []
client = MockClient.return_value.__enter__.return_value
resp = client.stream.return_value.__enter__.return_value
resp.iter_lines.return_value = iter(['{"message":{"content":""},"done":true}'])
list(stream_answer("q", cfg=CFG, pool=self._pool(), since="2025-09-01"))
assert mock_retrieve.call_args.kwargs["since"] == "2025-09-01"
@patch("llm.rag.httpx.Client")
@patch("llm.rag.retrieve")
def test_http_error_propagates(self, mock_retrieve, MockClient):
mock_retrieve.return_value = []
stream_cm = MockClient.return_value.__enter__.return_value.stream.return_value
resp = stream_cm.__enter__.return_value
client = MockClient.return_value.__enter__.return_value
resp = client.stream.return_value.__enter__.return_value
resp.raise_for_status.side_effect = RuntimeError("ollama down")
with __import__("pytest").raises(RuntimeError, match="ollama down"):
list(stream_answer("q", cfg=CFG, pool=MagicMock()))
with pytest.raises(RuntimeError, match="ollama down"):
list(stream_answer("q", cfg=CFG, pool=self._pool()))