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

feat(cli): stack llm index --collection all + stack llm hosts (refs #572 #654)

Browse Source
This commit is contained in:
kert
2026-09-03 10:48:06 -04:00
parent 94a44c9283
commit d3ae6e9300
4 changed files with 182 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

107
src/cli/llm.py
View File

@@ -1,52 +1,107 @@
"""stack llm — local RAG over comments + corpus (index + chat serve)."""
"""stack llm — local RAG over the library (index + fleet + chat serve)."""
import typer
app = typer.Typer(no_args_is_help=True)
_COLLECTIONS = ("comments", "rules", "corpus")
def _docs_for(collection: str, store, docket: str, keys: tuple[str, ...]):
from llm.source import (
ZoteroPdfIndex,
iter_comment_docs,
iter_corpus_docs,
iter_rule_docs,
)
if collection == "comments":
return iter_comment_docs(store, docket=docket)
if collection == "rules":
return iter_rule_docs(store, keys=keys)
from conf import ROOT, path
zotero = ZoteroPdfIndex.snapshot(
path("db.zotero"), path("storage.zotero"), ROOT / ".state" / "llm"
)
return iter_corpus_docs(store, zotero=zotero)
@app.command()
def index(
collection: str = typer.Option(
"comments", help="Which collection: comments | corpus | rules."
"comments", help="Which collection: comments | rules | corpus | all."
),
docket: str = typer.Option("", help="Limit comments to one docket id."),
key: list[str] = typer.Option([], "--key", help="Limit rules to these item keys."),
force: bool = typer.Option(False, help="Re-embed even when unchanged."),
limit: int = typer.Option(0, help="Stop after N docs (0 = all)."),
limit: int = typer.Option(0, help="Stop after N docs per collection (0 = all)."),
) -> None:
"""Embed comments/corpus/rules into pgvector (incremental, resumable)."""
"""Embed comments/rules/corpus into pgvector (incremental, resumable).
Comments are processed newest-posted first; rules come from their FR
paragraph anchor maps (run `stack bib fr-grab` first for exact links);
corpus = every non-comment bib item incl. Zotero-only storage PDFs.
"""
import itertools
from conf.connect import bib
from llm import config as llm_config
from llm.index import index_docs
from llm.pool import HostPool
from llm.source import iter_comment_docs, iter_corpus_docs, iter_rule_docs
targets = _COLLECTIONS if collection == "all" else (collection,)
if any(t not in _COLLECTIONS for t in targets):
raise typer.BadParameter("collection must be comments, rules, corpus or all")
cfg = llm_config.load()
store = bib()
if collection == "comments":
docs = iter_comment_docs(store, docket=docket)
elif collection == "corpus":
docs = iter_corpus_docs(store)
elif collection == "rules":
docs = iter_rule_docs(store, keys=tuple(key))
else:
raise typer.BadParameter("collection must be 'comments', 'corpus' or 'rules'")
if limit:
docs = itertools.islice(docs, limit)
stats = index_docs(
docs,
collection=collection,
cfg=cfg,
pool=HostPool.from_config(cfg),
force=force,
)
typer.echo(
f"indexed={stats['indexed']} skipped={stats['skipped']} "
f"chunks={stats['chunks']}"
)
for target in targets:
docs = _docs_for(target, store, docket, tuple(key))
if limit:
docs = itertools.islice(docs, limit)
stats = index_docs(
docs,
collection=target,
cfg=cfg,
pool=HostPool.from_config(cfg),
force=force,
)
typer.echo(
f"{target}: indexed={stats['indexed']} skipped={stats['skipped']} "
f"chunks={stats['chunks']}"
)
@app.command()
def hosts() -> None:
"""Show the Ollama fleet: declared VRAM, liveness, models, and which
host + model would answer a chat right now.
Source .env first (`set -a; . ./.env; set +a`) so LLM_OLLAMA_HOSTS
lists the whole fleet; the compose service gets it from its environment.
"""
from llm import config as llm_config
from llm.pool import HostPool, pick_model
cfg = llm_config.load()
pool = HostPool.from_config(cfg)
declared = pool.hosts
try:
live = pool.check(cfg.instruct_model)
except RuntimeError as e:
typer.echo(str(e))
raise typer.Exit(1)
for row in pool.status():
typer.echo(
f"{row['host']:<32} {row['vram_gb']:>5.0f} GB up "
f"{', '.join(row['models'])}"
)
for h in declared:
if h not in live:
state = f"up (no {cfg.instruct_model})" if pool.reachable(h) else "DOWN"
typer.echo(f"{h:<32} {pool.vram(h):>5.0f} GB {state}")
with pool.acquire_generation() as host:
typer.echo(f"generation → {pick_model(cfg, pool, host)} @ {host}")
@app.command()

4
src/llm/pool.py
View File

@@ -52,6 +52,10 @@ class HostPool:
def vram(self, host: str) -> float:
return self._vram.get(host.rstrip("/"), 0.0)
def reachable(self, host: str) -> bool:
"""True when the last ``check()`` got a model list from ``host``."""
return host.rstrip("/") in self._models
def serves(self, host: str, model: str) -> bool:
names = self._models.get(host.rstrip("/"), set())
return model in names or model.split(":")[0] in names

99
tests/cli/test_llm_exercise.py
View File

@@ -46,13 +46,20 @@ class TestIndexComments:
class TestIndexCorpus:
@patch("llm.source.ZoteroPdfIndex.snapshot")
@patch("llm.source.iter_corpus_docs")
@patch("llm.pool.HostPool.from_config")
@patch("llm.index.index_docs")
@patch("conf.connect.bib")
@patch("llm.config.load")
def test_corpus_collection(
self, mock_load, mock_bib, mock_index_docs, mock_from_config, mock_iter
self,
mock_load,
mock_bib,
mock_index_docs,
mock_from_config,
mock_iter,
mock_snapshot,
):
cfg = MagicMock()
mock_load.return_value = cfg
@@ -61,14 +68,54 @@ class TestIndexCorpus:
mock_iter.return_value = iter(["doc1"])
mock_from_config.return_value = MagicMock()
mock_index_docs.return_value = _STATS
zot = MagicMock()
mock_snapshot.return_value = zot
result = runner.invoke(app, ["index", "--collection", "corpus"])
assert result.exit_code == 0
mock_iter.assert_called_once_with(store)
mock_iter.assert_called_once_with(store, zotero=zot)
kwargs = mock_index_docs.call_args.kwargs
assert kwargs["collection"] == "corpus"
assert "indexed=3 skipped=1 chunks=7" in result.output
assert "corpus: indexed=3 skipped=1 chunks=7" in result.output
class TestIndexAll:
@patch("llm.source.ZoteroPdfIndex.snapshot")
@patch("llm.source.iter_corpus_docs")
@patch("llm.source.iter_rule_docs")
@patch("llm.source.iter_comment_docs")
@patch("llm.pool.HostPool.from_config")
@patch("llm.index.index_docs")
@patch("conf.connect.bib")
@patch("llm.config.load")
def test_all_runs_comments_rules_corpus_in_order(
self,
mock_load,
mock_bib,
mock_index_docs,
mock_from_config,
mock_comments,
mock_rules,
mock_corpus,
mock_snapshot,
):
mock_load.return_value = MagicMock()
mock_bib.return_value = MagicMock()
mock_from_config.return_value = MagicMock()
mock_index_docs.return_value = _STATS
for m in (mock_comments, mock_rules, mock_corpus):
m.return_value = iter(["d"])
result = runner.invoke(app, ["index", "--collection", "all"])
assert result.exit_code == 0
assert [c.kwargs["collection"] for c in mock_index_docs.call_args_list] == [
"comments",
"rules",
"corpus",
]
assert result.output.count("indexed=3") == 3
class TestIndexRules:
@@ -110,7 +157,7 @@ class TestIndexBadCollection:
result = runner.invoke(app, ["index", "--collection", "bogus"])
assert result.exit_code == 2
assert "collection must be 'comments', 'corpus' or 'rules'" in result.output
assert "collection must be comments, rules, corpus or all" in result.output
class TestIndexLimit:
@@ -143,3 +190,47 @@ class TestServe:
mock_run.assert_called_once_with(
"llm.api:app", host="0.0.0.0", port=9000, log_level="info"
)
class TestHosts:
@patch("llm.pool.pick_model", return_value="qwen2.5:32b")
@patch("llm.pool.HostPool.status")
@patch("llm.pool.HostPool.check")
@patch("llm.config.load")
def test_lists_fleet_and_generation_pick(
self, mock_load, mock_check, mock_status, _pm
):
from llm.config import LlmConfig
rack, rig = "http://ollama:11434", "http://rig:11434"
mock_load.return_value = LlmConfig(
ollama_hosts=(rack, rig),
host_vram={rack: 12, rig: 24},
embed_model="e",
instruct_model="qwen2.5:14b",
embed_dim=768,
build_ann_index=False,
pg_host="x",
pg_port=5432,
pg_db="llm",
pg_user="llm",
)
mock_check.return_value = [rig]
mock_status.return_value = [
{"host": rig, "vram_gb": 24.0, "models": ["qwen2.5:32b"]}
]
result = runner.invoke(app, ["hosts"])
assert result.exit_code == 0
assert "http://rig:11434" in result.output and "24 GB" in result.output
assert "http://ollama:11434" in result.output and "DOWN" in result.output
assert "generation" in result.output and "qwen2.5:32b @" in result.output
@patch("llm.pool.HostPool.check", side_effect=RuntimeError("no Ollama host"))
@patch("llm.config.load")
def test_no_live_hosts_exits_1(self, mock_load, _check):
mock_load.return_value = MagicMock(ollama_hosts=("http://h:1",), host_vram={})
result = runner.invoke(app, ["hosts"])
assert result.exit_code == 1
assert "no Ollama host" in result.output

2
tests/llm/test_pool.py
View File

@@ -209,6 +209,8 @@ class TestGeneration:
assert pool.serves(H1, "big") is True
assert pool.serves(H1, "big:latest") is True
assert pool.serves(H1, "nope") is False
assert pool.reachable(H1) is True
assert pool.reachable(H2) is False
assert pool.status() == [
{"host": H1, "vram_gb": 0.0, "models": ["big:latest", "chat:latest"]}
]