bstein/atlasbot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

atlasbot: add quick/smart bot config

Browse Source
This commit is contained in:
Brad Stein 2026-01-28 12:57:50 -03:00
parent fcca4107db
commit 9adc25ad86
5 changed files with 56 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
atlasbot/config.py
View File

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

3
atlasbot/llm/prompts.py
View File

@ -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 = (

4
atlasbot/main.py
View File

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

30
atlasbot/matrix/bot.py
View File

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

1
tests/test_engine.py
View File

@ -25,6 +25,7 @@ def _settings() -> Settings:
room_alias="",
server_name="",
bot_mentions=(),
matrix_bots=(),
ollama_url="",
ollama_model="base",
ollama_model_fast="fast",