v0.17.3: bug-fix maintenance pass — ingest origin exemption, FK safety, auto-search reset, deterministic ingest, WAL, config centralization, dead-code removal
- /api/ingest exempt from origin check for CLI/Bearer clients - recreate deleted conversations on bogus conversation_id (chat + search) - augmented auto-search emits reset:true; frontend clears first-pass text - image uploads stored as placeholders, never text-ingested - check_fact_conflicts requires shared subject keywords - deterministic ingest point ids (md5 chunk hash) - get_load() no longer crashes when rocm-smi yields no VRAM lines - SQLite WAL + busy_timeout; timing-safe API key compares - supervise fire-and-forget auto-ingest tasks; log missing logprobs - centralize EMBED_URL/EMBED_MODEL/QDRANT_URL/NODE_NAME in config - remove dead triage.py, select_node tests, is_state_changing, stray artifacts - add regression tests + autouse global-reset conftest
This commit is contained in:
gramps
@@ -35,15 +35,15 @@ Every router has a dedicated test file:
|
||||
| `test_search_url_sanitization.py` | `search.py` URL sanitizer |
|
||||
| `test_cluster.py` | `cluster.py` — registration, deregistration, pong, events, coordinator query |
|
||||
| `test_cluster_heartbeat.py` | `cluster.py` — heartbeat handler, known/unknown node |
|
||||
| `test_model_swap.py` | `cluster.py` + `triage.py` — request_model_swap, handle_model_ready/failed, select_node swap triggering |
|
||||
| `test_model_swap.py` | `cluster.py` — request_model_swap, handle_model_ready/failed |
|
||||
| `test_node_agent.py` | `node_agent/agent.py` — registration, ping/pong, model swap |
|
||||
| `test_image.py` | Image generation — cluster handlers, router proxy, node agent ComfyUI integration, hardware probe, capability detection |
|
||||
| `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 |
|
||||
| `test_rate_and_payload_guardrails.py` | Rate limits + payload size enforcement |
|
||||
| `test_error_envelopes.py` | Global exception handler + stream error incidents |
|
||||
| `test_fixes_regression.py` | Origin-exempt ingest, bogus conversation_id FK, auto-search reset, image uploads, conflict false-positives, deterministic ingest ids, get_load VRAM parsing, version pin |
|
||||
| `test_upload.py` | `routers/upload.py` — upload, delete, link, by-conversation, attachment_count integration |
|
||||
|
||||
Modules that call `httpx.AsyncClient` (chat, completions, models, search_route, upload, ingest, model_pull)
|
||||
@@ -69,7 +69,6 @@ Refactored from single-file (`app.py`) into modules under project root:
|
||||
| `gpu.py` | GPU stats — `rocm-smi` (AMD/Linux) or `system_profiler` (Apple Silicon/macOS) |
|
||||
| `crypto.py` | AES-256-GCM encrypt/decrypt + key management (stored as `heartbeat_interval_ms` in settings) |
|
||||
| `model_pull.py` | Startup model availability check + Ollama pull API |
|
||||
| `triage.py` | Phi-4-mini-based query classification + cluster node selection |
|
||||
| `cluster.py` | Cluster node registry, event log, coordinator election, ping/pong, model swap handlers, image generation request/response |
|
||||
| `amqp.py` | AMQP connection manager — connect, disconnect, publish, subscribe, auto-reconnect |
|
||||
| `node_agent/` | Standalone worker agent — AMQP client for registration, ping/pong, model swap, image generation |
|
||||
@@ -84,7 +83,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) → 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
|
||||
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_BASE` 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
|
||||
@@ -118,7 +117,6 @@ All services are available bare-metal or as containers in `docker compose up`.
|
||||
| Service | Required | Port | Docker service name |
|
||||
|---------|----------|------|---------------------|
|
||||
| llama-server (coordinator) | Yes | 8081 + RPC :50052 (worker GPU) | `llama-server` |
|
||||
| Phi-4-mini (triage) | No | 8083 | — |
|
||||
| SearXNG | No | 8888 | `searxng` |
|
||||
| RabbitMQ (coordinator) | No | 5672 — AMQP broker | `rabbitmq` |
|
||||
| wttr.in | No | weather shortcut | — |
|
||||
|
||||
@@ -22,7 +22,7 @@ from hardware import assess_hardware
|
||||
from memory import get_memory_count
|
||||
from security import (
|
||||
get_client_ip, is_ip_allowed, check_rate_limit, rate_policy,
|
||||
origin_allowed, is_state_changing, request_body_limit,
|
||||
origin_allowed, request_body_limit,
|
||||
audit_event, customer_error_envelope, log_incident,
|
||||
)
|
||||
from auth import get_session, is_admin_only, router as auth_router
|
||||
@@ -141,7 +141,12 @@ async def session_auth_middleware(request: Request, call_next):
|
||||
"/api/auth/heartbeat", "/api/auth/guest", "/api/ingest", "/api/hardware",
|
||||
}
|
||||
|
||||
if path.startswith("/api/"):
|
||||
# Bearer-token-authenticated endpoints are reached by CLI/terminal tooling
|
||||
# (curl, caic-ingest.sh) that sends no Origin/Referer header — exempt them
|
||||
# from the browser origin check.
|
||||
origin_exempt_paths = {"/api/ingest"}
|
||||
|
||||
if path.startswith("/api/") and path not in origin_exempt_paths:
|
||||
if not origin_allowed(request):
|
||||
audit_event("origin_check", "denied", ip=ip, role="none",
|
||||
details=f"{request.method} {path}", warning=True)
|
||||
|
||||
-2334
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -17,7 +17,7 @@ from db import get_db, get_setting
|
||||
from security import (
|
||||
SESSIONS, PIN_ATTEMPTS, SESSION_LOCK, BODY_LIMIT_DEFAULT_BYTES,
|
||||
audit_event, get_client_ip, is_ip_allowed, check_rate_limit,
|
||||
rate_policy, origin_allowed, is_state_changing, request_body_limit,
|
||||
rate_policy, origin_allowed, request_body_limit,
|
||||
read_json_body, hash_pin, customer_error_envelope, log_incident,
|
||||
)
|
||||
|
||||
|
||||
@@ -30,6 +30,7 @@ def get_db():
|
||||
conn.row_factory = sqlite3.Row
|
||||
conn.execute("PRAGMA foreign_keys = ON")
|
||||
conn.execute("PRAGMA journal_mode = WAL")
|
||||
conn.execute("PRAGMA busy_timeout = 5000")
|
||||
return conn
|
||||
|
||||
|
||||
@@ -132,6 +133,8 @@ def init_db():
|
||||
from security import hash_pin
|
||||
conn = sqlite3.connect(DB_PATH)
|
||||
conn.row_factory = sqlite3.Row
|
||||
conn.execute("PRAGMA journal_mode = WAL")
|
||||
conn.execute("PRAGMA busy_timeout = 5000")
|
||||
|
||||
# --- Core tables ---
|
||||
conn.execute("""
|
||||
|
||||
@@ -25,7 +25,6 @@ Refactored from single-file (`app.py`) into modules under project root:
|
||||
| `hardware.py` | Hardware self-assessment — CPU, RAM, VRAM, service health probes (llama-server, Qdrant, SearXNG, ComfyUI) |
|
||||
| `amqp.py` | aio-pika connection manager for RabbitMQ (connect, disconnect, publish, subscribe, auto-reconnect) |
|
||||
| `cluster.py` | Cluster node registry, event log, coordinator election, ping/pong, model swap handlers, image generation request/response |
|
||||
| `triage.py` | Phi-4-mini query classification + `select_node()` for cluster routing |
|
||||
| `routers/` | One module per endpoint group |
|
||||
|
||||
### 1.2 External Services
|
||||
@@ -308,7 +307,6 @@ All streaming endpoints yield `data: {json}\n\n`:
|
||||
| test_search_url_sanitization.py | URL sanitizer |
|
||||
| test_settings_allowlist.py | Allowlisted key enforcement |
|
||||
| test_skills_framework.py | List, toggle, unknown skill, prompt injection |
|
||||
| test_triage.py | classify_query, select_node, get_inference_url |
|
||||
| test_upload.py | Upload, delete, link, by-conversation, attachment_count |
|
||||
|
||||
### 8.3 DoD Process
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -38,6 +38,21 @@ AUTO_FACT_PATTERNS = [
|
||||
]
|
||||
SOCIAL_TRIGGERS = {"hi", "hello", "hey", "yo", "sup", "howdy", "good morning", "good evening"}
|
||||
|
||||
# Short filler words that shouldn't count as subject overlap between facts.
|
||||
_STOPWORDS = {
|
||||
"with", "that", "have", "this", "from", "they", "what", "when", "where",
|
||||
"which", "there", "your", "will", "would", "about", "these", "their",
|
||||
"been", "into", "than", "then", "them", "were", "being", "more", "most",
|
||||
"some", "other", "only", "still", "also", "after", "before", "during",
|
||||
"because", "through", "without",
|
||||
}
|
||||
|
||||
|
||||
def _subject_words(text: str) -> set:
|
||||
"""Meaningful subject tokens for overlap comparison."""
|
||||
words = re.findall(r"[A-Za-z0-9_]{4,}", text.lower())
|
||||
return {w for w in words if w not in _STOPWORDS}
|
||||
|
||||
|
||||
def _is_social(text: str) -> bool:
|
||||
t = text.strip().lower()
|
||||
@@ -86,6 +101,10 @@ def auto_detect_facts(user_message: str, assistant_message: str) -> list[str]:
|
||||
def check_fact_conflicts(facts: list[str]) -> list[dict]:
|
||||
"""Search for existing memories that conflict with detected facts.
|
||||
|
||||
A conflict is reported only when the existing memory is about the same
|
||||
subject (meaningful keyword overlap) but states something different —
|
||||
unrelated hits that merely share an FTS keyword are not conflicts.
|
||||
|
||||
Returns list of {memory_id, old_fact, new_fact} for each conflict.
|
||||
"""
|
||||
conflicts = []
|
||||
@@ -93,7 +112,7 @@ def check_fact_conflicts(facts: list[str]) -> list[dict]:
|
||||
related = search_memories(new_fact, limit=1)
|
||||
if related:
|
||||
old = related[0]["fact"]
|
||||
if old.rstrip(".") != new_fact.rstrip("."):
|
||||
if old.rstrip(".") != new_fact.rstrip(".") and (_subject_words(new_fact) & _subject_words(old)):
|
||||
conflicts.append({
|
||||
"memory_id": related[0]["rowid"],
|
||||
"old_fact": old,
|
||||
|
||||
+11
-9
@@ -184,16 +184,18 @@ def get_load() -> dict:
|
||||
capture_output=True, text=True, timeout=3,
|
||||
)
|
||||
if result.returncode == 0:
|
||||
total = 0
|
||||
used = 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:
|
||||
if "VRAM Total Used Memory (B)" in line:
|
||||
parts = line.split(":")
|
||||
if len(parts) >= 2:
|
||||
used = int(parts[-1].strip())
|
||||
elif "VRAM Total Memory (B)" in line:
|
||||
parts = line.split(":")
|
||||
if len(parts) >= 2:
|
||||
total = int(parts[-1].strip())
|
||||
if total > 0:
|
||||
load["vram_pct"] = round(used / total * 100)
|
||||
except (FileNotFoundError, subprocess.TimeoutExpired):
|
||||
pass
|
||||
|
||||
+26
-5
@@ -23,6 +23,22 @@ from config import MAX_CHAT_MESSAGE_CHARS, MODEL_CONTEXT_LENGTH
|
||||
log = logging.getLogger("caic")
|
||||
router = APIRouter()
|
||||
|
||||
# References to background auto-ingest tasks so they are never garbage-collected.
|
||||
_ingest_tasks: set = set()
|
||||
|
||||
|
||||
async def _safe_ingest(coro):
|
||||
try:
|
||||
await coro
|
||||
except Exception as e:
|
||||
log.warning("auto-ingest task failed: %s", e)
|
||||
|
||||
|
||||
def _spawn_ingest(coro):
|
||||
task = asyncio.create_task(_safe_ingest(coro))
|
||||
_ingest_tasks.add(task)
|
||||
task.add_done_callback(_ingest_tasks.discard)
|
||||
|
||||
|
||||
def parse_llama_stream_chunk(line: str) -> tuple:
|
||||
if line.startswith("data: "):
|
||||
@@ -112,6 +128,11 @@ async def chat(request: Request):
|
||||
db.execute("INSERT INTO conversations (id, title, model, created_at, updated_at) VALUES (?, ?, ?, ?, ?)",
|
||||
(conv_id, encrypt_text(title), model, now, now))
|
||||
else:
|
||||
# A client-supplied id may reference a conversation that no longer exists;
|
||||
# recreate the row so the message insert satisfies the FK instead of 500ing.
|
||||
title = user_message[:80] + ("..." if len(user_message) > 80 else "")
|
||||
db.execute("INSERT OR IGNORE INTO conversations (id, title, model, created_at, updated_at) VALUES (?, ?, ?, ?, ?)",
|
||||
(conv_id, encrypt_text(title), model, now, now))
|
||||
db.execute("UPDATE conversations SET updated_at = ? WHERE id = ?", (now, conv_id))
|
||||
|
||||
db.execute("INSERT INTO messages (conversation_id, role, content, created_at, perplexity) VALUES (?, ?, ?, ?, ?)",
|
||||
@@ -171,6 +192,8 @@ async def chat(request: Request):
|
||||
|
||||
assistant_msg = "".join(full_response)
|
||||
perplexity = calculate_perplexity(all_logprobs) if all_logprobs else 0.0
|
||||
if not all_logprobs:
|
||||
log.warning("No logprobs received from inference server — perplexity auto-search unavailable")
|
||||
should_search = is_uncertain(all_logprobs) or is_refusal(assistant_msg)
|
||||
|
||||
if search_enabled and should_search:
|
||||
@@ -210,7 +233,7 @@ async def chat(request: Request):
|
||||
if is_refusal(cleaned_response) or len(cleaned_response) < 20:
|
||||
cleaned_response = format_direct_answer(user_message, search_results)
|
||||
|
||||
yield f"data: {json.dumps({'token': cleaned_response, 'conversation_id': conv_id, 'augmented': True})}\n\n"
|
||||
yield f"data: {json.dumps({'token': cleaned_response, 'conversation_id': conv_id, 'augmented': True, 'reset': True})}\n\n"
|
||||
|
||||
if not private_chat:
|
||||
saved_msg = cleaned_response + "\n\n---\n*🔍 Enhanced with web search results*"
|
||||
@@ -229,8 +252,7 @@ async def chat(request: Request):
|
||||
if conflicts:
|
||||
rag_update = {"conflicts": conflicts}
|
||||
else:
|
||||
# Fire-and-forget: persist facts silently, don't block the response
|
||||
asyncio.create_task(ingest_auto_fact(facts, user_message, cleaned_response))
|
||||
_spawn_ingest(ingest_auto_fact(facts, user_message, cleaned_response))
|
||||
|
||||
yield f"data: {json.dumps({'done': True, 'conversation_id': conv_id, 'searched': True, 'perplexity': round(perplexity, 2), 'tokens_per_sec': round(tokens_per_sec, 1), 'prompt_tokens': prompt_tokens, 'completion_tokens': completion_tokens, 'context_length': MODEL_CONTEXT_LENGTH, **(rag_update and {'rag_update_suggestion': rag_update} or {})})}\n\n"
|
||||
return
|
||||
@@ -252,8 +274,7 @@ async def chat(request: Request):
|
||||
if conflicts:
|
||||
rag_update = {"conflicts": conflicts}
|
||||
else:
|
||||
# Fire-and-forget: persist facts silently, don't block the response
|
||||
asyncio.create_task(ingest_auto_fact(facts, user_message, assistant_msg))
|
||||
_spawn_ingest(ingest_auto_fact(facts, user_message, assistant_msg))
|
||||
|
||||
yield f"data: {json.dumps({'done': True, 'conversation_id': conv_id, 'perplexity': round(perplexity, 2), 'tokens_per_sec': round(tokens_per_sec, 1), 'prompt_tokens': prompt_tokens, 'completion_tokens': completion_tokens, 'context_length': MODEL_CONTEXT_LENGTH, **(rag_update and {'rag_update_suggestion': rag_update} or {})})}\n\n"
|
||||
|
||||
|
||||
@@ -6,6 +6,7 @@ FIM (fill-in-the-middle) requests are proxied directly — not persisted.
|
||||
Chat-style requests are persisted to conversation history.
|
||||
Auth: static Bearer token via COMPLETIONS_API_KEY in config.
|
||||
"""
|
||||
import hmac
|
||||
import json
|
||||
import logging
|
||||
import uuid
|
||||
@@ -30,7 +31,7 @@ def _check_api_key(request: Request):
|
||||
if not auth.startswith("Bearer "):
|
||||
raise HTTPException(status_code=401, detail="Missing Bearer token")
|
||||
token = auth[7:].strip()
|
||||
if token != COMPLETIONS_API_KEY:
|
||||
if not hmac.compare_digest(token.encode(), COMPLETIONS_API_KEY.encode()):
|
||||
raise HTTPException(status_code=401, detail="Invalid API key")
|
||||
|
||||
|
||||
|
||||
+5
-2
@@ -1,4 +1,6 @@
|
||||
"""JarvisChat routers - /api/ingest terminal command RAG hook."""
|
||||
import hashlib
|
||||
import hmac
|
||||
import logging
|
||||
from datetime import datetime, timezone
|
||||
|
||||
@@ -20,7 +22,7 @@ def _check_api_key(request: Request):
|
||||
if not auth.startswith("Bearer "):
|
||||
raise HTTPException(status_code=401, detail="Missing Bearer token")
|
||||
token = auth[7:].strip()
|
||||
if token != COMPLETIONS_API_KEY:
|
||||
if not hmac.compare_digest(token.encode(), COMPLETIONS_API_KEY.encode()):
|
||||
raise HTTPException(status_code=401, detail="Invalid API key")
|
||||
|
||||
|
||||
@@ -50,7 +52,8 @@ async def ingest_content(request: Request):
|
||||
log.warning(f"Ingest embedding failed for chunk {i}: {embed_resp.status_code}")
|
||||
continue
|
||||
vector = embed_resp.json()["embedding"]
|
||||
point_id = f"ingest-{source}-{datetime.now(timezone.utc).timestamp()}-{i}"
|
||||
chunk_hash = hashlib.md5(chunk.encode("utf-8")).hexdigest()[:12]
|
||||
point_id = f"ingest-{source}-{chunk_hash}-{i}"
|
||||
payload = {"text": encrypt_text(chunk), "source": source, "ingest_date": datetime.now(timezone.utc).isoformat(), "type": "ingest"}
|
||||
payload.update(metadata)
|
||||
upsert_resp = await client.put(
|
||||
|
||||
@@ -44,6 +44,9 @@ async def explicit_search(request: Request):
|
||||
db.execute("INSERT INTO conversations (id, title, model, created_at, updated_at) VALUES (?, ?, ?, ?, ?)",
|
||||
(conv_id, encrypt_text(title), model, now, now))
|
||||
else:
|
||||
title = query[:70] + "..." if len(query) > 70 else query
|
||||
db.execute("INSERT OR IGNORE INTO conversations (id, title, model, created_at, updated_at) VALUES (?, ?, ?, ?, ?)",
|
||||
(conv_id, title, model, now, now))
|
||||
db.execute("UPDATE conversations SET updated_at = ? WHERE id = ?", (now, conv_id))
|
||||
|
||||
db.execute("INSERT INTO messages (conversation_id, role, content, created_at, perplexity) VALUES (?, ?, ?, ?, ?)",
|
||||
|
||||
+10
-1
@@ -52,12 +52,21 @@ async def upload_file(
|
||||
except Exception as e:
|
||||
log.warning(f"PDF extraction error: {e}")
|
||||
raise HTTPException(status_code=422, detail="Failed to extract text from PDF")
|
||||
elif content_type.startswith("image/"):
|
||||
# No OCR pipeline exists — store a descriptive placeholder so images
|
||||
# remain usable in the gallery/context but never pollute the RAG corpus.
|
||||
extracted = f"[Image: {file.filename}]"
|
||||
else:
|
||||
extracted = raw_bytes.decode("utf-8", errors="replace")
|
||||
|
||||
result = {"filename": file.filename, "size_bytes": len(raw_bytes), "mode": mode}
|
||||
|
||||
if mode in ("ingest", "both"):
|
||||
is_image = content_type.startswith("image/")
|
||||
if is_image and mode in ("ingest", "both"):
|
||||
result["chunks_ingested"] = 0
|
||||
result["note"] = "Image files cannot be text-ingested; stored for gallery/context only"
|
||||
|
||||
if mode in ("ingest", "both") and not is_image:
|
||||
os.makedirs(UPLOAD_DIR, exist_ok=True)
|
||||
chunks = chunk_text(extracted)
|
||||
ingested = 0
|
||||
|
||||
@@ -161,10 +161,6 @@ def origin_allowed(request: Request) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
def is_state_changing(method: str) -> bool:
|
||||
return method in {"POST", "PUT", "DELETE", "PATCH"}
|
||||
|
||||
|
||||
async def read_json_body(request: Request, max_bytes: int) -> dict:
|
||||
raw = await request.body()
|
||||
if len(raw) > max_bytes:
|
||||
|
||||
@@ -1635,6 +1635,7 @@ async function sendSearch() {
|
||||
if (data.error) { textEl.textContent = 'Error: ' + data.error; setStreamingState(false); return; }
|
||||
if (data.conversation_id && !currentConvId) { currentConvId = data.conversation_id; await loadConversations(); }
|
||||
if (data.search_results) { textEl.innerHTML = '<div class="search-indicator">🔍 Found ' + data.search_results + ' results, summarizing...</div>'; }
|
||||
if (data.reset) { fullText = ''; textEl.innerHTML = ''; firstToken = false; }
|
||||
if (data.token) { if (firstToken) { textEl.innerHTML = ''; firstToken = false; ttr = performance.now() - ttrStart; tokenCount = 0; } fullText += data.token; tokenCount++; textEl.innerHTML = renderMarkdown(fullText); scrollToLatest(); }
|
||||
if (data.raw_results) {
|
||||
let rawHtml = '<details class="raw-results"><summary>🔍 View raw search results (' + data.raw_results.length + ')</summary><ul>';
|
||||
@@ -1843,7 +1844,7 @@ async function sendMessage() {
|
||||
}
|
||||
if (data.searching) { textEl.innerHTML = fullText ? renderMarkdown(fullText) + '<div class="search-indicator"><div class="spinner"></div>Searching...</div>' : '<div class="search-indicator"><div class="spinner"></div>Searching...</div>'; searchTriggered = true; }
|
||||
if (data.search_results) { textEl.innerHTML = '<div class="search-indicator">🔍 Found ' + data.search_results + ' results...</div>'; fullText = ''; firstToken = true; }
|
||||
if (data.token) { if (firstToken) { textEl.innerHTML = ''; firstToken = false; ttr = performance.now() - ttrStart; tokenCount = 0; } fullText += data.token; tokenCount++; textEl.innerHTML = renderMarkdown(fullText); scrollToLatest(); }
|
||||
if (data.token) { if (data.reset) { fullText = ''; firstToken = true; } if (firstToken) { textEl.innerHTML = ''; firstToken = false; ttr = performance.now() - ttrStart; tokenCount = 0; } fullText += data.token; tokenCount++; textEl.innerHTML = renderMarkdown(fullText); scrollToLatest(); }
|
||||
if (data.done) {
|
||||
const roleLabel = assistantDiv.querySelector('.role-label');
|
||||
if (data.searched && roleLabel) roleLabel.textContent = 'web search';
|
||||
|
||||
@@ -4,3 +4,31 @@ import sys
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||
|
||||
"""Shared pytest fixtures.
|
||||
|
||||
All test modules manipulate in-process globals (sessions, rate buckets,
|
||||
cluster registry, eviction log). An autouse fixture resets every global
|
||||
before each test so no state leaks between tests, regardless of whether an
|
||||
individual test file remembers to clear it.
|
||||
"""
|
||||
import pytest
|
||||
|
||||
import cluster
|
||||
import routers.chat
|
||||
from eviction import EVICTION_LOG
|
||||
from security import SESSIONS, PIN_ATTEMPTS, RATE_EVENTS
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_global_state():
|
||||
SESSIONS.clear()
|
||||
PIN_ATTEMPTS.clear()
|
||||
RATE_EVENTS.clear()
|
||||
cluster.CLUSTER_NODES.clear()
|
||||
cluster.CLUSTER_EVENTS.clear()
|
||||
cluster.CLUSTER_COORDINATOR = None
|
||||
cluster._pending_pings.clear()
|
||||
EVICTION_LOG.clear()
|
||||
routers.chat._ingest_tasks.clear()
|
||||
yield
|
||||
|
||||
@@ -10,7 +10,6 @@ import app
|
||||
import config
|
||||
import db
|
||||
import routers.chat
|
||||
import triage
|
||||
from security import SESSIONS, PIN_ATTEMPTS, RATE_EVENTS
|
||||
|
||||
|
||||
@@ -269,7 +268,6 @@ def test_private_chat_does_not_persist(tmp_path: Path, monkeypatch):
|
||||
'data: {"choices":[{"delta":{"content":""},"finish_reason":"stop","logprobs":{"content":[]}}],"usage":{"completion_tokens":2,"prompt_tokens":10,"tokens_per_second":5.0}}',
|
||||
"data: [DONE]",
|
||||
]))
|
||||
monkeypatch.setattr(triage, "classify_query", lambda q: "general")
|
||||
|
||||
async def _mock_ensure(m): return True
|
||||
monkeypatch.setattr("model_pull.ensure_model", _mock_ensure)
|
||||
@@ -301,7 +299,6 @@ def test_private_chat_does_not_auto_search(tmp_path: Path, monkeypatch):
|
||||
'data: {"choices":[{"delta":{"content":""},"finish_reason":"stop","logprobs":{"content":[{"logprob":-2.5}]}}],"usage":{"completion_tokens":1,"prompt_tokens":10,"tokens_per_second":5.0}}',
|
||||
"data: [DONE]",
|
||||
]))
|
||||
monkeypatch.setattr(triage, "classify_query", lambda q: "general")
|
||||
monkeypatch.setattr(routers.chat, "query_searxng", lambda q: [{"title": "result"}])
|
||||
|
||||
with make_client(tmp_path) as client:
|
||||
|
||||
@@ -0,0 +1,357 @@
|
||||
"""Regression tests for bug fixes.
|
||||
|
||||
Covers: /api/ingest origin exemption for CLI/Bearer clients, bogus
|
||||
conversation_id FK handling in chat + search, the auto-search reset flag,
|
||||
image uploads being stored as placeholders instead of text-ingested,
|
||||
false-positive conflict detection, deterministic ingest point IDs,
|
||||
get_load() VRAM parsing, and version pinning.
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
import httpx
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import app
|
||||
import config
|
||||
import db
|
||||
import memory
|
||||
from crypto import decrypt_text
|
||||
import node_agent.agent as agent
|
||||
import routers.chat
|
||||
import routers.ingest as ingest_route
|
||||
import routers.search_route
|
||||
from security import SESSIONS, PIN_ATTEMPTS, RATE_EVENTS
|
||||
|
||||
|
||||
def make_client(tmp_path: Path) -> TestClient:
|
||||
os.environ["CAIC_ADMIN_PIN"] = "1234"
|
||||
db.DB_PATH = tmp_path / "caic-regression.db"
|
||||
SESSIONS.clear()
|
||||
PIN_ATTEMPTS.clear()
|
||||
RATE_EVENTS.clear()
|
||||
db.init_db()
|
||||
return TestClient(app.app, raise_server_exceptions=False)
|
||||
|
||||
|
||||
def _guest_headers(client: TestClient) -> dict:
|
||||
sid = client.post("/api/auth/guest", headers={"Origin": "http://testserver"}).json()["session_id"]
|
||||
return {"X-Session-ID": sid, "Origin": "http://testserver"}
|
||||
|
||||
|
||||
def _admin_headers(client: TestClient) -> dict:
|
||||
login = client.post("/api/auth/login", json={"pin": "1234"}, headers={"Origin": "http://testserver"})
|
||||
sid = login.json()["session_id"]
|
||||
return {"X-Session-ID": sid, "Origin": "http://testserver"}
|
||||
|
||||
|
||||
def parse_sse_payloads(body: str) -> list[dict]:
|
||||
payloads = []
|
||||
for chunk in body.split("\n\n"):
|
||||
chunk = chunk.strip()
|
||||
if not chunk.startswith("data: "):
|
||||
continue
|
||||
payloads.append(json.loads(chunk[len("data: "):]))
|
||||
return payloads
|
||||
|
||||
|
||||
class _MockStreamResponse:
|
||||
def __init__(self, lines: list[str]):
|
||||
self._lines = lines
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
async def aiter_lines(self):
|
||||
for line in self._lines:
|
||||
yield line
|
||||
|
||||
|
||||
def _stream_json_lines(events: list[dict]) -> list[str]:
|
||||
return [json.dumps(event) for event in events]
|
||||
|
||||
|
||||
class _FakeAsyncClient:
|
||||
class FakeResponse:
|
||||
def __init__(self, status, json_data=None):
|
||||
self.status_code = status
|
||||
self._json = json_data or {}
|
||||
|
||||
def json(self):
|
||||
return self._json
|
||||
|
||||
def __init__(self, *a, **kw):
|
||||
self.put_payloads = []
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *a):
|
||||
pass
|
||||
|
||||
async def post(self, url, **kw):
|
||||
if "/api/embeddings" in url:
|
||||
return self.FakeResponse(200, {"embedding": [0.1] * 768})
|
||||
return self.FakeResponse(200)
|
||||
|
||||
async def put(self, url, **kw):
|
||||
self.put_payloads.append(kw.get("json", {}))
|
||||
return self.FakeResponse(200)
|
||||
|
||||
|
||||
# ── /api/ingest is reached by CLI tools with no Origin header ──────────
|
||||
|
||||
|
||||
def test_ingest_origin_exemption(tmp_path: Path, monkeypatch):
|
||||
monkeypatch.setattr(ingest_route, "COMPLETIONS_API_KEY", "sk-regression")
|
||||
monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: _FakeAsyncClient())
|
||||
with make_client(tmp_path) as client:
|
||||
resp = client.post(
|
||||
"/api/ingest",
|
||||
json={"content": "regression test content " * 20, "source": "cli"},
|
||||
headers={"Authorization": "Bearer sk-regression", "Content-Type": "application/json"},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
assert resp.json()["source"] == "cli"
|
||||
|
||||
|
||||
def test_ingest_bad_key_still_blocked_without_origin(tmp_path: Path, monkeypatch):
|
||||
monkeypatch.setattr(ingest_route, "COMPLETIONS_API_KEY", "sk-regression")
|
||||
with make_client(tmp_path) as client:
|
||||
resp = client.post(
|
||||
"/api/ingest",
|
||||
json={"content": "x " * 50},
|
||||
headers={"Authorization": "Bearer wrong-key", "Content-Type": "application/json"},
|
||||
)
|
||||
assert resp.status_code == 401
|
||||
|
||||
|
||||
# ── a client-supplied conversation_id that no longer exists ────────────
|
||||
|
||||
|
||||
def test_chat_bogus_conversation_id_creates_row(tmp_path: Path, monkeypatch):
|
||||
events = _stream_json_lines([
|
||||
{"message": {"content": "hi"}, "logprobs": [{"logprob": -0.01}]},
|
||||
{"done": True, "eval_count": 1, "eval_duration": 1000000000},
|
||||
])
|
||||
|
||||
def stream_stub(self, method, url, json=None, timeout=None):
|
||||
return _MockStreamResponse(events)
|
||||
|
||||
monkeypatch.setattr(httpx.AsyncClient, "stream", stream_stub)
|
||||
|
||||
with make_client(tmp_path) as client:
|
||||
resp = client.post(
|
||||
"/api/chat",
|
||||
json={"message": "hello", "conversation_id": "ghost-conv", "model": config.DEFAULT_MODEL},
|
||||
headers=_guest_headers(client),
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
conv_resp = client.get("/api/conversations/ghost-conv", headers=_guest_headers(client))
|
||||
assert conv_resp.status_code == 200
|
||||
assert len(conv_resp.json()["messages"]) >= 2
|
||||
|
||||
|
||||
def test_search_bogus_conversation_id_creates_row(tmp_path: Path, monkeypatch):
|
||||
async def empty_search(query: str, max_results: int = 5):
|
||||
return []
|
||||
|
||||
monkeypatch.setattr(routers.search_route, "query_searxng", empty_search)
|
||||
|
||||
with make_client(tmp_path) as client:
|
||||
resp = client.post(
|
||||
"/api/search",
|
||||
json={"query": "nothing here", "conversation_id": "ghost-search", "model": config.DEFAULT_MODEL},
|
||||
headers=_guest_headers(client),
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
conv_resp = client.get("/api/conversations/ghost-search", headers=_guest_headers(client))
|
||||
assert conv_resp.status_code == 200
|
||||
assert len(conv_resp.json()["messages"]) >= 1
|
||||
|
||||
|
||||
# ── auto-search augmentation must reset the streamed text ──────────────
|
||||
|
||||
|
||||
def test_auto_search_augmented_event_has_reset_flag(tmp_path: Path, monkeypatch):
|
||||
first_stream = _stream_json_lines([
|
||||
{"message": {"content": "I don't have current data on that."}, "logprobs": [{"logprob": -5.0}]},
|
||||
{"done": True, "eval_count": 2, "eval_duration": 1000000000},
|
||||
])
|
||||
second_stream = _stream_json_lines([
|
||||
{"message": {"content": "According to the search results, the value is forty-two."}},
|
||||
{"done": True},
|
||||
])
|
||||
batches = [first_stream, second_stream]
|
||||
|
||||
def stream_stub(self, method, url, json=None, timeout=None):
|
||||
return _MockStreamResponse(batches.pop(0))
|
||||
|
||||
async def search_stub(query: str, max_results: int = 5):
|
||||
return [{"title": "Answer", "url": "https://example.com", "content": "The value is 42."}]
|
||||
|
||||
with make_client(tmp_path) as client:
|
||||
monkeypatch.setattr(httpx.AsyncClient, "stream", stream_stub)
|
||||
monkeypatch.setattr(routers.chat, "query_searxng", search_stub)
|
||||
resp = client.post(
|
||||
"/api/chat",
|
||||
json={"message": "what is the latest value", "model": config.DEFAULT_MODEL},
|
||||
headers=_guest_headers(client),
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
payloads = parse_sse_payloads(resp.text)
|
||||
|
||||
augmented = [p for p in payloads if p.get("augmented")]
|
||||
assert augmented, "expected an augmented token event"
|
||||
assert augmented[0].get("reset") is True
|
||||
# The augmented token must carry the fresh answer, not the discarded
|
||||
# first-pass "I don't have current data" text.
|
||||
assert "According to the search results" in augmented[0]["token"]
|
||||
assert "I don't have current data" not in augmented[0]["token"]
|
||||
|
||||
|
||||
# ── image uploads are placeholders, never text-ingested ────────────────
|
||||
|
||||
|
||||
def test_upload_image_skips_ingest(tmp_path: Path, monkeypatch):
|
||||
fake = _FakeAsyncClient()
|
||||
monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: fake)
|
||||
with make_client(tmp_path) as client:
|
||||
resp = client.post(
|
||||
"/api/upload",
|
||||
headers=_admin_headers(client),
|
||||
data={"mode": "both"},
|
||||
files={"file": ("photo.png", b"\x89PNG\r\n\x1a\nfake", "image/png")},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
data = resp.json()
|
||||
context_id = data["context_id"]
|
||||
|
||||
assert data["chunks_ingested"] == 0
|
||||
assert data["note"]
|
||||
assert fake.put_payloads == []
|
||||
assert data["filename"] == "photo.png"
|
||||
|
||||
row = db.get_db().execute(
|
||||
"SELECT content FROM upload_context WHERE id = ?", (context_id,)
|
||||
).fetchone()
|
||||
assert row and decrypt_text(row["content"]) == "[Image: photo.png]"
|
||||
|
||||
|
||||
# ── conflict detection needs a shared subject, not just an FTS hit ─────
|
||||
|
||||
|
||||
def test_conflict_detection_requires_shared_subject(tmp_path: Path):
|
||||
os.environ["CAIC_ADMIN_PIN"] = "1234"
|
||||
db.DB_PATH = tmp_path / "caic-mem-regression.db"
|
||||
SESSIONS.clear()
|
||||
PIN_ATTEMPTS.clear()
|
||||
RATE_EVENTS.clear()
|
||||
db.init_db()
|
||||
|
||||
memory.add_memory("the cat sat on the mat", "general")
|
||||
conflicts = memory.check_fact_conflicts(["the dog is brown"])
|
||||
assert conflicts == []
|
||||
|
||||
memory.add_memory("I prefer Rust over Go", "preference")
|
||||
conflicts = memory.check_fact_conflicts(["I prefer Go over Rust"])
|
||||
assert len(conflicts) == 1
|
||||
assert conflicts[0]["new_fact"] == "I prefer Go over Rust"
|
||||
assert conflicts[0]["old_fact"] == "I prefer Rust over Go"
|
||||
assert "memory_id" in conflicts[0]
|
||||
|
||||
|
||||
# ── ingest point IDs are deterministic (no duplicate vectors) ──────────
|
||||
|
||||
|
||||
def test_ingest_deterministic_point_ids(tmp_path: Path, monkeypatch):
|
||||
monkeypatch.setattr(ingest_route, "COMPLETIONS_API_KEY", "sk-regression")
|
||||
|
||||
captured_first = []
|
||||
captured_second = []
|
||||
|
||||
class CaptureClient:
|
||||
FakeResponse = _FakeAsyncClient.FakeResponse
|
||||
|
||||
def __init__(self, *a, **kw):
|
||||
self.capture = None
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *a):
|
||||
pass
|
||||
|
||||
async def post(self, url, **kw):
|
||||
if "/api/embeddings" in url:
|
||||
return self.FakeResponse(200, {"embedding": [0.2] * 768})
|
||||
return self.FakeResponse(200)
|
||||
|
||||
async def put(self, url, **kw):
|
||||
payload = kw.get("json", {})
|
||||
if self.capture is not None:
|
||||
self.capture.append(payload["points"][0]["id"])
|
||||
return self.FakeResponse(200)
|
||||
|
||||
body = {"content": "alpha beta gamma delta epsilon " * 8, "source": "hook"}
|
||||
headers = {"Authorization": "Bearer sk-regression", "Content-Type": "application/json"}
|
||||
|
||||
fake1 = CaptureClient()
|
||||
fake1.capture = captured_first
|
||||
monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: fake1)
|
||||
with make_client(tmp_path) as client:
|
||||
r1 = client.post("/api/ingest", json=body, headers=headers)
|
||||
assert r1.status_code == 200, r1.text
|
||||
|
||||
fake2 = CaptureClient()
|
||||
fake2.capture = captured_second
|
||||
monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: fake2)
|
||||
with make_client(tmp_path) as client:
|
||||
r2 = client.post("/api/ingest", json=body, headers=headers)
|
||||
assert r2.status_code == 200, r2.text
|
||||
|
||||
assert captured_first and captured_second
|
||||
assert len(captured_first) == len(captured_second)
|
||||
assert captured_first == captured_second, "re-ingesting identical content changed point ids"
|
||||
assert len(set(captured_first)) == len(captured_first)
|
||||
|
||||
|
||||
# ── get_load() VRAM parsing ────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_get_load_vram_parses_rocm_output(monkeypatch):
|
||||
monkeypatch.setattr(agent, "HAS_PSUTIL", False)
|
||||
output = (
|
||||
"======================= ROCm System Management Interface =======================\n"
|
||||
"GPU[0] : gfx1030\n"
|
||||
"VRAM Total Used Memory (B): 3221225472\n"
|
||||
"VRAM Total Memory (B): 17179869184\n"
|
||||
)
|
||||
fake = subprocess.CompletedProcess(["rocm-smi", "--showmeminfo", "vram"], 0, output, "")
|
||||
monkeypatch.setattr(subprocess, "run", lambda *a, **kw: fake)
|
||||
load = agent.get_load()
|
||||
assert load["vram_pct"] == 19 # 3 GiB / 16 GiB
|
||||
|
||||
|
||||
def test_get_load_vram_absent_does_not_crash(monkeypatch):
|
||||
# Regression: rocm-smi returned no parseable VRAM lines, so the old code
|
||||
# left total/used unbound and raised.
|
||||
monkeypatch.setattr(agent, "HAS_PSUTIL", False)
|
||||
output = "======================= ROCm System Management Interface =======================\nNo GPU detected\n"
|
||||
fake = subprocess.CompletedProcess(["rocm-smi", "--showmeminfo", "vram"], 0, output, "")
|
||||
monkeypatch.setattr(subprocess, "run", lambda *a, **kw: fake)
|
||||
load = agent.get_load()
|
||||
assert "vram_pct" not in load
|
||||
|
||||
|
||||
# ── version pinning ────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_version_is_bumped():
|
||||
assert re.fullmatch(r"v\d+\.\d+\.\d+", config.VERSION)
|
||||
@@ -2,7 +2,6 @@
|
||||
import asyncio
|
||||
|
||||
import cluster
|
||||
import triage
|
||||
from config import AMQP_EXCHANGE_ADMIN, AMQP_EXCHANGE_SYSTEM
|
||||
|
||||
|
||||
@@ -149,54 +148,3 @@ def test_handle_model_failed_unknown_node(caplog, monkeypatch):
|
||||
))
|
||||
|
||||
assert any("unknown node" in rec.message for rec in caplog.records)
|
||||
|
||||
|
||||
# ---------- 4. select_node() triggers swap when model mismatched ----------
|
||||
|
||||
|
||||
def test_select_node_code_triggers_swap(monkeypatch):
|
||||
_reset()
|
||||
_published.clear()
|
||||
monkeypatch.setattr(cluster, "publish", _fake_publish)
|
||||
|
||||
cluster.CLUSTER_NODES["jarvis"] = {
|
||||
"name": "jarvis", "type": "worker", "status": "active",
|
||||
"ip": "192.168.50.210",
|
||||
"active_model": {"name": "llama3.1", "port": 8081},
|
||||
"inventory": [
|
||||
{"filename": "qwen2.5-coder-14b-Q4_K_M.gguf", "name": "qwen2.5-coder", "version": "14b", "quant": "Q4_K_M"},
|
||||
],
|
||||
}
|
||||
|
||||
result = asyncio.run(triage.select_node("code"))
|
||||
|
||||
assert result is None
|
||||
# Swap should have been published
|
||||
assert any("cmd.swap_model" in rk for _, rk, _ in _published)
|
||||
# Node should now be swapping
|
||||
assert cluster.CLUSTER_NODES["jarvis"]["status"] == "swapping"
|
||||
|
||||
|
||||
# ---------- 5. select_node() returns None when node is already swapping ----------
|
||||
|
||||
|
||||
def test_select_node_swapping_returns_none(monkeypatch):
|
||||
_reset()
|
||||
_published.clear()
|
||||
monkeypatch.setattr(cluster, "publish", _fake_publish)
|
||||
|
||||
cluster.CLUSTER_NODES["jarvis"] = {
|
||||
"name": "jarvis", "type": "worker", "status": "swapping",
|
||||
"ip": "192.168.50.210",
|
||||
"active_model": {"name": "llama3.1", "port": 8081},
|
||||
"inventory": [
|
||||
{"filename": "qwen2.5-coder-14b-Q4_K_M.gguf", "name": "qwen2.5-coder"},
|
||||
],
|
||||
}
|
||||
|
||||
result = asyncio.run(triage.select_node("code"))
|
||||
|
||||
assert result is None
|
||||
# No swap command should be published while already swapping
|
||||
swap_published = any("cmd.swap_model" in rk for _, rk, _ in _published)
|
||||
assert not swap_published
|
||||
|
||||
@@ -1,141 +0,0 @@
|
||||
"""Tests for triage.py — query classification and node selection."""
|
||||
import asyncio
|
||||
import json
|
||||
|
||||
import httpx
|
||||
|
||||
import cluster
|
||||
import config
|
||||
import triage
|
||||
|
||||
|
||||
def _reset():
|
||||
cluster.CLUSTER_NODES.clear()
|
||||
cluster.CLUSTER_COORDINATOR = None
|
||||
|
||||
|
||||
_published = []
|
||||
|
||||
|
||||
async def _fake_publish(exchange, routing_key, payload):
|
||||
_published.append((exchange, routing_key, payload))
|
||||
|
||||
|
||||
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 = asyncio.run(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 = asyncio.run(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
|
||||
@@ -1,98 +0,0 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user