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:
@@ -28,6 +28,8 @@ Every router has a dedicated test file:
|
|||||||
| `test_profile.py` | `routers/profile.py` — get, update, default, length validation |
|
| `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_route.py` | `routers/search_route.py` — explicit search flow, no results, errors |
|
||||||
| `test_search_url_sanitization.py` | `search.py` URL sanitizer |
|
| `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_settings_allowlist.py` | `routers/settings.py` — allowlisted key enforcement |
|
||||||
| `test_skills_framework.py` | `routers/skills.py` — list, toggle, unknown skill, prompt injection |
|
| `test_skills_framework.py` | `routers/skills.py` — list, toggle, unknown skill, prompt injection |
|
||||||
| `test_ip_allowlist.py` | IP allowlist helper + middleware |
|
| `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 |
|
| `search.py` | SearXNG integration, perplexity scoring, refusal detection |
|
||||||
| `rag.py` | Qdrant vector search + system prompt assembly + chunk_text() helper |
|
| `rag.py` | Qdrant vector search + system prompt assembly + chunk_text() helper |
|
||||||
| `gpu.py` | AMD GPU stats via `rocm-smi` |
|
| `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) |
|
| `routers/` | One module per endpoint group (chat, search, skills, completions, upload, ingest) |
|
||||||
|
|
||||||
### Entrypoint / API keys
|
### Entrypoint / API keys
|
||||||
@@ -66,7 +70,7 @@ Refactored from single-file (`app.py`) into modules under project root:
|
|||||||
|
|
||||||
### Key flows
|
### 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
|
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
|
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
|
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 |
|
| Service | Required | Port |
|
||||||
|---------|----------|------|
|
|---------|----------|------|
|
||||||
| llama-server (coordinator) | Yes | 8081 + RPC :50052 (worker GPU) |
|
| llama-server (coordinator) | Yes | 8081 + RPC :50052 (worker GPU) |
|
||||||
|
| Phi-4-mini (triage) | No | 8083 |
|
||||||
| SearXNG | No | 8888 |
|
| SearXNG | No | 8888 |
|
||||||
| wttr.in | No | weather shortcut |
|
| wttr.in | No | weather shortcut |
|
||||||
| rocm-smi | No | AMD GPU stats |
|
| rocm-smi | No | AMD GPU stats |
|
||||||
|
|||||||
@@ -95,7 +95,11 @@ Developer wiki: [docs/wiki/Home.md](docs/wiki/Home.md)
|
|||||||
│ └── logo.png # Logo image (optional)
|
│ └── logo.png # Logo image (optional)
|
||||||
├── templates/
|
├── templates/
|
||||||
│ └── index.html # Frontend
|
│ └── 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
|
## Requirements
|
||||||
@@ -309,7 +313,7 @@ Settings are stored in the `settings` table and include:
|
|||||||
python3 -m pytest tests/ -v
|
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
|
## License
|
||||||
|
|
||||||
|
|||||||
@@ -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.
|
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.
|
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.
|
||||||
|
|
||||||
|
|||||||
@@ -64,6 +64,7 @@ async def handle_registration(exchange: str, routing_key: str, payload: dict) ->
|
|||||||
"name": node_name,
|
"name": node_name,
|
||||||
"type": node_type,
|
"type": node_type,
|
||||||
"status": "active",
|
"status": "active",
|
||||||
|
"ip": payload.get("ip"),
|
||||||
"capabilities": payload.get("capabilities", []),
|
"capabilities": payload.get("capabilities", []),
|
||||||
"active_model": payload.get("active_model"),
|
"active_model": payload.get("active_model"),
|
||||||
"load": payload.get("load"),
|
"load": payload.get("load"),
|
||||||
|
|||||||
@@ -9,7 +9,7 @@ import logging
|
|||||||
|
|
||||||
log = logging.getLogger("caic")
|
log = logging.getLogger("caic")
|
||||||
|
|
||||||
VERSION = "v0.14.0"
|
VERSION = "v0.15.0"
|
||||||
OLLAMA_BASE = os.environ.get("OLLAMA_BASE", "http://localhost:11434")
|
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")
|
LLAMA_SERVER_BASE = os.environ.get("LLAMA_SERVER_BASE", "http://192.168.50.108:8081")
|
||||||
SEARXNG_BASE = "http://localhost:8888"
|
SEARXNG_BASE = "http://localhost:8888"
|
||||||
@@ -102,6 +102,11 @@ ALLOWED_SETTINGS_KEYS = {
|
|||||||
"skills_enabled",
|
"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 ---
|
||||||
PERPLEXITY_THRESHOLD = 15.0
|
PERPLEXITY_THRESHOLD = 15.0
|
||||||
|
|
||||||
|
|||||||
@@ -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<name>.+?)-(?P<version>[^-]+)-(?P<quant>Q\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())
|
||||||
@@ -0,0 +1,3 @@
|
|||||||
|
aio-pika>=9.0.0
|
||||||
|
psutil>=5.9.0
|
||||||
|
httpx>=0.27.0
|
||||||
+5
-3
@@ -8,10 +8,11 @@ import httpx
|
|||||||
from fastapi import APIRouter, HTTPException, Request
|
from fastapi import APIRouter, HTTPException, Request
|
||||||
from fastapi.responses import StreamingResponse
|
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 db import get_db, get_upload_context
|
||||||
from memory import process_remember_command
|
from memory import process_remember_command
|
||||||
from rag import build_system_prompt
|
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,
|
from search import (calculate_perplexity, is_uncertain, is_refusal,
|
||||||
clean_hedging, format_search_results, format_direct_answer,
|
clean_hedging, format_search_results, format_direct_answer,
|
||||||
extract_search_query, query_searxng)
|
extract_search_query, query_searxng)
|
||||||
@@ -122,6 +123,7 @@ async def chat(request: Request):
|
|||||||
full_response = []
|
full_response = []
|
||||||
all_logprobs = []
|
all_logprobs = []
|
||||||
tokens_per_sec = 0.0
|
tokens_per_sec = 0.0
|
||||||
|
inference_base = await _get_inference_url(user_message)
|
||||||
|
|
||||||
if remember_response:
|
if remember_response:
|
||||||
yield f"data: {json.dumps({'token': remember_response + chr(10) + chr(10), 'conversation_id': conv_id})}\n\n"
|
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:
|
async with httpx.AsyncClient() as client:
|
||||||
try:
|
try:
|
||||||
async with client.stream(
|
async with client.stream(
|
||||||
"POST", f"{LLAMA_SERVER_BASE}/v1/chat/completions",
|
"POST", f"{inference_base}/v1/chat/completions",
|
||||||
json=upstream_payload,
|
json=upstream_payload,
|
||||||
timeout=httpx.Timeout(300.0, connect=10.0),
|
timeout=httpx.Timeout(300.0, connect=10.0),
|
||||||
) as resp:
|
) as resp:
|
||||||
@@ -168,7 +170,7 @@ async def chat(request: Request):
|
|||||||
|
|
||||||
augmented_response = []
|
augmented_response = []
|
||||||
async with client.stream(
|
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},
|
json={"model": model, "messages": augmented_messages, "stream": True},
|
||||||
timeout=httpx.Timeout(300.0, connect=10.0),
|
timeout=httpx.Timeout(300.0, connect=10.0),
|
||||||
) as resp2:
|
) as resp2:
|
||||||
|
|||||||
@@ -9,9 +9,16 @@ import app
|
|||||||
import config
|
import config
|
||||||
import db
|
import db
|
||||||
import routers.chat
|
import routers.chat
|
||||||
|
import triage
|
||||||
from security import SESSIONS, PIN_ATTEMPTS, RATE_EVENTS
|
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:
|
def make_client(tmp_path: Path) -> TestClient:
|
||||||
os.environ["CAIC_ADMIN_PIN"] = "1234"
|
os.environ["CAIC_ADMIN_PIN"] = "1234"
|
||||||
db.DB_PATH = tmp_path / "caic-streaming.db"
|
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):
|
def test_chat_stream_emits_tokens_and_done(tmp_path: Path, monkeypatch):
|
||||||
|
_mock_triage_url(monkeypatch)
|
||||||
with make_client(tmp_path) as client:
|
with make_client(tmp_path) as client:
|
||||||
sid = client.post("/api/auth/guest", headers={"Origin": "http://testserver"}).json()[
|
sid = client.post("/api/auth/guest", headers={"Origin": "http://testserver"}).json()[
|
||||||
"session_id"
|
"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):
|
def test_chat_auto_search_trigger_emits_search_events(tmp_path: Path, monkeypatch):
|
||||||
|
_mock_triage_url(monkeypatch)
|
||||||
with make_client(tmp_path) as client:
|
with make_client(tmp_path) as client:
|
||||||
sid = client.post("/api/auth/guest", headers={"Origin": "http://testserver"}).json()[
|
sid = client.post("/api/auth/guest", headers={"Origin": "http://testserver"}).json()[
|
||||||
"session_id"
|
"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):
|
def test_chat_with_upload_context_id_injects_document(tmp_path: Path, monkeypatch):
|
||||||
|
_mock_triage_url(monkeypatch)
|
||||||
captured_payload = {}
|
captured_payload = {}
|
||||||
|
|
||||||
def stream_stub(self, method, url, json=None, timeout=None):
|
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):
|
def test_chat_with_expired_upload_context_id_silent(tmp_path: Path, monkeypatch):
|
||||||
|
_mock_triage_url(monkeypatch)
|
||||||
captured_payload = {}
|
captured_payload = {}
|
||||||
|
|
||||||
def stream_stub(self, method, url, json=None, timeout=None):
|
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):
|
def test_memory_command_paths_remember_and_forget(tmp_path: Path, monkeypatch):
|
||||||
|
_mock_triage_url(monkeypatch)
|
||||||
with make_client(tmp_path) as client:
|
with make_client(tmp_path) as client:
|
||||||
sid = client.post("/api/auth/guest", headers={"Origin": "http://testserver"}).json()[
|
sid = client.post("/api/auth/guest", headers={"Origin": "http://testserver"}).json()[
|
||||||
"session_id"
|
"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
|
||||||
@@ -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
|
||||||
Reference in New Issue
Block a user