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:
@@ -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