atlasbot/tests/test_engine.py

import asyncio

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, replies: list[str]) -> None:
        self._replies = replies
        self.calls: list[str] = []

    async def chat(self, messages, *, model=None):
        self.calls.append(model or "")
        return self._replies.pop(0)


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,
        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,
    )


def test_engine_answer_basic():
    llm = FakeLLM(
        [
            '{"needs_snapshot": true}',
            '[{"name":"primary","question":"What is Atlas?","relevance":90}]',
            "Based on the snapshot, Atlas has 22 nodes.",
            '{"confidence":80,"relevance":90,"satisfaction":85,"hallucination_risk":"low"}',
            "Atlas has 22 nodes and is healthy.",
        ]
    )
    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