atlasbot/tests/test_engine.py

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 atlasbot.config import Settings


class FakeLLM:
    def __init__(self) -> None:
        self.calls: list[str] = []

    async def chat(self, messages, *, model=None, timeout_sec=None):
        prompt = messages[-1]["content"]
        self.calls.append(prompt)
        if "normalized" in prompt and "keywords" in prompt:
            return '{"normalized":"What is Atlas?","keywords":["atlas"]}'
        if "needs_snapshot" in prompt:
            return '{"needs_snapshot": true, "answer_style": "direct"}'
        if "sub-questions" in prompt:
            return '[{"id":"q1","question":"What is Atlas?","priority":1}]'
        if "sub-question" in prompt:
            return "Atlas has 22 nodes."
        if "final response" in prompt:
            return "Atlas has 22 nodes."
        if "Score response quality" in prompt:
            return '{"confidence":80,"relevance":90,"satisfaction":85,"hallucination_risk":"low"}'
        if "claims list" in prompt:
            return '{"claims": []}'
        return "{}"


class SlowFakeLLM(FakeLLM):
    async def chat(self, messages, *, model=None, timeout_sec=None):
        await asyncio.sleep(0.02)
        return await super().chat(messages, model=model, timeout_sec=timeout_sec)


def _settings() -> Settings:
    return Settings(
        matrix_base="",
        auth_base="",
        bot_user="",
        bot_pass="",
        room_alias="",
        server_name="",
        bot_mentions=(),
        matrix_bots=(),
        ollama_url="",
        ollama_model="base",
        ollama_model_fast="fast",
        ollama_model_smart="smart",
        ollama_model_genius="genius",
        ollama_fallback_model="",
        ollama_timeout_sec=1.0,
        ollama_retries=0,
        ollama_api_key="",
        http_port=8090,
        internal_token="",
        kb_dir="",
        vm_url="",
        ariadne_state_url="",
        ariadne_state_token="",
        snapshot_ttl_sec=30,
        thinking_interval_sec=30,
        quick_time_budget_sec=15.0,
        smart_time_budget_sec=45.0,
        genius_time_budget_sec=180.0,
        conversation_ttl_sec=300,
        snapshot_pin_enabled=False,
        queue_enabled=False,
        nats_url="",
        nats_stream="",
        nats_subject="",
        nats_result_bucket="",
        fast_max_angles=1,
        smart_max_angles=1,
        genius_max_angles=1,
        fast_max_candidates=1,
        smart_max_candidates=1,
        genius_max_candidates=1,
        fast_llm_calls_max=9,
        smart_llm_calls_max=17,
        genius_llm_calls_max=32,
        llm_limit_multiplier=1.5,
        state_db_path="/tmp/atlasbot_test_state.db",
    )


def test_engine_answer_basic():
    llm = FakeLLM()
    settings = _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_quick_mode_time_budget_guard():
    llm = SlowFakeLLM()
    settings = replace(_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