165 lines
6.2 KiB
Python
165 lines
6.2 KiB
Python
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import os
|
|
from pathlib import Path
|
|
from unittest import TestCase, mock
|
|
|
|
|
|
BOT_PATH = Path(__file__).resolve().parents[1] / "atlasbot" / "bot.py"
|
|
|
|
|
|
def load_bot_module():
|
|
env = {
|
|
"BOT_USER": "atlas-smart",
|
|
"BOT_PASS": "smart-pass",
|
|
"BOT_USER_QUICK": "atlas-quick",
|
|
"BOT_PASS_QUICK": "quick-pass",
|
|
"BOT_USER_SMART": "atlas-smart",
|
|
"BOT_PASS_SMART": "smart-pass",
|
|
"BOT_USER_GENIUS": "atlas-genius",
|
|
"BOT_PASS_GENIUS": "genius-pass",
|
|
"OLLAMA_URL": "http://ollama.invalid",
|
|
"OLLAMA_MODEL": "base-model",
|
|
"ATLASBOT_MODEL_FAST": "fast-model",
|
|
"ATLASBOT_MODEL_SMART": "smart-model",
|
|
"ATLASBOT_MODEL_GENIUS": "genius-model",
|
|
"ATLASBOT_QUICK_TIME_BUDGET_SEC": "15",
|
|
"ATLASBOT_SMART_TIME_BUDGET_SEC": "45",
|
|
"ATLASBOT_GENIUS_TIME_BUDGET_SEC": "180",
|
|
"KB_DIR": "",
|
|
"VM_URL": "http://vm.invalid",
|
|
"ARIADNE_STATE_URL": "",
|
|
"ARIADNE_STATE_TOKEN": "",
|
|
}
|
|
with mock.patch.dict(os.environ, env, clear=False):
|
|
spec = importlib.util.spec_from_file_location("atlasbot_bot", BOT_PATH)
|
|
module = importlib.util.module_from_spec(spec)
|
|
assert spec.loader is not None
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
class AtlasbotModeTests(TestCase):
|
|
def setUp(self):
|
|
self.bot = load_bot_module()
|
|
|
|
def test_bot_accounts_include_genius_mode(self):
|
|
accounts = self.bot._bot_accounts()
|
|
by_user = {account["user"]: account["mode"] for account in accounts}
|
|
|
|
self.assertEqual(by_user["atlas-quick"], "fast")
|
|
self.assertEqual(by_user["atlas-smart"], "smart")
|
|
self.assertEqual(by_user["atlas-genius"], "genius")
|
|
|
|
def test_objective_cluster_question_uses_fact_pack_without_llm(self):
|
|
fact_lines = [
|
|
"hottest_cpu: longhorn-system (6.69)",
|
|
"hottest_ram: longhorn-system (36.05 GB)",
|
|
]
|
|
|
|
with (
|
|
mock.patch.object(self.bot, "_fact_pack_lines", return_value=fact_lines),
|
|
mock.patch.object(self.bot, "_ollama_call", side_effect=AssertionError("LLM should not be called")),
|
|
):
|
|
reply = self.bot.open_ended_answer(
|
|
"what is the hottest cpu node in titan lab currently?",
|
|
inventory=[],
|
|
snapshot=None,
|
|
workloads=[],
|
|
history_lines=[],
|
|
mode="smart",
|
|
allow_tools=True,
|
|
)
|
|
|
|
self.assertIn("longhorn-system", reply)
|
|
self.assertIn("Confidence:", reply)
|
|
|
|
def test_subjective_genius_answer_uses_genius_model(self):
|
|
fact_lines = [
|
|
"hottest_cpu: longhorn-system (6.69)",
|
|
"worker_nodes: titan-01, titan-02, titan-03",
|
|
]
|
|
captured: dict[str, object] = {}
|
|
|
|
def fake_ollama_call(hist_key, prompt, *, context, use_history=True, system_override=None, model=None, timeout=None):
|
|
captured["model"] = model
|
|
captured["timeout"] = timeout
|
|
captured["context"] = context
|
|
return "The worker spread stands out because Titan keeps meaningful capacity on the same cluster. Confidence: high"
|
|
|
|
with (
|
|
mock.patch.object(self.bot, "_fact_pack_lines", return_value=fact_lines),
|
|
mock.patch.object(self.bot, "_ollama_call", side_effect=fake_ollama_call),
|
|
):
|
|
reply = self.bot.open_ended_answer(
|
|
"what stands out about titan lab?",
|
|
inventory=[],
|
|
snapshot=None,
|
|
workloads=[],
|
|
history_lines=[],
|
|
mode="genius",
|
|
allow_tools=True,
|
|
context='Cluster snapshot (JSON): {"injected":true}',
|
|
)
|
|
|
|
self.assertIn("The worker spread stands out", reply)
|
|
self.assertEqual(captured["model"], "genius-model")
|
|
self.assertLessEqual(float(captured["timeout"]), 180.0)
|
|
self.assertIn('Cluster snapshot (JSON): {"injected":true}', str(captured["context"]))
|
|
|
|
def test_mode_timeouts_stay_within_budgets(self):
|
|
fact_lines = [
|
|
"hottest_cpu: longhorn-system (6.69)",
|
|
"worker_nodes: titan-01, titan-02, titan-03",
|
|
]
|
|
seen: list[tuple[str, float]] = []
|
|
|
|
def fake_ollama_call(hist_key, prompt, *, context, use_history=True, system_override=None, model=None, timeout=None):
|
|
seen.append((str(model), float(timeout or 0)))
|
|
return "Atlas has a clear standout because the worker spread is healthy. Confidence: high"
|
|
|
|
with (
|
|
mock.patch.object(self.bot, "_fact_pack_lines", return_value=fact_lines),
|
|
mock.patch.object(self.bot, "_ollama_call", side_effect=fake_ollama_call),
|
|
):
|
|
for mode in ("fast", "smart", "genius"):
|
|
reply = self.bot.open_ended_answer(
|
|
"what stands out about titan lab?",
|
|
inventory=[],
|
|
snapshot=None,
|
|
workloads=[],
|
|
history_lines=[],
|
|
mode=mode,
|
|
allow_tools=True,
|
|
)
|
|
self.assertIn("Confidence:", reply)
|
|
|
|
self.assertEqual([model for model, _ in seen], ["fast-model", "smart-model", "genius-model"])
|
|
self.assertLessEqual(seen[0][1], 15.0)
|
|
self.assertLessEqual(seen[1][1], 45.0)
|
|
self.assertLessEqual(seen[2][1], 180.0)
|
|
|
|
def test_llm_timeout_still_returns_a_conclusion(self):
|
|
fact_lines = [
|
|
"worker_nodes: titan-01, titan-02, titan-03",
|
|
"hottest_cpu: longhorn-system (6.69)",
|
|
]
|
|
|
|
with (
|
|
mock.patch.object(self.bot, "_fact_pack_lines", return_value=fact_lines),
|
|
mock.patch.object(self.bot, "_ollama_call", side_effect=TimeoutError("simulated timeout")),
|
|
):
|
|
reply = self.bot.open_ended_answer(
|
|
"what stands out about the worker nodes?",
|
|
inventory=[],
|
|
snapshot=None,
|
|
workloads=[],
|
|
history_lines=[],
|
|
mode="genius",
|
|
allow_tools=True,
|
|
)
|
|
|
|
self.assertIn("worker nodes", reply.lower())
|
|
self.assertIn("Confidence:", reply)
|