atlasbot/atlasbot/config.py

179 lines
6.8 KiB
Python

import os
from dataclasses import dataclass
def _env_bool(name: str, default: str = "false") -> bool:
value = os.getenv(name, default).strip().lower()
return value in {"1", "true", "yes", "y", "on"}
def _env_int(name: str, default: str) -> int:
raw = os.getenv(name, default)
try:
return int(raw)
except (TypeError, ValueError):
return int(default)
def _env_float(name: str, default: str) -> float:
raw = os.getenv(name, default)
try:
return float(raw)
except (TypeError, ValueError):
return float(default)
@dataclass(frozen=True)
class Settings:
matrix_base: str
auth_base: str
bot_user: str
bot_pass: str
room_alias: str
server_name: str
bot_mentions: tuple[str, ...]
matrix_bots: tuple["MatrixBotConfig", ...]
ollama_url: str
ollama_model: str
ollama_model_fast: str
ollama_model_smart: str
ollama_model_genius: str
ollama_fallback_model: str
ollama_timeout_sec: float
ollama_retries: int
ollama_api_key: str
http_port: int
internal_token: str
kb_dir: str
vm_url: str
ariadne_state_url: str
ariadne_state_token: str
snapshot_ttl_sec: int
thinking_interval_sec: int
conversation_ttl_sec: int
snapshot_pin_enabled: bool
queue_enabled: bool
nats_url: str
nats_stream: str
nats_subject: str
nats_result_bucket: str
fast_max_angles: int
smart_max_angles: int
genius_max_angles: int
fast_max_candidates: int
smart_max_candidates: int
genius_max_candidates: int
fast_llm_calls_max: int
smart_llm_calls_max: int
genius_llm_calls_max: int
llm_limit_multiplier: float
state_db_path: str = "/data/atlasbot_state.db"
nats_timeout_fast_sec: int = 180
nats_timeout_smart_sec: int = 480
nats_timeout_genius_sec: int = 900
debug_pipeline: bool = False
@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"))
genius_user = os.getenv("BOT_USER_GENIUS", "").strip()
genius_pass = os.getenv("BOT_PASS_GENIUS", "").strip()
if genius_user and genius_pass:
bots.append(MatrixBotConfig(genius_user, genius_pass, (genius_user,), "genius"))
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(
[
item.strip()
for item in os.getenv("BOT_MENTIONS", "atlasbot,atlas-quick,atlas-smart,atlas-genius").split(",")
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"),
bot_user=os.getenv("BOT_USER", ""),
bot_pass=os.getenv("BOT_PASS", ""),
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"),
ollama_model_smart=os.getenv("ATLASBOT_MODEL_SMART", "qwen2.5:14b-instruct"),
ollama_model_genius=os.getenv("ATLASBOT_MODEL_GENIUS", "qwen2.5:14b-instruct"),
ollama_fallback_model=os.getenv("OLLAMA_FALLBACK_MODEL", ""),
ollama_timeout_sec=_env_float("OLLAMA_TIMEOUT_SEC", "480"),
ollama_retries=_env_int("OLLAMA_RETRIES", "1"),
ollama_api_key=os.getenv("CHAT_API_KEY", ""),
http_port=_env_int("ATLASBOT_HTTP_PORT", "8090"),
internal_token=os.getenv("ATLASBOT_INTERNAL_TOKEN", "")
or os.getenv("CHAT_API_HOMEPAGE", ""),
kb_dir=os.getenv("KB_DIR", ""),
vm_url=os.getenv("VM_URL", "http://victoria-metrics-single-server.monitoring.svc.cluster.local:8428"),
ariadne_state_url=os.getenv("ARIADNE_STATE_URL", ""),
ariadne_state_token=os.getenv("ARIADNE_STATE_TOKEN", ""),
snapshot_ttl_sec=_env_int("ATLASBOT_SNAPSHOT_TTL_SEC", "30"),
thinking_interval_sec=_env_int("ATLASBOT_THINKING_INTERVAL_SEC", "30"),
conversation_ttl_sec=_env_int("ATLASBOT_CONVERSATION_TTL_SEC", "900"),
snapshot_pin_enabled=_env_bool("ATLASBOT_SNAPSHOT_PIN_ENABLED", "false"),
state_db_path=os.getenv("ATLASBOT_STATE_DB", "/data/atlasbot_state.db"),
queue_enabled=_env_bool("ATLASBOT_QUEUE_ENABLED", "false"),
nats_url=os.getenv("ATLASBOT_NATS_URL", "nats://nats.nats.svc.cluster.local:4222"),
nats_stream=os.getenv("ATLASBOT_NATS_STREAM", "atlasbot"),
nats_subject=os.getenv("ATLASBOT_NATS_SUBJECT", "atlasbot.requests"),
nats_result_bucket=os.getenv("ATLASBOT_NATS_RESULTS", "atlasbot_results"),
nats_timeout_fast_sec=_env_int("ATLASBOT_NATS_TIMEOUT_FAST_SEC", "180"),
nats_timeout_smart_sec=_env_int("ATLASBOT_NATS_TIMEOUT_SMART_SEC", "480"),
nats_timeout_genius_sec=_env_int("ATLASBOT_NATS_TIMEOUT_GENIUS_SEC", "900"),
fast_max_angles=_env_int("ATLASBOT_FAST_MAX_ANGLES", "2"),
smart_max_angles=_env_int("ATLASBOT_SMART_MAX_ANGLES", "5"),
genius_max_angles=_env_int("ATLASBOT_GENIUS_MAX_ANGLES", "9"),
fast_max_candidates=_env_int("ATLASBOT_FAST_MAX_CANDIDATES", "2"),
smart_max_candidates=_env_int("ATLASBOT_SMART_MAX_CANDIDATES", "6"),
genius_max_candidates=_env_int("ATLASBOT_GENIUS_MAX_CANDIDATES", "10"),
fast_llm_calls_max=_env_int("ATLASBOT_FAST_LLM_CALLS_MAX", "24"),
smart_llm_calls_max=_env_int("ATLASBOT_SMART_LLM_CALLS_MAX", "48"),
genius_llm_calls_max=_env_int("ATLASBOT_GENIUS_LLM_CALLS_MAX", "96"),
llm_limit_multiplier=_env_float("ATLASBOT_LLM_LIMIT_MULTIPLIER", "1.5"),
debug_pipeline=_env_bool("ATLASBOT_DEBUG_PIPELINE", "false"),
)