89 lines
2.7 KiB
Python
89 lines
2.7 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) -> None:
|
|
self.calls: list[str] = []
|
|
|
|
async def chat(self, messages, *, model=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 "{}"
|
|
|
|
|
|
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,
|
|
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,
|
|
)
|
|
|
|
|
|
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
|