80 lines
2.7 KiB
Python
80 lines
2.7 KiB
Python
"""Answer-engine regression tests."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from dataclasses import replace
|
|
|
|
from atlasbot.engine.answerer import AnswerEngine
|
|
from atlasbot.knowledge.loader import KnowledgeBase
|
|
from atlasbot.snapshot.builder import SnapshotProvider
|
|
from testing.fakes import FakeLLM, SlowFakeLLM, build_test_settings
|
|
|
|
|
|
def test_engine_answer_basic() -> None:
|
|
"""The quick path should answer from the fact sheet."""
|
|
|
|
llm = FakeLLM()
|
|
settings = build_test_settings()
|
|
kb = KnowledgeBase("")
|
|
snapshot = SnapshotProvider(settings)
|
|
engine = AnswerEngine(settings, llm, kb, snapshot)
|
|
|
|
result = asyncio.run(engine.answer("What is Atlas?", mode="quick"))
|
|
assert "Atlas has 22 nodes" in result.reply
|
|
|
|
|
|
def test_smart_mode_uses_factsheet_path() -> None:
|
|
"""Smart mode should stay on the factsheet branch for direct cluster questions."""
|
|
|
|
llm = FakeLLM()
|
|
settings = build_test_settings()
|
|
kb = KnowledgeBase("")
|
|
snapshot = SnapshotProvider(settings)
|
|
engine = AnswerEngine(settings, llm, kb, snapshot)
|
|
|
|
result = asyncio.run(engine.answer("What is the most demanding system in titan lab currently?", mode="smart"))
|
|
assert "Atlas has 22 nodes" in result.reply
|
|
assert "time budget" not in result.reply.lower()
|
|
|
|
|
|
def test_genius_mode_uses_factsheet_path() -> None:
|
|
"""Genius mode should also return the factsheet answer for the same query."""
|
|
|
|
llm = FakeLLM()
|
|
settings = build_test_settings()
|
|
kb = KnowledgeBase("")
|
|
snapshot = SnapshotProvider(settings)
|
|
engine = AnswerEngine(settings, llm, kb, snapshot)
|
|
|
|
result = asyncio.run(engine.answer("What is the most demanding system in titan lab currently?", mode="genius"))
|
|
assert "Atlas has 22 nodes" in result.reply
|
|
assert "time budget" not in result.reply.lower()
|
|
|
|
|
|
def test_plain_math_question_is_rejected_for_cluster_modes() -> None:
|
|
"""The bot should keep users on cluster questions instead of generic math."""
|
|
|
|
llm = FakeLLM()
|
|
settings = build_test_settings()
|
|
kb = KnowledgeBase("")
|
|
snapshot = SnapshotProvider(settings)
|
|
engine = AnswerEngine(settings, llm, kb, snapshot)
|
|
|
|
result = asyncio.run(engine.answer("what is 2+2?", mode="quick"))
|
|
assert "focus on Titan cluster operations" in result.reply
|
|
|
|
|
|
def test_quick_mode_time_budget_guard() -> None:
|
|
"""A slow model call should trip the quick-mode budget guard."""
|
|
|
|
llm = SlowFakeLLM()
|
|
settings = replace(build_test_settings(), quick_time_budget_sec=0.01)
|
|
kb = KnowledgeBase("")
|
|
snapshot = SnapshotProvider(settings)
|
|
engine = AnswerEngine(settings, llm, kb, snapshot)
|
|
|
|
result = asyncio.run(engine.answer("What is Atlas?", mode="quick"))
|
|
assert "time budget" in result.reply
|
|
assert result.meta.get("time_budget_hit") is True
|