diff --git a/AGENTS.md b/AGENTS.md index d16e8c3..cb70033 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -28,6 +28,8 @@ Every router has a dedicated test file: | `test_profile.py` | `routers/profile.py` — get, update, default, length validation | | `test_search_route.py` | `routers/search_route.py` — explicit search flow, no results, errors | | `test_search_url_sanitization.py` | `search.py` URL sanitizer | +| `test_node_agent.py` | `node_agent/agent.py` — registration, ping/pong, model swap | +| `test_triage.py` | `triage.py` — classify_query, select_node, get_inference_url | | `test_settings_allowlist.py` | `routers/settings.py` — allowlisted key enforcement | | `test_skills_framework.py` | `routers/skills.py` — list, toggle, unknown skill, prompt injection | | `test_ip_allowlist.py` | IP allowlist helper + middleware | @@ -55,6 +57,8 @@ Refactored from single-file (`app.py`) into modules under project root: | `search.py` | SearXNG integration, perplexity scoring, refusal detection | | `rag.py` | Qdrant vector search + system prompt assembly + chunk_text() helper | | `gpu.py` | AMD GPU stats via `rocm-smi` | +| `triage.py` | Phi-4-mini-based query classification + cluster node selection | +| `node_agent/` | Standalone worker agent — AMQP client for registration, ping/pong, model swap | | `routers/` | One module per endpoint group (chat, search, skills, completions, upload, ingest) | ### Entrypoint / API keys @@ -66,7 +70,7 @@ Refactored from single-file (`app.py`) into modules under project root: ### Key flows -1. **`/api/chat`** → `process_remember_command()` intercepts "remember that..." / "forget about..." first → optional `upload_context_id` fetches document text from SQLite → `build_system_prompt()` (profile + FTS5 memory + Qdrant RAG + preset + skills + uploaded doc) → stream from llama-server with `logprobs: true` → if perplexity > 15.0 OR `REFUSAL_PATTERNS` match, re-query with SearXNG results +1. **`/api/chat`** → `process_remember_command()` intercepts "remember that..." / "forget about..." first → optional `upload_context_id` fetches document text from SQLite → `build_system_prompt()` (profile + FTS5 memory + Qdrant RAG + preset + skills + uploaded doc) → triage classifies query (general/code/search/rag) → `select_node()` picks best worker → stream from chosen node with `logprobs: true` → if perplexity > 15.0 OR `REFUSAL_PATTERNS` match, re-query with SearXNG results 2. **`/api/search`** → bypasses perplexity/refusal, queries SearXNG directly → summarizes via llama-server 3. **`/v1/chat/completions`** → OpenAI-compatible for Continue.dev/IDE integration; FIM requests proxied without persistence 4. **`/api/upload`** → multipart file upload, PDF/text extraction, `mode=(context|ingest|both)`, stores SQLite context (1hr expiry) + Qdrant upsert @@ -97,6 +101,7 @@ The upstream request includes `"logprobs": true`. `parse_llama_stream_chunk()` e | Service | Required | Port | |---------|----------|------| | llama-server (coordinator) | Yes | 8081 + RPC :50052 (worker GPU) | +| Phi-4-mini (triage) | No | 8083 | | SearXNG | No | 8888 | | wttr.in | No | weather shortcut | | rocm-smi | No | AMD GPU stats | diff --git a/README.md b/README.md index f4ad12a..2c7edc7 100644 --- a/README.md +++ b/README.md @@ -95,7 +95,11 @@ Developer wiki: [docs/wiki/Home.md](docs/wiki/Home.md) │ └── logo.png # Logo image (optional) ├── templates/ │ └── index.html # Frontend -└── tests/ # 148 pytest tests +├── node_agent/ +│ ├── agent.py # Standalone worker agent (AMQP client) +│ └── requirements.txt +├── triage.py # Query classification + cluster node selection +└── tests/ # 168 pytest tests ``` ## Requirements @@ -309,7 +313,7 @@ Settings are stored in the `settings` table and include: python3 -m pytest tests/ -v ``` -All 148 tests use `tmp_path` fixtures + monkeypatched `httpx.AsyncClient`/`aio-pika`. No external services needed. +All 168 tests use `tmp_path` fixtures + monkeypatched `httpx.AsyncClient`/`aio-pika`. No external services needed. ## License diff --git a/TASKS.md b/TASKS.md index d417bb3..61eb09d 100644 --- a/TASKS.md +++ b/TASKS.md @@ -667,7 +667,7 @@ Run full test suite. All existing tests must continue to pass. --- -## TASK 12 — Roadmap N4: Worker Node Registration Publisher (Worker Side) +## ~~TASK 12 — Roadmap N4: Worker Node Registration Publisher (Worker Side) [DONE]~~ This task creates the worker node AMQP client that runs on worker (192.168.50.210). It is a standalone Python script — not part of the jC FastAPI app — that runs as a systemd service on worker. @@ -757,7 +757,7 @@ Run full test suite. All existing tests must continue to pass. --- -## TASK 13 — Roadmap N5: Query Routing via AMQP + Phi-4-mini Triage +## ~~TASK 13 — Roadmap N5: Query Routing via AMQP + Phi-4-mini Triage [DONE]~~ This task wires the cluster into jC's chat flow. When a query arrives at `/api/chat`, instead of always routing to the hardcoded `LLAMA_SERVER_BASE`, jC now routes to the best available cluster node based on query context. diff --git a/cluster.py b/cluster.py index deac9a8..1ef3df5 100644 --- a/cluster.py +++ b/cluster.py @@ -64,6 +64,7 @@ async def handle_registration(exchange: str, routing_key: str, payload: dict) -> "name": node_name, "type": node_type, "status": "active", + "ip": payload.get("ip"), "capabilities": payload.get("capabilities", []), "active_model": payload.get("active_model"), "load": payload.get("load"), diff --git a/config.py b/config.py index aff95d0..4e34d2d 100644 --- a/config.py +++ b/config.py @@ -9,7 +9,7 @@ import logging log = logging.getLogger("caic") -VERSION = "v0.14.0" +VERSION = "v0.15.0" OLLAMA_BASE = os.environ.get("OLLAMA_BASE", "http://localhost:11434") LLAMA_SERVER_BASE = os.environ.get("LLAMA_SERVER_BASE", "http://192.168.50.108:8081") SEARXNG_BASE = "http://localhost:8888" @@ -102,6 +102,11 @@ ALLOWED_SETTINGS_KEYS = { "skills_enabled", } +# --- Triage / query routing --- +TRIAGE_BASE = os.environ.get("CAIC_TRIAGE_BASE", "http://127.0.0.1:8083/v1") +TRIAGE_TIMEOUT = 10 +FALLBACK_TO_DEFAULT = True + # --- Perplexity --- PERPLEXITY_THRESHOLD = 15.0 diff --git a/node_agent/agent.py b/node_agent/agent.py new file mode 100644 index 0000000..20ce706 --- /dev/null +++ b/node_agent/agent.py @@ -0,0 +1,412 @@ +""" +cAIc — Worker node agent. + +Standalone AMQP client that registers with the cAIc coordinator, +responds to pings, and handles model swap commands. + +## Config file: /etc/caic-node-agent.conf + +```ini +[agent] +# hostname — defaults to socket.gethostname() +node_name = jarvis +# LAN IP — defaults from socket +node_ip = 192.168.50.210 +# "worker" (fixed) +node_type = worker +# comma-separated capability list +capabilities = llm +# RabbitMQ URL on coordinator +amqp_url = amqp://caic:password@192.168.50.108:5672/caic +# port llama-server listens on +llama_port = 8081 +# path to GGUF model files +models_dir = /var/lib/caic/models +# currently active model filename +active_model = llama3.1-latest-Q4_K_M.gguf +``` + +## systemd unit: /etc/systemd/system/caic-node-agent.service + +```ini +[Unit] +Description=cAIc Worker Node Agent +After=network.target rabbitmq.service +Wants=rabbitmq.service + +[Service] +Type=simple +User=root +Group=root +WorkingDirectory=/opt/caic +ExecStart=/usr/bin/python3 /opt/caic/node_agent/agent.py +Restart=always +RestartSec=5 + +[Install] +WantedBy=multi-user.target +``` +""" +import asyncio +import json +import logging +import os +import socket +import subprocess +import sys +import time +from configparser import ConfigParser +from datetime import datetime, timezone +from pathlib import Path +from typing import Optional + +try: + import psutil + HAS_PSUTIL = True +except ImportError: + HAS_PSUTIL = False + +try: + import aio_pika + from aio_pika import DeliveryMode, ExchangeType + HAS_AIO_PIKA = True +except ImportError: + HAS_AIO_PIKA = False + +try: + import httpx + HAS_HTTPX = True +except ImportError: + HAS_HTTPX = False + +log = logging.getLogger("caic") +CONFIG_PATH = "/etc/caic-node-agent.conf" + +# ── data types ────────────────────────────────────────────────────────── + +class AgentConfig: + def __init__(self): + self.node_name: str = socket.gethostname() + self.node_ip: str = "127.0.0.1" + self.node_type: str = "worker" + self.capabilities: list[str] = ["llm"] + self.amqp_url: str = "amqp://caic:password@localhost:5672/caic" + self.llama_port: int = 8081 + self.models_dir: str = "/var/lib/caic/models" + self.active_model: str = "" + + @classmethod + def from_ini(cls, path: str = CONFIG_PATH) -> "AgentConfig": + cfg = cls() + parser = ConfigParser() + if not os.path.exists(path): + log.warning("config %s not found, using defaults", path) + return cfg + parser.read(path) + sec = "agent" + if parser.has_section(sec): + cfg.node_name = parser.get(sec, "node_name", fallback=cfg.node_name) + cfg.node_ip = parser.get(sec, "node_ip", fallback=cfg.node_ip) + cfg.node_type = parser.get(sec, "node_type", fallback=cfg.node_type) + raw_caps = parser.get(sec, "capabilities", fallback="llm") + cfg.capabilities = [c.strip() for c in raw_caps.split(",") if c.strip()] + cfg.amqp_url = parser.get(sec, "amqp_url", fallback=cfg.amqp_url) + cfg.llama_port = parser.getint(sec, "llama_port", fallback=cfg.llama_port) + cfg.models_dir = parser.get(sec, "models_dir", fallback=cfg.models_dir) + cfg.active_model = parser.get(sec, "active_model", fallback=cfg.active_model) + return cfg + + +class ModelInfo: + def __init__(self, filename: str, name: str = "", version: str = "", quant: str = ""): + self.filename = filename + self.name = name + self.version = version + self.quant = quant + self.path = "" + + def to_dict(self) -> dict: + return { + "name": self.name, + "version": self.version, + "quant": self.quant, + "filename": self.filename, + } + + +# ── model discovery ───────────────────────────────────────────────────── + +_MODEL_PATTERN = None # lazy compile + + +def discover_models(models_dir: str) -> list[dict]: + import re + global _MODEL_PATTERN + if _MODEL_PATTERN is None: + _MODEL_PATTERN = re.compile( + r"^(?P.+?)-(?P[^-]+)-(?PQ\d+_K_[A-Z]+|IQ\d_[A-Z]+|fp\d+)\.gguf$" + ) + root = Path(models_dir) + if not root.is_dir(): + log.warning("models_dir %s not found", models_dir) + return [] + + results = [] + for fpath in sorted(root.glob("*.gguf")): + m = _MODEL_PATTERN.match(fpath.name) + if m: + info = ModelInfo( + filename=fpath.name, + name=m.group("name"), + version=m.group("version"), + quant=m.group("quant"), + ) + info.path = str(fpath) + results.append(info.to_dict()) + else: + log.debug("skipping unrecognized model filename: %s", fpath.name) + return results + + +# ── load reporting ────────────────────────────────────────────────────── + +def get_load() -> dict: + load = {} + if HAS_PSUTIL: + load["cpu_pct"] = round(psutil.cpu_percent(interval=0.5)) + load["ram_pct"] = round(psutil.virtual_memory().percent) + try: + result = subprocess.run( + ["rocm-smi", "--showmeminfo", "vram"], + capture_output=True, text=True, timeout=3, + ) + if result.returncode == 0: + for line in result.stdout.splitlines(): + if "VRAM Total" in line: + parts = line.split() + if len(parts) >= 3: + total = int(parts[-1]) + elif "VRAM Used" in line: + parts = line.split() + if len(parts) >= 3: + used = int(parts[-1]) + if total and total > 0: + load["vram_pct"] = round(used / total * 100) + except (FileNotFoundError, subprocess.TimeoutExpired): + pass + return load + + +# ── AMQP helpers ──────────────────────────────────────────────────────── + +async def declare_exchanges(channel) -> tuple: + admin = await channel.declare_exchange("jc.admin", ExchangeType.TOPIC, durable=True) + system = await channel.declare_exchange("jc.system", ExchangeType.TOPIC, durable=True) + return admin, system + + +async def publish(channel, exchange, routing_key, payload): + body = json.dumps(payload).encode() + msg = aio_pika.Message(body, delivery_mode=DeliveryMode.PERSISTENT) + await exchange.publish(msg, routing_key) + + +# ── registration ──────────────────────────────────────────────────────── + +def build_registration_payload(cfg: AgentConfig, inventory: list[dict]) -> dict: + model_dict = None + if cfg.active_model: + for inv in inventory: + if inv["filename"] == cfg.active_model: + model_dict = {**inv, "port": cfg.llama_port} + break + return { + "node_name": cfg.node_name, + "node_type": cfg.node_type, + "ip": cfg.node_ip, + "capabilities": cfg.capabilities, + "active_model": model_dict, + "inventory": inventory, + } + + +# ── ping / pong ───────────────────────────────────────────────────────── + +async def handle_ping(cfg: AgentConfig, channel, exchange, msg: aio_pika.IncomingMessage): + async with msg.process(): + try: + payload = json.loads(msg.body.decode()) + except json.JSONDecodeError: + return + correlation_id = payload.get("correlation_id") + if not correlation_id: + return + + load = get_load() + now = datetime.now(timezone.utc).isoformat() + "Z" + pong = { + "node_name": cfg.node_name, + "type": "pong", + "correlation_id": correlation_id, + "status": "active", + "active_model": None, # simplified; could read current + "load": load, + "timestamp": now, + } + await publish(channel, exchange, f"node.{cfg.node_name}.pong", pong) + + +# ── model swap ────────────────────────────────────────────────────────── + +async def handle_swap_model(cfg: AgentConfig, channel, exchanges, msg: aio_pika.IncomingMessage): + admin_ex, system_ex = exchanges + async with msg.process(): + try: + payload = json.loads(msg.body.decode()) + except json.JSONDecodeError: + return + + model_filename = payload.get("model_filename") + if not model_filename: + log.error("swap_model missing model_filename") + return + + log.info("swapping model to %s", model_filename) + + # 1. Stop current llama-server + log.info("stopping llama-server") + subprocess.run(["systemctl", "stop", "llama-server"], check=False) + + # 2. Update config + _update_config_active_model(cfg, model_filename) + + # 3. Start llama-server + log.info("starting llama-server") + subprocess.run(["systemctl", "start", "llama-server"], check=False) + + # 4. Poll health endpoint + healthy = await _wait_for_llama(cfg.llama_port, timeout=120, interval=2) + + now = datetime.now(timezone.utc).isoformat() + "Z" + if healthy: + result_payload = { + "node_name": cfg.node_name, + "type": "model_ready", + "active_model": model_filename, + "port": cfg.llama_port, + "timestamp": now, + } + log.info("model swap successful: %s", model_filename) + else: + result_payload = { + "node_name": cfg.node_name, + "type": "model_failed", + "active_model": model_filename, + "port": cfg.llama_port, + "error": "llama-server did not become healthy within 120s", + "timestamp": now, + } + log.error("model swap failed: %s", model_filename) + + await publish(channel, system_ex, f"node.{cfg.node_name}.{result_payload['type']}", result_payload) + + +def _update_config_active_model(cfg: AgentConfig, model_filename: str): + parser = ConfigParser() + if os.path.exists(CONFIG_PATH): + parser.read(CONFIG_PATH) + if not parser.has_section("agent"): + parser.add_section("agent") + parser.set("agent", "active_model", model_filename) + with open(CONFIG_PATH, "w") as f: + parser.write(f) + cfg.active_model = model_filename + + +async def _wait_for_llama(port: int, timeout: int = 120, interval: int = 2) -> bool: + if not HAS_HTTPX: + log.warning("httpx not installed, skipping health check") + return True + deadline = time.time() + timeout + url = f"http://localhost:{port}/v1/models" + async with httpx.AsyncClient() as client: + while time.time() < deadline: + try: + resp = await client.get(url, timeout=5) + if resp.status_code == 200: + return True + except (httpx.ConnectError, httpx.TimeoutException): + pass + await asyncio.sleep(interval) + return False + + +# ── main ──────────────────────────────────────────────────────────────── + +async def amain(): + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s caic[%(process)d]: %(levelname)s %(message)s", + ) + log.info("cAIc node agent starting") + + if not HAS_AIO_PIKA: + log.error("aio-pika not installed") + sys.exit(1) + + cfg = AgentConfig.from_ini() + log.info("node_name=%s node_ip=%s", cfg.node_name, cfg.node_ip) + + inventory = discover_models(cfg.models_dir) + log.info("discovered %d models", len(inventory)) + + # Connect to AMQP + conn = await aio_pika.connect_robust(cfg.amqp_url) + channel = await conn.channel() + admin_ex, system_ex = await declare_exchanges(channel) + log.info("connected to AMQP broker") + + # Publish registration + reg_payload = build_registration_payload(cfg, inventory) + await publish(channel, admin_ex, f"node.{cfg.node_name}.register", reg_payload) + log.info("registration published") + + # Wait for admission + response_queue = await channel.declare_queue("", exclusive=True) + await response_queue.bind(admin_ex, f"node.{cfg.node_name}.admitted") + await response_queue.bind(admin_ex, f"node.{cfg.node_name}.rejected") + + admitted = False + async with response_queue.iterator() as iterator: + async for message in iterator: + async with message.process(): + payload = json.loads(message.body.decode()) + if payload.get("type") == "admitted": + log.info("admitted to cluster") + admitted = True + break + else: + log.error("rejected: %s", payload.get("reason")) + sys.exit(1) + + if not admitted: + log.error("no admission response received") + sys.exit(1) + + # Set up ping consumer + ping_queue = await channel.declare_queue("", exclusive=True) + await ping_queue.bind(admin_ex, f"node.{cfg.node_name}.ping") + await ping_queue.consume(lambda msg: handle_ping(cfg, channel, admin_ex, msg)) + + # Set up swap consumer + swap_queue = await channel.declare_queue("", exclusive=True) + await swap_queue.bind(admin_ex, f"node.{cfg.node_name}.cmd.swap_model") + await swap_queue.consume(lambda msg: handle_swap_model(cfg, channel, (admin_ex, system_ex), msg)) + + log.info("listening for pings and commands") + # Run forever + await asyncio.Event().wait() + + +if __name__ == "__main__": + asyncio.run(amain()) diff --git a/node_agent/requirements.txt b/node_agent/requirements.txt new file mode 100644 index 0000000..3b80356 --- /dev/null +++ b/node_agent/requirements.txt @@ -0,0 +1,3 @@ +aio-pika>=9.0.0 +psutil>=5.9.0 +httpx>=0.27.0 diff --git a/routers/chat.py b/routers/chat.py index 8de2d23..a6fc7ca 100644 --- a/routers/chat.py +++ b/routers/chat.py @@ -8,10 +8,11 @@ import httpx from fastapi import APIRouter, HTTPException, Request from fastapi.responses import StreamingResponse -from config import DEFAULT_MODEL, LLAMA_SERVER_BASE +from config import DEFAULT_MODEL from db import get_db, get_upload_context from memory import process_remember_command from rag import build_system_prompt +from triage import get_inference_url as _get_inference_url from search import (calculate_perplexity, is_uncertain, is_refusal, clean_hedging, format_search_results, format_direct_answer, extract_search_query, query_searxng) @@ -122,6 +123,7 @@ async def chat(request: Request): full_response = [] all_logprobs = [] tokens_per_sec = 0.0 + inference_base = await _get_inference_url(user_message) if remember_response: yield f"data: {json.dumps({'token': remember_response + chr(10) + chr(10), 'conversation_id': conv_id})}\n\n" @@ -129,7 +131,7 @@ async def chat(request: Request): async with httpx.AsyncClient() as client: try: async with client.stream( - "POST", f"{LLAMA_SERVER_BASE}/v1/chat/completions", + "POST", f"{inference_base}/v1/chat/completions", json=upstream_payload, timeout=httpx.Timeout(300.0, connect=10.0), ) as resp: @@ -168,7 +170,7 @@ async def chat(request: Request): augmented_response = [] async with client.stream( - "POST", f"{LLAMA_SERVER_BASE}/v1/chat/completions", + "POST", f"{inference_base}/v1/chat/completions", json={"model": model, "messages": augmented_messages, "stream": True}, timeout=httpx.Timeout(300.0, connect=10.0), ) as resp2: diff --git a/tests/test_chat_streaming_and_memory_paths.py b/tests/test_chat_streaming_and_memory_paths.py index 47d797f..d8cfd1b 100644 --- a/tests/test_chat_streaming_and_memory_paths.py +++ b/tests/test_chat_streaming_and_memory_paths.py @@ -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" diff --git a/tests/test_node_agent.py b/tests/test_node_agent.py new file mode 100644 index 0000000..4e360b0 --- /dev/null +++ b/tests/test_node_agent.py @@ -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" diff --git a/tests/test_triage.py b/tests/test_triage.py new file mode 100644 index 0000000..86ae99c --- /dev/null +++ b/tests/test_triage.py @@ -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 diff --git a/triage.py b/triage.py new file mode 100644 index 0000000..8dd2dcc --- /dev/null +++ b/triage.py @@ -0,0 +1,79 @@ +"""cAIc — Query triage and cluster node selection.""" +import logging + +import httpx + +from config import TRIAGE_BASE, TRIAGE_TIMEOUT, LLAMA_SERVER_BASE + +log = logging.getLogger("caic") + +_CLASSIFICATION_PROMPT = """Classify the following user query into exactly one category. Respond with only the category name. + +Categories: +- general: everyday questions, chitchat, creative writing, advice, explanations +- code: programming, debugging, code generation, technical questions about software +- search: questions about current events, real-time information, weather, news, specific things that may have changed since training +- rag: questions about specific documents, personal data, notes, memory, uploaded content + +Query: {query} +Category:""" + + +async def classify_query(query: str) -> str: + try: + async with httpx.AsyncClient() as client: + resp = await client.post( + f"{TRIAGE_BASE}/chat/completions", + json={ + "model": "phi-4-mini", + "messages": [ + {"role": "system", "content": "You are a query classifier. Respond with exactly one word."}, + {"role": "user", "content": _CLASSIFICATION_PROMPT.format(query=query)}, + ], + "temperature": 0.0, + "max_tokens": 10, + }, + timeout=TRIAGE_TIMEOUT, + ) + text = resp.json()["choices"][0]["message"]["content"].strip().lower() + valid = {"general", "code", "search", "rag"} + for v in valid: + if v in text: + return v + except Exception: + log.warning("triage classify_query failed, falling back to general", exc_info=True) + return "general" + + +def select_node(classification: str) -> dict | None: + from cluster import CLUSTER_NODES + + for node in CLUSTER_NODES.values(): + if node.get("status") != "active": + continue + am = node.get("active_model") or {} + name = (am.get("name") or "").lower() + if classification == "code": + if "coder" in name or "qwen" in name: + return node + elif classification == "general": + if "mistral" in name or "llama" in name: + return node + else: + return None + return None + + +async def get_inference_url(query: str) -> str: + if not query: + return LLAMA_SERVER_BASE + classification = await classify_query(query) + if classification in ("search", "rag"): + return LLAMA_SERVER_BASE + node = select_node(classification) + if node: + am = node.get("active_model") or {} + port = am.get("port", 8081) + ip = node.get("ip") or "127.0.0.1" + return f"http://{ip}:{port}/v1" + return LLAMA_SERVER_BASE