74 lines
2.0 KiB
Python
74 lines
2.0 KiB
Python
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_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,
|
|
fast_max_candidates=1,
|
|
smart_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
|