272d653ef9
F2's live check needs to index just the 6 CPT bib items, not the whole corpus — --key was wired only into the rules collection (iter_rule_refs); the corpus branch (_refs_for -> iter_corpus_refs) ignored it entirely, so "--collection corpus --key ..." silently indexed everything. iter_corpus_refs takes an optional keys tuple (same Python-side post-filter iter_rule_refs already uses) and _refs_for now passes --key through on the corpus path too.
400 lines
12 KiB
Python
400 lines
12 KiB
Python
"""Exercise cli/llm.py's `index` and `serve` commands with mocked backends."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from typer.testing import CliRunner
|
|
|
|
from cli.llm import app
|
|
|
|
runner = CliRunner()
|
|
|
|
_STATS = {
|
|
"indexed": 3,
|
|
"skipped": 1,
|
|
"chunks": 7,
|
|
"fingerprint_skipped": 0,
|
|
"hash_skipped": 0,
|
|
"docket_complete": 0,
|
|
}
|
|
|
|
|
|
class TestIndexComments:
|
|
@patch("llm.index.docket_complete", return_value={})
|
|
@patch("llm.index._engine")
|
|
@patch("llm.source.iter_comment_refs")
|
|
@patch("llm.pool.HostPool.from_config")
|
|
@patch("llm.index.index_refs")
|
|
@patch("conf.connect.bib")
|
|
@patch("llm.config.load")
|
|
def test_default_collection_is_comments(
|
|
self,
|
|
mock_load,
|
|
mock_bib,
|
|
mock_index_refs,
|
|
mock_from_config,
|
|
mock_iter,
|
|
_engine,
|
|
_complete,
|
|
):
|
|
cfg = MagicMock()
|
|
mock_load.return_value = cfg
|
|
store = MagicMock()
|
|
store.sealed_dockets.return_value = {}
|
|
mock_bib.return_value = store
|
|
docs = iter(["doc1", "doc2"])
|
|
mock_iter.return_value = docs
|
|
pool = MagicMock()
|
|
mock_from_config.return_value = pool
|
|
mock_index_refs.return_value = _STATS
|
|
|
|
result = runner.invoke(app, ["index"])
|
|
|
|
assert result.exit_code == 0, result.output
|
|
mock_iter.assert_called_once_with(store, docket="", skip_dockets=set())
|
|
mock_index_refs.assert_called_once()
|
|
kwargs = mock_index_refs.call_args.kwargs
|
|
assert kwargs["collection"] == "comments"
|
|
assert kwargs["cfg"] is cfg
|
|
assert kwargs["pool"] is pool
|
|
assert kwargs["force"] is False
|
|
assert "indexed=3 skipped=1 chunks=7" in result.output
|
|
|
|
|
|
class TestIndexCorpus:
|
|
@patch("llm.index.docket_complete", return_value={})
|
|
@patch("llm.index._engine")
|
|
@patch("llm.source.ZoteroPdfIndex.lazy")
|
|
@patch("llm.source.iter_corpus_refs")
|
|
@patch("llm.pool.HostPool.from_config")
|
|
@patch("llm.index.index_refs")
|
|
@patch("conf.connect.bib")
|
|
@patch("llm.config.load")
|
|
def test_corpus_collection(
|
|
self,
|
|
mock_load,
|
|
mock_bib,
|
|
mock_index_refs,
|
|
mock_from_config,
|
|
mock_iter,
|
|
mock_lazy,
|
|
_engine,
|
|
_complete,
|
|
):
|
|
cfg = MagicMock()
|
|
mock_load.return_value = cfg
|
|
store = MagicMock()
|
|
store.sealed_dockets.return_value = {}
|
|
mock_bib.return_value = store
|
|
mock_iter.return_value = iter(["doc1"])
|
|
mock_from_config.return_value = MagicMock()
|
|
mock_index_refs.return_value = _STATS
|
|
zot = MagicMock()
|
|
mock_lazy.return_value = zot
|
|
|
|
result = runner.invoke(app, ["index", "--collection", "corpus"])
|
|
|
|
assert result.exit_code == 0, result.output
|
|
mock_iter.assert_called_once_with(store, keys=(), zotero=zot)
|
|
kwargs = mock_index_refs.call_args.kwargs
|
|
assert kwargs["collection"] == "corpus"
|
|
assert "corpus: indexed=3 skipped=1 chunks=7" in result.output
|
|
|
|
@patch("llm.index.docket_complete", return_value={})
|
|
@patch("llm.index._engine")
|
|
@patch("llm.source.ZoteroPdfIndex.lazy")
|
|
@patch("llm.source.iter_corpus_refs")
|
|
@patch("llm.pool.HostPool.from_config")
|
|
@patch("llm.index.index_refs")
|
|
@patch("conf.connect.bib")
|
|
@patch("llm.config.load")
|
|
def test_key_option_limits_corpus_to_those_items(
|
|
self,
|
|
mock_load,
|
|
mock_bib,
|
|
mock_index_refs,
|
|
mock_from_config,
|
|
mock_iter,
|
|
mock_lazy,
|
|
_engine,
|
|
_complete,
|
|
):
|
|
mock_load.return_value = MagicMock()
|
|
store = MagicMock()
|
|
store.sealed_dockets.return_value = {}
|
|
mock_bib.return_value = store
|
|
mock_iter.return_value = iter(["doc1"])
|
|
mock_from_config.return_value = MagicMock()
|
|
mock_index_refs.return_value = _STATS
|
|
zot = MagicMock()
|
|
mock_lazy.return_value = zot
|
|
|
|
result = runner.invoke(
|
|
app,
|
|
[
|
|
"index",
|
|
"--collection",
|
|
"corpus",
|
|
"--key",
|
|
"GQGTPGYV",
|
|
"--key",
|
|
"T6H4ZPEQ",
|
|
],
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
mock_iter.assert_called_once_with(
|
|
store, keys=("GQGTPGYV", "T6H4ZPEQ"), zotero=zot
|
|
)
|
|
|
|
|
|
class TestIndexAll:
|
|
@patch("llm.index.docket_complete", return_value={})
|
|
@patch("llm.index._engine")
|
|
@patch("llm.source.ZoteroPdfIndex.lazy")
|
|
@patch("llm.source.iter_corpus_refs")
|
|
@patch("llm.source.iter_rule_refs")
|
|
@patch("llm.source.iter_comment_refs")
|
|
@patch("llm.pool.HostPool.from_config")
|
|
@patch("llm.index.index_refs")
|
|
@patch("conf.connect.bib")
|
|
@patch("llm.config.load")
|
|
def test_all_runs_comments_rules_corpus_in_order(
|
|
self,
|
|
mock_load,
|
|
mock_bib,
|
|
mock_index_refs,
|
|
mock_from_config,
|
|
mock_comments,
|
|
mock_rules,
|
|
mock_corpus,
|
|
mock_lazy,
|
|
_engine,
|
|
_complete,
|
|
):
|
|
mock_load.return_value = MagicMock()
|
|
store = MagicMock()
|
|
store.sealed_dockets.return_value = {}
|
|
mock_bib.return_value = store
|
|
mock_from_config.return_value = MagicMock()
|
|
mock_index_refs.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, result.output
|
|
assert [c.kwargs["collection"] for c in mock_index_refs.call_args_list] == [
|
|
"comments",
|
|
"rules",
|
|
"corpus",
|
|
]
|
|
assert result.output.count("indexed=3") == 3
|
|
|
|
|
|
class TestIndexRules:
|
|
@patch("llm.index.docket_complete", return_value={})
|
|
@patch("llm.index._engine")
|
|
@patch("llm.source.iter_rule_refs")
|
|
@patch("llm.pool.HostPool.from_config")
|
|
@patch("llm.index.index_refs")
|
|
@patch("conf.connect.bib")
|
|
@patch("llm.config.load")
|
|
def test_rules_collection(
|
|
self,
|
|
mock_load,
|
|
mock_bib,
|
|
mock_index_refs,
|
|
mock_from_config,
|
|
mock_iter,
|
|
_engine,
|
|
_complete,
|
|
):
|
|
cfg = MagicMock()
|
|
mock_load.return_value = cfg
|
|
store = MagicMock()
|
|
store.sealed_dockets.return_value = {}
|
|
mock_bib.return_value = store
|
|
mock_iter.return_value = iter(["doc1"])
|
|
mock_from_config.return_value = MagicMock()
|
|
mock_index_refs.return_value = _STATS
|
|
|
|
result = runner.invoke(
|
|
app,
|
|
["index", "--collection", "rules", "--key", "abc123", "--key", "def456"],
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
mock_iter.assert_called_once_with(store, keys=("abc123", "def456"))
|
|
kwargs = mock_index_refs.call_args.kwargs
|
|
assert kwargs["collection"] == "rules"
|
|
assert "indexed=3 skipped=1 chunks=7" in result.output
|
|
|
|
|
|
class TestIndexBadCollection:
|
|
@patch("conf.connect.bib")
|
|
@patch("llm.config.load")
|
|
def test_bad_collection_raises_bad_parameter(self, mock_load, mock_bib):
|
|
mock_load.return_value = MagicMock()
|
|
mock_bib.return_value = MagicMock()
|
|
|
|
result = runner.invoke(app, ["index", "--collection", "bogus"])
|
|
|
|
assert result.exit_code == 2
|
|
assert "collection must be comments, rules, corpus or all" in result.output
|
|
|
|
|
|
class TestIndexLimit:
|
|
@patch("llm.index.docket_complete", return_value={})
|
|
@patch("llm.index._engine")
|
|
@patch("llm.source.iter_comment_refs")
|
|
@patch("llm.pool.HostPool.from_config")
|
|
@patch("llm.index.index_refs")
|
|
@patch("conf.connect.bib")
|
|
@patch("llm.config.load")
|
|
def test_limit_truncates_docs(
|
|
self,
|
|
mock_load,
|
|
mock_bib,
|
|
mock_index_refs,
|
|
mock_from_config,
|
|
mock_iter,
|
|
_engine,
|
|
_complete,
|
|
):
|
|
mock_load.return_value = MagicMock()
|
|
store = MagicMock()
|
|
store.sealed_dockets.return_value = {}
|
|
mock_bib.return_value = store
|
|
mock_iter.return_value = iter([f"doc{i}" for i in range(5)])
|
|
mock_from_config.return_value = MagicMock()
|
|
mock_index_refs.return_value = _STATS
|
|
|
|
result = runner.invoke(app, ["index", "--limit", "2"])
|
|
|
|
assert result.exit_code == 0, result.output
|
|
refs_arg = mock_index_refs.call_args.args[0]
|
|
assert list(refs_arg) == ["doc0", "doc1"]
|
|
assert mock_index_refs.call_args.kwargs["mark_complete"] is False
|
|
|
|
|
|
class TestRestamp:
|
|
"""F6: `stack llm restamp` must never rebuild the HNSW index — only
|
|
the cheap metadata GIN indexes, and not even those under --dry-run."""
|
|
|
|
@patch("os.path.exists", return_value=False)
|
|
@patch("pfs.anchors.code_family_index", return_value={})
|
|
@patch("llm.restamp.restamp")
|
|
@patch("llm.migrate.ensure_metadata_indexes")
|
|
@patch("llm.migrate.migrate")
|
|
@patch("llm.index._engine")
|
|
@patch("llm.config.load")
|
|
def test_dry_run_issues_no_ddl(
|
|
self,
|
|
mock_load,
|
|
mock_engine,
|
|
mock_migrate,
|
|
mock_ensure_metadata,
|
|
mock_restamp,
|
|
mock_code_index,
|
|
mock_exists,
|
|
):
|
|
mock_load.return_value = MagicMock(duckdb_replica="nope.duckdb")
|
|
mock_restamp.return_value = {
|
|
"collection": "rules",
|
|
"scanned": 0,
|
|
"updated": 0,
|
|
"seconds": 0.0,
|
|
}
|
|
|
|
result = runner.invoke(app, ["restamp", "--collection", "rules", "--dry-run"])
|
|
|
|
assert result.exit_code == 0, result.output
|
|
mock_ensure_metadata.assert_not_called()
|
|
assert mock_restamp.call_args.kwargs["dry_run"] is True
|
|
|
|
@patch("os.path.exists", return_value=False)
|
|
@patch("pfs.anchors.code_family_index", return_value={})
|
|
@patch("llm.restamp.restamp")
|
|
@patch("llm.migrate.ensure_metadata_indexes")
|
|
@patch("llm.migrate.migrate")
|
|
@patch("llm.index._engine")
|
|
@patch("llm.config.load")
|
|
def test_real_run_calls_ensure_metadata_indexes_not_hnsw(
|
|
self,
|
|
mock_load,
|
|
mock_engine,
|
|
mock_migrate,
|
|
mock_ensure_metadata,
|
|
mock_restamp,
|
|
mock_code_index,
|
|
mock_exists,
|
|
):
|
|
mock_load.return_value = MagicMock(duckdb_replica="nope.duckdb")
|
|
mock_restamp.return_value = {
|
|
"collection": "rules",
|
|
"scanned": 0,
|
|
"updated": 0,
|
|
"seconds": 0.0,
|
|
}
|
|
|
|
result = runner.invoke(app, ["restamp", "--collection", "rules"])
|
|
|
|
assert result.exit_code == 0, result.output
|
|
mock_ensure_metadata.assert_called_once()
|
|
|
|
|
|
class TestServe:
|
|
@patch("uvicorn.run")
|
|
def test_serve_starts_uvicorn(self, mock_run):
|
|
result = runner.invoke(app, ["serve", "--host", "0.0.0.0", "--port", "9000"])
|
|
assert result.exit_code == 0
|
|
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
|