gramps
9d1fd44d7f
- request_model_swap() publishes cmd.swap_model on jc.admin, sets node status to swapping - handle_model_ready() updates active_model from inventory, restores active status - handle_model_failed() sets node status to error with failure detail - select_node() made async with two-pass logic: find ready node or trigger swap - inventory field stored in node records for swap target lookup - Both handlers subscribed on jc.system via SUBSCRIBE_TABLE
99 lines
3.5 KiB
Python
99 lines
3.5 KiB
Python
"""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")
|
|
|
|
_IDEAL_MODEL_MAP = {
|
|
"code": {"name_contains": ["coder", "qwen"]},
|
|
"general": {"name_contains": ["mistral", "llama"]},
|
|
}
|
|
|
|
_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"
|
|
|
|
|
|
async def select_node(classification: str) -> dict | None:
|
|
from cluster import CLUSTER_NODES
|
|
|
|
if classification in ("search", "rag"):
|
|
return None
|
|
|
|
ideal = _IDEAL_MODEL_MAP.get(classification, {})
|
|
ideal_contains = ideal.get("name_contains", [])
|
|
|
|
# First pass: find an active node with the right model already loaded
|
|
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 any(ideal in name for ideal in ideal_contains):
|
|
return node
|
|
|
|
# Second pass: find an active node that can swap to the right model
|
|
for node in CLUSTER_NODES.values():
|
|
if node.get("status") != "active":
|
|
continue
|
|
inventory = node.get("inventory") or []
|
|
for inv in inventory:
|
|
inv_name = (inv.get("name") or "").lower()
|
|
if any(ideal in inv_name for ideal in ideal_contains):
|
|
from cluster import request_model_swap
|
|
await request_model_swap(node["name"], inv["filename"])
|
|
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 = await 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
|