From 9adc25ad86513d372235ae570f12b93b609b0ea2 Mon Sep 17 00:00:00 2001 From: Brad Stein Date: Wed, 28 Jan 2026 12:57:50 -0300 Subject: [PATCH] atlasbot: add quick/smart bot config --- atlasbot/config.py | 34 ++++++++++++++++++++++++++++++++++ atlasbot/llm/prompts.py | 3 ++- atlasbot/main.py | 4 ++-- atlasbot/matrix/bot.py | 30 +++++++++++++++++------------- tests/test_engine.py | 1 + 5 files changed, 56 insertions(+), 16 deletions(-) diff --git a/atlasbot/config.py b/atlasbot/config.py index 3519cb8..6bc0c28 100644 --- a/atlasbot/config.py +++ b/atlasbot/config.py @@ -31,6 +31,7 @@ class Settings: room_alias: str server_name: str bot_mentions: tuple[str, ...] + matrix_bots: tuple["MatrixBotConfig", ...] ollama_url: str ollama_model: str @@ -64,6 +65,37 @@ class Settings: smart_max_candidates: int +@dataclass(frozen=True) +class MatrixBotConfig: + username: str + password: str + mentions: tuple[str, ...] + mode: str + + +def _load_matrix_bots(bot_mentions: tuple[str, ...]) -> tuple[MatrixBotConfig, ...]: + bots: list[MatrixBotConfig] = [] + quick_user = os.getenv("BOT_USER_QUICK", "").strip() + quick_pass = os.getenv("BOT_PASS_QUICK", "").strip() + if quick_user and quick_pass: + bots.append(MatrixBotConfig(quick_user, quick_pass, (quick_user,), "quick")) + + smart_user = os.getenv("BOT_USER_SMART", "").strip() + smart_pass = os.getenv("BOT_PASS_SMART", "").strip() + if smart_user and smart_pass: + bots.append(MatrixBotConfig(smart_user, smart_pass, (smart_user,), "smart")) + + if bots: + return tuple(bots) + + legacy_user = os.getenv("BOT_USER", "").strip() + legacy_pass = os.getenv("BOT_PASS", "").strip() + if legacy_user and legacy_pass: + mentions = bot_mentions or (legacy_user,) + bots.append(MatrixBotConfig(legacy_user, legacy_pass, mentions, "")) + return tuple(bots) + + def load_settings() -> Settings: bot_mentions = tuple( @@ -73,6 +105,7 @@ def load_settings() -> Settings: if item.strip() ] ) + matrix_bots = _load_matrix_bots(bot_mentions) return Settings( matrix_base=os.getenv("MATRIX_BASE", "http://othrys-synapse-matrix-synapse:8008"), auth_base=os.getenv("AUTH_BASE", "http://matrix-authentication-service:8080"), @@ -81,6 +114,7 @@ def load_settings() -> Settings: room_alias=os.getenv("ROOM_ALIAS", "#othrys:live.bstein.dev"), server_name=os.getenv("MATRIX_SERVER_NAME", "live.bstein.dev"), bot_mentions=bot_mentions, + matrix_bots=matrix_bots, ollama_url=os.getenv("OLLAMA_URL", "http://ollama.ai.svc.cluster.local:11434"), ollama_model=os.getenv("OLLAMA_MODEL", "qwen2.5:14b-instruct"), ollama_model_fast=os.getenv("ATLASBOT_MODEL_FAST", "qwen2.5:14b-instruct"), diff --git a/atlasbot/llm/prompts.py b/atlasbot/llm/prompts.py index 0b01cf7..654f2e5 100644 --- a/atlasbot/llm/prompts.py +++ b/atlasbot/llm/prompts.py @@ -32,7 +32,8 @@ SCORE_PROMPT = ( SYNTHESIZE_PROMPT = ( "Synthesize a final response from the best candidates. " "Use a natural, helpful tone with light reasoning. " - "Avoid lists unless the user asked for lists." + "Avoid lists unless the user asked for lists. " + "Do not include confidence scores or evaluation metadata." ) STOCK_SYSTEM = ( diff --git a/atlasbot/main.py b/atlasbot/main.py index c800391..203f862 100644 --- a/atlasbot/main.py +++ b/atlasbot/main.py @@ -47,8 +47,8 @@ async def main() -> None: server = uvicorn.Server(uvicorn.Config(api.app, host="0.0.0.0", port=settings.http_port, log_level="info")) tasks = [] - if settings.bot_user and settings.bot_pass: - tasks.append(asyncio.create_task(MatrixBot(settings, engine, answer_handler).run())) + for bot in settings.matrix_bots: + tasks.append(asyncio.create_task(MatrixBot(settings, bot, engine, answer_handler).run())) tasks.append(asyncio.create_task(server.serve())) await asyncio.gather(*tasks) diff --git a/atlasbot/matrix/bot.py b/atlasbot/matrix/bot.py index 147ed67..00e99b9 100644 --- a/atlasbot/matrix/bot.py +++ b/atlasbot/matrix/bot.py @@ -6,7 +6,7 @@ from urllib.parse import quote import httpx -from atlasbot.config import Settings +from atlasbot.config import MatrixBotConfig, Settings from collections.abc import Awaitable, Callable from atlasbot.engine.answerer import AnswerEngine, AnswerResult @@ -15,14 +15,15 @@ log = logging.getLogger(__name__) class MatrixClient: - def __init__(self, settings: Settings) -> None: + def __init__(self, settings: Settings, bot: MatrixBotConfig) -> None: self._settings = settings + self._bot = bot async def login(self) -> str: payload = { "type": "m.login.password", - "identifier": {"type": "m.id.user", "user": self._settings.bot_user}, - "password": self._settings.bot_pass, + "identifier": {"type": "m.id.user", "user": self._bot.username}, + "password": self._bot.password, } url = f"{self._settings.auth_base}/_matrix/client/v3/login" async with httpx.AsyncClient(timeout=15.0) as client: @@ -77,12 +78,14 @@ class MatrixBot: def __init__( self, settings: Settings, + bot: MatrixBotConfig, engine: AnswerEngine, answer_handler: Callable[[str, str, Callable[[str, str], None] | None], Awaitable[AnswerResult]] | None = None, ) -> None: self._settings = settings + self._bot = bot self._engine = engine - self._client = MatrixClient(settings) + self._client = MatrixClient(settings, bot) self._answer_handler = answer_handler async def run(self) -> None: @@ -121,9 +124,9 @@ class MatrixBot: content = event.get("content") or {} body = content.get("body") or "" sender = event.get("sender") or "" - if sender.endswith(f"/{self._settings.bot_user}") or sender == self._settings.bot_user: + if sender.endswith(f"/{self._bot.username}") or sender == self._bot.username: continue - mode, question = _extract_mode(body, self._settings.bot_mentions) + mode, question = _extract_mode(body, self._bot.mentions, self._bot.mode) if not question: continue await self._client.send_message(token, room_id, "Thinking…") @@ -169,15 +172,16 @@ class MatrixBot: task.cancel() -def _extract_mode(body: str, mentions: tuple[str, ...]) -> tuple[str, str]: +def _extract_mode(body: str, mentions: tuple[str, ...], default_mode: str) -> tuple[str, str]: lower = body.lower() for mention in mentions: if mention and mention.lower() in lower: - mode = "quick" - if "atlas-smart" in lower or "smart" in lower: - mode = "smart" - if "atlas-quick" in lower or "quick" in lower: - mode = "quick" + mode = default_mode or "quick" + if not default_mode: + if "atlas-smart" in lower or "smart" in lower: + mode = "smart" + if "atlas-quick" in lower or "quick" in lower: + mode = "quick" cleaned = body for tag in mentions: cleaned = cleaned.replace(tag, "") diff --git a/tests/test_engine.py b/tests/test_engine.py index e65bf33..f825ab9 100644 --- a/tests/test_engine.py +++ b/tests/test_engine.py @@ -25,6 +25,7 @@ def _settings() -> Settings: room_alias="", server_name="", bot_mentions=(), + matrix_bots=(), ollama_url="", ollama_model="base", ollama_model_fast="fast",