Tasks 12 + 13: worker node agent + Phi-4-mini query triage
Task 12 — node_agent/agent.py: standalone AMQP worker agent config reader, model discovery, registration publisher, ping/pong handler, model swap with systemctl + health poll (14 tests) Task 13 — triage.py: query classification + cluster node selection classify_query() routes to Phi-4-mini at :8083 select_node() picks best worker by model affinity get_inference_url() replaces hardcoded LLAMA_SERVER_BASE cluster.py stores node ip for URL construction (6 tests + 5 existing chat tests mocked for triage) 168 tests passing (+20 new)
This commit is contained in:
@@ -9,9 +9,16 @@ import app
|
||||
import config
|
||||
import db
|
||||
import routers.chat
|
||||
import triage
|
||||
from security import SESSIONS, PIN_ATTEMPTS, RATE_EVENTS
|
||||
|
||||
|
||||
def _mock_triage_url(monkeypatch, url: str = config.LLAMA_SERVER_BASE):
|
||||
async def fake_url(query: str) -> str:
|
||||
return url
|
||||
monkeypatch.setattr(routers.chat, "_get_inference_url", fake_url)
|
||||
|
||||
|
||||
def make_client(tmp_path: Path) -> TestClient:
|
||||
os.environ["CAIC_ADMIN_PIN"] = "1234"
|
||||
db.DB_PATH = tmp_path / "caic-streaming.db"
|
||||
@@ -53,6 +60,7 @@ def _stream_json_lines(events: list[dict]) -> list[str]:
|
||||
|
||||
|
||||
def test_chat_stream_emits_tokens_and_done(tmp_path: Path, monkeypatch):
|
||||
_mock_triage_url(monkeypatch)
|
||||
with make_client(tmp_path) as client:
|
||||
sid = client.post("/api/auth/guest", headers={"Origin": "http://testserver"}).json()[
|
||||
"session_id"
|
||||
@@ -88,6 +96,7 @@ def test_chat_stream_emits_tokens_and_done(tmp_path: Path, monkeypatch):
|
||||
|
||||
|
||||
def test_chat_auto_search_trigger_emits_search_events(tmp_path: Path, monkeypatch):
|
||||
_mock_triage_url(monkeypatch)
|
||||
with make_client(tmp_path) as client:
|
||||
sid = client.post("/api/auth/guest", headers={"Origin": "http://testserver"}).json()[
|
||||
"session_id"
|
||||
@@ -142,6 +151,7 @@ def test_chat_auto_search_trigger_emits_search_events(tmp_path: Path, monkeypatc
|
||||
|
||||
|
||||
def test_chat_with_upload_context_id_injects_document(tmp_path: Path, monkeypatch):
|
||||
_mock_triage_url(monkeypatch)
|
||||
captured_payload = {}
|
||||
|
||||
def stream_stub(self, method, url, json=None, timeout=None):
|
||||
@@ -173,6 +183,7 @@ def test_chat_with_upload_context_id_injects_document(tmp_path: Path, monkeypatc
|
||||
|
||||
|
||||
def test_chat_with_expired_upload_context_id_silent(tmp_path: Path, monkeypatch):
|
||||
_mock_triage_url(monkeypatch)
|
||||
captured_payload = {}
|
||||
|
||||
def stream_stub(self, method, url, json=None, timeout=None):
|
||||
@@ -205,6 +216,7 @@ def test_chat_with_expired_upload_context_id_silent(tmp_path: Path, monkeypatch)
|
||||
|
||||
|
||||
def test_memory_command_paths_remember_and_forget(tmp_path: Path, monkeypatch):
|
||||
_mock_triage_url(monkeypatch)
|
||||
with make_client(tmp_path) as client:
|
||||
sid = client.post("/api/auth/guest", headers={"Origin": "http://testserver"}).json()[
|
||||
"session_id"
|
||||
|
||||
@@ -0,0 +1,340 @@
|
||||
"""Tests for node_agent/agent.py — standalone worker agent."""
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
|
||||
import node_agent.agent as agent
|
||||
from node_agent.agent import (
|
||||
AgentConfig,
|
||||
ModelInfo,
|
||||
build_registration_payload,
|
||||
discover_models,
|
||||
get_load,
|
||||
handle_ping,
|
||||
handle_swap_model,
|
||||
)
|
||||
|
||||
|
||||
class FakeMsg:
|
||||
def __init__(self, body_dict: dict):
|
||||
self.body = json.dumps(body_dict).encode()
|
||||
|
||||
@asynccontextmanager
|
||||
async def process(self):
|
||||
yield
|
||||
|
||||
|
||||
class FakeExchange:
|
||||
def __init__(self, name=""):
|
||||
self.name = name
|
||||
self.published = []
|
||||
|
||||
async def publish(self, msg, routing_key):
|
||||
self.published.append((msg, routing_key))
|
||||
|
||||
|
||||
class FakeChannel:
|
||||
def __init__(self):
|
||||
self.exchanges = {}
|
||||
self.is_closed = False
|
||||
|
||||
async def get_exchange(self, name):
|
||||
if name not in self.exchanges:
|
||||
self.exchanges[name] = FakeExchange(name)
|
||||
return self.exchanges[name]
|
||||
|
||||
async def declare_exchange(self, name, typ, durable=True):
|
||||
self.exchanges[name] = FakeExchange(name)
|
||||
return self.exchanges[name]
|
||||
|
||||
async def declare_queue(self, name="", exclusive=True):
|
||||
return self
|
||||
|
||||
async def bind(self, exchange, routing_key):
|
||||
pass
|
||||
|
||||
async def close(self):
|
||||
self.is_closed = True
|
||||
|
||||
|
||||
# ── 1. Registration payload shape ────────────────────────────────────
|
||||
|
||||
|
||||
def test_registration_payload_shape():
|
||||
cfg = AgentConfig()
|
||||
cfg.node_name = "worker01"
|
||||
cfg.node_ip = "192.168.50.210"
|
||||
cfg.capabilities = ["llm"]
|
||||
cfg.active_model = "llama3.1-latest-Q4_K_M.gguf"
|
||||
cfg.llama_port = 8081
|
||||
|
||||
inventory = [{"filename": "llama3.1-latest-Q4_K_M.gguf", "name": "llama3.1", "quant": "Q4_K_M"}]
|
||||
payload = build_registration_payload(cfg, inventory)
|
||||
|
||||
assert payload["node_name"] == "worker01"
|
||||
assert payload["node_type"] == "worker"
|
||||
assert payload["ip"] == "192.168.50.210"
|
||||
assert payload["capabilities"] == ["llm"]
|
||||
assert "active_model" in payload
|
||||
assert payload["active_model"]["filename"] == "llama3.1-latest-Q4_K_M.gguf"
|
||||
assert payload["active_model"]["port"] == 8081
|
||||
assert len(payload["inventory"]) == 1
|
||||
|
||||
|
||||
def test_registration_payload_active_model_none():
|
||||
cfg = AgentConfig()
|
||||
cfg.active_model = ""
|
||||
payload = build_registration_payload(cfg, [])
|
||||
assert payload["active_model"] is None
|
||||
|
||||
|
||||
# ── 2. Model discovery ───────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_discover_models(tmp_path):
|
||||
models_dir = tmp_path / "models"
|
||||
models_dir.mkdir()
|
||||
# Create valid model files
|
||||
(models_dir / "llama3.1-latest-Q4_K_M.gguf").write_text("")
|
||||
(models_dir / "mistral-nemo-7b-Q6_K_L.gguf").write_text("")
|
||||
# Create file that doesn't match pattern
|
||||
(models_dir / "readme.txt").write_text("")
|
||||
# Create file with unrecognized naming
|
||||
(models_dir / "my_custom_model.gguf").write_text("")
|
||||
|
||||
result = discover_models(str(models_dir))
|
||||
assert len(result) == 2
|
||||
names = {m["name"] for m in result}
|
||||
assert "llama3.1" in names
|
||||
assert "mistral-nemo" in names
|
||||
|
||||
|
||||
def test_discover_models_no_directory(tmp_path):
|
||||
result = discover_models(str(tmp_path / "nonexistent"))
|
||||
assert result == []
|
||||
|
||||
|
||||
# ── 3. Config reading ────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_config_from_ini(tmp_path):
|
||||
ini = tmp_path / "caic-node-agent.conf"
|
||||
ini.write_text(
|
||||
"[agent]\n"
|
||||
"node_name = testnode\n"
|
||||
"node_ip = 10.0.0.5\n"
|
||||
"capabilities = llm,rag\n"
|
||||
"amqp_url = amqp://user:pass@host/vhost\n"
|
||||
"llama_port = 9090\n"
|
||||
"models_dir = /tmp/models\n"
|
||||
"active_model = test.gguf\n"
|
||||
)
|
||||
cfg = AgentConfig.from_ini(str(ini))
|
||||
assert cfg.node_name == "testnode"
|
||||
assert cfg.node_ip == "10.0.0.5"
|
||||
assert cfg.capabilities == ["llm", "rag"]
|
||||
assert cfg.amqp_url == "amqp://user:pass@host/vhost"
|
||||
assert cfg.llama_port == 9090
|
||||
assert cfg.models_dir == "/tmp/models"
|
||||
assert cfg.active_model == "test.gguf"
|
||||
|
||||
|
||||
def test_config_from_ini_missing_uses_defaults(tmp_path):
|
||||
cfg = AgentConfig.from_ini(str(tmp_path / "nonexistent.conf"))
|
||||
assert cfg.node_name != "" # socket.gethostname() returns something
|
||||
assert cfg.node_type == "worker"
|
||||
assert cfg.capabilities == ["llm"]
|
||||
|
||||
|
||||
# ── 4. Ping handler publishes pong ──────────────────────────────────
|
||||
|
||||
|
||||
def test_ping_handler_publishes_pong(monkeypatch):
|
||||
monkeypatch.setattr(agent, "HAS_PSUTIL", False)
|
||||
monkeypatch.setattr(agent, "HAS_AIO_PIKA", True)
|
||||
|
||||
cfg = AgentConfig()
|
||||
cfg.node_name = "testworker"
|
||||
|
||||
channel = FakeChannel()
|
||||
admin_ex = FakeExchange("jc.admin")
|
||||
channel.exchanges["jc.admin"] = admin_ex
|
||||
exchange = admin_ex
|
||||
|
||||
async def run():
|
||||
await handle_ping(cfg, channel, exchange, FakeMsg({
|
||||
"from": "coordinator",
|
||||
"node_name": "testworker",
|
||||
"type": "ping",
|
||||
"correlation_id": "abc-123",
|
||||
"timestamp": "2026-01-01T00:00:00Z",
|
||||
}))
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
assert len(exchange.published) == 1
|
||||
msg, rk = exchange.published[0]
|
||||
assert rk == "node.testworker.pong"
|
||||
payload = json.loads(msg.body)
|
||||
assert payload["type"] == "pong"
|
||||
assert payload["correlation_id"] == "abc-123"
|
||||
assert payload["node_name"] == "testworker"
|
||||
|
||||
|
||||
# ── 5. Model swap success path ──────────────────────────────────────
|
||||
|
||||
|
||||
def test_model_swap_success(monkeypatch):
|
||||
monkeypatch.setattr(agent, "HAS_AIO_PIKA", True)
|
||||
monkeypatch.setattr(agent, "HAS_HTTPX", True)
|
||||
|
||||
cfg = AgentConfig()
|
||||
cfg.node_name = "testworker"
|
||||
cfg.llama_port = 9999
|
||||
|
||||
captured_cmds = []
|
||||
|
||||
def fake_run(cmd, **kwargs):
|
||||
captured_cmds.append(cmd)
|
||||
return subprocess.CompletedProcess(cmd, 0, b"", b"")
|
||||
|
||||
monkeypatch.setattr(subprocess, "run", fake_run)
|
||||
|
||||
# Mock _wait_for_llama to succeed
|
||||
async def fake_wait(*a, **kw):
|
||||
return True
|
||||
monkeypatch.setattr(agent, "_wait_for_llama", fake_wait)
|
||||
|
||||
# Mock _update_config_active_model to no-op
|
||||
monkeypatch.setattr(agent, "_update_config_active_model", lambda c, m: None)
|
||||
|
||||
channel = FakeChannel()
|
||||
admin_ex = FakeExchange("jc.admin")
|
||||
system_ex = FakeExchange("jc.system")
|
||||
channel.exchanges["jc.admin"] = admin_ex
|
||||
channel.exchanges["jc.system"] = system_ex
|
||||
|
||||
asyncio.run(handle_swap_model(cfg, channel, (admin_ex, system_ex), FakeMsg({"model_filename": "new-model-Q4_K_M.gguf"})))
|
||||
|
||||
# Check systemctl calls
|
||||
assert len(captured_cmds) == 2
|
||||
assert captured_cmds[0] == ["systemctl", "stop", "llama-server"]
|
||||
assert captured_cmds[1] == ["systemctl", "start", "llama-server"]
|
||||
|
||||
# Check model_ready published on jc.system
|
||||
assert len(system_ex.published) == 1
|
||||
msg, rk = system_ex.published[0]
|
||||
assert rk == "node.testworker.model_ready"
|
||||
payload = json.loads(msg.body)
|
||||
assert payload["type"] == "model_ready"
|
||||
assert payload["active_model"] == "new-model-Q4_K_M.gguf"
|
||||
|
||||
|
||||
# ── 6. Model swap timeout path ──────────────────────────────────────
|
||||
|
||||
|
||||
def test_model_swap_timeout(monkeypatch):
|
||||
monkeypatch.setattr(agent, "HAS_AIO_PIKA", True)
|
||||
monkeypatch.setattr(agent, "HAS_HTTPX", True)
|
||||
|
||||
cfg = AgentConfig()
|
||||
cfg.node_name = "testworker"
|
||||
|
||||
monkeypatch.setattr(subprocess, "run", lambda cmd, **kw: subprocess.CompletedProcess(cmd, 0, b"", b""))
|
||||
|
||||
# Mock _wait_for_llama to fail
|
||||
async def fake_wait_fail(*a, **kw):
|
||||
return False
|
||||
monkeypatch.setattr(agent, "_wait_for_llama", fake_wait_fail)
|
||||
monkeypatch.setattr(agent, "_update_config_active_model", lambda c, m: None)
|
||||
|
||||
channel = FakeChannel()
|
||||
admin_ex = FakeExchange("jc.admin")
|
||||
system_ex = FakeExchange("jc.system")
|
||||
channel.exchanges["jc.admin"] = admin_ex
|
||||
channel.exchanges["jc.system"] = system_ex
|
||||
|
||||
asyncio.run(handle_swap_model(cfg, channel, (admin_ex, system_ex), FakeMsg({"model_filename": "broken-model-Q4_K_M.gguf"})))
|
||||
|
||||
assert len(system_ex.published) == 1
|
||||
msg, rk = system_ex.published[0]
|
||||
assert rk == "node.testworker.model_failed"
|
||||
payload = json.loads(msg.body)
|
||||
assert payload["type"] == "model_failed"
|
||||
assert "error" in payload
|
||||
|
||||
|
||||
# ── 7. Load reporting ───────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_get_load_with_psutil(monkeypatch):
|
||||
class FakePsutil:
|
||||
@staticmethod
|
||||
def cpu_percent(interval=0.5):
|
||||
return 42.0
|
||||
|
||||
@staticmethod
|
||||
def virtual_memory():
|
||||
class VM:
|
||||
percent = 65.0
|
||||
return VM()
|
||||
|
||||
monkeypatch.setattr(agent, "HAS_PSUTIL", True)
|
||||
monkeypatch.setattr(agent, "psutil", FakePsutil)
|
||||
|
||||
# Mock subprocess to fail (no rocm-smi)
|
||||
def fake_run(*a, **kw):
|
||||
raise FileNotFoundError("rocm-smi not found")
|
||||
monkeypatch.setattr(subprocess, "run", fake_run)
|
||||
|
||||
load = get_load()
|
||||
assert load["cpu_pct"] == 42
|
||||
assert load["ram_pct"] == 65
|
||||
assert "vram_pct" not in load
|
||||
|
||||
|
||||
def test_get_load_without_psutil(monkeypatch):
|
||||
monkeypatch.setattr(agent, "HAS_PSUTIL", False)
|
||||
monkeypatch.setattr(subprocess, "run", lambda *a, **kw: (_ for _ in ()).throw(FileNotFoundError("")))
|
||||
load = get_load()
|
||||
assert "cpu_pct" not in load
|
||||
|
||||
|
||||
# ── 8. Agent idle after admission (no heartbeat) ────────────────────
|
||||
|
||||
|
||||
def test_no_heartbeat_timer():
|
||||
"""Agent has no background heartbeat mechanism. This test asserts that
|
||||
the codebase contains no heartbeat-related logic in the agent itself."""
|
||||
import inspect
|
||||
source = inspect.getsource(agent)
|
||||
# There should be no heartbeat timer in the agent
|
||||
assert "heartbeat" not in source.lower(), \
|
||||
"agent.py should contain no heartbeat logic"
|
||||
|
||||
|
||||
# ── 9. Config writer for model swap ─────────────────────────────────
|
||||
|
||||
|
||||
def test_update_config_active_model(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(agent, "CONFIG_PATH", str(tmp_path / "caic-node-agent.conf"))
|
||||
cfg = AgentConfig()
|
||||
cfg.active_model = "old.gguf"
|
||||
agent._update_config_active_model(cfg, "new.gguf")
|
||||
assert cfg.active_model == "new.gguf"
|
||||
content = (tmp_path / "caic-node-agent.conf").read_text()
|
||||
assert "new.gguf" in content
|
||||
|
||||
|
||||
# ── 10. ModelInfo to_dict ────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_model_info_to_dict():
|
||||
info = ModelInfo(filename="test-Q4_K_M.gguf", name="test", version="latest", quant="Q4_K_M")
|
||||
d = info.to_dict()
|
||||
assert d["name"] == "test"
|
||||
assert d["quant"] == "Q4_K_M"
|
||||
assert d["filename"] == "test-Q4_K_M.gguf"
|
||||
@@ -0,0 +1,133 @@
|
||||
"""Tests for triage.py — query classification and node selection."""
|
||||
import json
|
||||
|
||||
import httpx
|
||||
|
||||
import cluster
|
||||
import config
|
||||
import triage
|
||||
|
||||
|
||||
def _reset():
|
||||
cluster.CLUSTER_NODES.clear()
|
||||
cluster.CLUSTER_COORDINATOR = None
|
||||
|
||||
|
||||
class _MockPostResponse:
|
||||
def __init__(self, json_data: dict, status_code: int = 200):
|
||||
self._json_data = json_data
|
||||
self.status_code = status_code
|
||||
|
||||
def json(self):
|
||||
return self._json_data
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
|
||||
class _MockPostContext:
|
||||
def __init__(self, response: _MockPostResponse):
|
||||
self._response = response
|
||||
|
||||
async def __aenter__(self):
|
||||
return self._response
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
|
||||
# ---------- 1. classify_query returns valid classification ----------
|
||||
|
||||
|
||||
def test_classify_returns_valid(monkeypatch):
|
||||
async def post_stub(self, url, json=None, timeout=None):
|
||||
return _MockPostResponse({
|
||||
"choices": [{"message": {"content": "code"}}]
|
||||
})
|
||||
|
||||
monkeypatch.setattr(httpx.AsyncClient, "post", post_stub)
|
||||
|
||||
result = __import__("asyncio").run(triage.classify_query("write a python function"))
|
||||
assert result == "code"
|
||||
|
||||
|
||||
# ---------- 2. classify_query on error returns "general" ----------
|
||||
|
||||
|
||||
def test_classify_error_returns_general(monkeypatch):
|
||||
async def post_stub(self, url, json=None, timeout=None):
|
||||
raise httpx.ConnectError("connection refused")
|
||||
|
||||
monkeypatch.setattr(httpx.AsyncClient, "post", post_stub)
|
||||
|
||||
result = __import__("asyncio").run(triage.classify_query("any question"))
|
||||
assert result == "general"
|
||||
|
||||
|
||||
# ---------- 3. select_node("code") returns coder node ----------
|
||||
|
||||
|
||||
def test_select_node_code_returns_coder():
|
||||
_reset()
|
||||
cluster.CLUSTER_NODES["coder01"] = {
|
||||
"name": "coder01", "type": "worker", "status": "active",
|
||||
"ip": "192.168.50.210",
|
||||
"active_model": {"name": "qwen2.5-coder-14b", "port": 8082},
|
||||
}
|
||||
cluster.CLUSTER_NODES["general01"] = {
|
||||
"name": "general01", "type": "worker", "status": "active",
|
||||
"ip": "192.168.50.211",
|
||||
"active_model": {"name": "llama3.1", "port": 8081},
|
||||
}
|
||||
|
||||
node = triage.select_node("code")
|
||||
assert node is not None
|
||||
assert node["name"] == "coder01"
|
||||
|
||||
|
||||
# ---------- 4. select_node("general") with no matching node returns None ----------
|
||||
|
||||
|
||||
def test_select_node_general_no_match_returns_none():
|
||||
_reset()
|
||||
cluster.CLUSTER_NODES["coder01"] = {
|
||||
"name": "coder01", "type": "worker", "status": "active",
|
||||
"active_model": {"name": "qwen2.5-coder-14b", "port": 8082},
|
||||
}
|
||||
node = triage.select_node("general")
|
||||
assert node is None
|
||||
|
||||
|
||||
# ---------- 5. get_inference_url with coder node ----------
|
||||
|
||||
|
||||
def test_get_inference_url_with_coder_node(monkeypatch):
|
||||
_reset()
|
||||
async def fake_classify(query: str) -> str:
|
||||
return "code"
|
||||
monkeypatch.setattr(triage, "classify_query", fake_classify)
|
||||
|
||||
cluster.CLUSTER_NODES["coder01"] = {
|
||||
"name": "coder01", "type": "worker", "status": "active",
|
||||
"ip": "192.168.50.210",
|
||||
"active_model": {"name": "qwen2.5-coder-14b", "port": 8082},
|
||||
}
|
||||
|
||||
url = __import__("asyncio").run(triage.get_inference_url("write a loop in rust"))
|
||||
assert url == "http://192.168.50.210:8082/v1"
|
||||
|
||||
|
||||
# ---------- 6. get_inference_url with no nodes returns LLAMA_SERVER_BASE ----------
|
||||
|
||||
|
||||
def test_get_inference_url_no_nodes(monkeypatch):
|
||||
_reset()
|
||||
async def fake_classify(query: str) -> str:
|
||||
return "code"
|
||||
monkeypatch.setattr(triage, "classify_query", fake_classify)
|
||||
|
||||
url = __import__("asyncio").run(triage.get_inference_url("any question"))
|
||||
assert url == config.LLAMA_SERVER_BASE
|
||||
Reference in New Issue
Block a user