diff --git a/assistant_bot/application.py b/assistant_bot/application.py index 898a20f..86929f0 100644 --- a/assistant_bot/application.py +++ b/assistant_bot/application.py @@ -12,22 +12,7 @@ from telegram.ext import ( from .ai import AIClient from .auth import authentication_guard -from .config import ( - get_assistant_mode, - get_assistant_password, - get_bot_token, - get_db_path, - get_local_timezone, - get_ollama_base_url, - get_voice_max_duration_seconds, - get_whisper_compute_type, - get_whisper_device, - get_whisper_language, - get_whisper_model, - get_yandex_cloud_folder, - get_yandex_stt_language, - get_yandex_stt_model, -) +from .config import AppSettings, load_app_settings from .handlers import ( ask_command, help_command, @@ -42,6 +27,7 @@ from .handlers import ( ) from .jobs import post_init, post_shutdown from .ollama import OllamaClient +from .services import SERVICES_KEY, ApplicationServices from .storage import AssistantStorage from .speech import SpeechRecognizer, SpeechTranscriber from .yandex_ai import YandexAIClient, YandexSpeechRecognizer @@ -55,42 +41,51 @@ def configure_logging() -> None: logging.getLogger("httpx").setLevel(logging.WARNING) -def create_application() -> Application: - assistant_password = get_assistant_password() - storage = AssistantStorage(get_db_path()) - mode = get_assistant_mode() +def create_services(settings: AppSettings) -> ApplicationServices: + storage = AssistantStorage(settings.db_path) + ai_client: AIClient + speech_recognizer: SpeechTranscriber + if settings.mode == "local": + ai_client = OllamaClient(settings.ollama_base_url) + speech_recognizer = SpeechRecognizer( + model_name=settings.whisper_model, + device=settings.whisper_device, + compute_type=settings.whisper_compute_type, + language=settings.whisper_language, + ) + else: + folder_id = settings.yandex_cloud_folder + if folder_id is None: + raise RuntimeError( + "YANDEX_CLOUD_FOLDER is required in yandex mode." + ) + ai_client = YandexAIClient(folder_id=folder_id) + speech_recognizer = YandexSpeechRecognizer( + folder_id=folder_id, + language=settings.yandex_stt_language, + model=settings.yandex_stt_model, + ) + + return ApplicationServices( + storage=storage, + ai_client=ai_client, + timezone=settings.timezone, + speech_recognizer=speech_recognizer, + voice_max_duration_seconds=settings.voice_max_duration_seconds, + assistant_password=settings.assistant_password, + ) + + +def create_application(settings: AppSettings | None = None) -> Application: + settings = settings or load_app_settings() application = ( Application.builder() - .token(get_bot_token()) + .token(settings.bot_token) .post_init(post_init) .post_shutdown(post_shutdown) .build() ) - application.bot_data["storage"] = storage - application.bot_data["assistant_password"] = assistant_password - ai_client: AIClient - speech_recognizer: SpeechTranscriber - if mode == "local": - ai_client = OllamaClient(get_ollama_base_url()) - speech_recognizer = SpeechRecognizer( - model_name=get_whisper_model(), - device=get_whisper_device(), - compute_type=get_whisper_compute_type(), - language=get_whisper_language(), - ) - else: - folder_id = get_yandex_cloud_folder() - ai_client = YandexAIClient(folder_id=folder_id) - speech_recognizer = YandexSpeechRecognizer( - folder_id=folder_id, - language=get_yandex_stt_language(), - model=get_yandex_stt_model(), - ) - - application.bot_data["ai_client"] = ai_client - application.bot_data["timezone"] = get_local_timezone() - application.bot_data["speech_recognizer"] = speech_recognizer - application.bot_data["voice_max_duration"] = get_voice_max_duration_seconds() + application.bot_data[SERVICES_KEY] = create_services(settings) application.add_handler(TypeHandler(Update, authentication_guard), group=-1) application.add_handler(CommandHandler("start", start)) diff --git a/assistant_bot/auth.py b/assistant_bot/auth.py index 9da1ddd..cd46213 100644 --- a/assistant_bot/auth.py +++ b/assistant_bot/auth.py @@ -4,7 +4,12 @@ from telegram import Update from telegram.constants import ChatType from telegram.ext import ApplicationHandlerStop, ContextTypes -from .telegram_utils import get_storage, reply_markdown, require_user_id +from .telegram_utils import ( + get_services, + get_storage, + reply_markdown, + require_user_id, +) def passwords_match(candidate: str, expected: str) -> bool: @@ -37,7 +42,7 @@ async def authentication_guard( raise ApplicationHandlerStop candidate = message.text.strip() if message and message.text else "" - password = str(context.application.bot_data["assistant_password"]) + password = get_services(context).assistant_password if candidate and passwords_match(candidate, password): storage.authorize_user(user_id) await reply_markdown( diff --git a/assistant_bot/config.py b/assistant_bot/config.py index 0c292ca..d2b93e6 100644 --- a/assistant_bot/config.py +++ b/assistant_bot/config.py @@ -1,5 +1,6 @@ import logging import os +from dataclasses import dataclass from pathlib import Path from zoneinfo import ZoneInfo, ZoneInfoNotFoundError @@ -46,6 +47,26 @@ MAX_AGENT_STEPS = 5 CONVERSATION_HISTORY_LIMIT = 12 +@dataclass(frozen=True) +class AppSettings: + bot_token: str + assistant_password: str + db_path: Path + mode: str + timezone: ZoneInfo + voice_max_duration_seconds: int + ollama_base_url: str + ollama_model: str + yandex_cloud_folder: str | None + yandex_cloud_model: str + yandex_stt_model: str + yandex_stt_language: str + whisper_model: str + whisper_device: str + whisper_compute_type: str + whisper_language: str | None + + def load_env_file(path: Path = ENV_FILE) -> None: if not path.exists(): return @@ -68,6 +89,128 @@ def load_env_file(path: Path = ENV_FILE) -> None: os.environ[key] = value +def load_app_settings() -> AppSettings: + """Load and validate the complete startup configuration once.""" + load_env_file() + + token = os.getenv(TOKEN_ENV_NAME) + if not token: + raise RuntimeError( + f"Set {TOKEN_ENV_NAME} in environment or .env file." + ) + + password = os.getenv(PASSWORD_ENV_NAME, "").strip() + if not password: + raise RuntimeError( + f"Set {PASSWORD_ENV_NAME} in environment or .env file." + ) + + mode = os.getenv(MODE_ENV_NAME, DEFAULT_MODE).strip().lower() + if mode not in {"local", "yandex"}: + raise RuntimeError( + f"{MODE_ENV_NAME} must be either 'local' or 'yandex', got {mode!r}." + ) + + raw_db_path = os.getenv(DB_ENV_NAME) + if raw_db_path: + db_path = Path(raw_db_path).expanduser() + if not db_path.is_absolute(): + db_path = PROJECT_ROOT / db_path + else: + db_path = DEFAULT_DB_FILE + + timezone_name = os.getenv(TIMEZONE_ENV_NAME, DEFAULT_TIMEZONE) + try: + local_timezone = ZoneInfo(timezone_name) + except ZoneInfoNotFoundError: + logger.warning( + "Unknown timezone %s, falling back to UTC", + timezone_name, + ) + local_timezone = ZoneInfo("UTC") + + raw_voice_duration = os.getenv( + VOICE_MAX_DURATION_ENV_NAME, + str(DEFAULT_VOICE_MAX_DURATION_SECONDS), + ) + try: + voice_duration = int(raw_voice_duration) + except ValueError: + voice_duration = 0 + if voice_duration <= 0: + logger.warning( + "%s must be a positive integer, using %s", + VOICE_MAX_DURATION_ENV_NAME, + DEFAULT_VOICE_MAX_DURATION_SECONDS, + ) + voice_duration = DEFAULT_VOICE_MAX_DURATION_SECONDS + + yandex_folder = os.getenv( + YANDEX_CLOUD_FOLDER_ENV_NAME, + "", + ).strip().strip("/") + if mode == "yandex" and not yandex_folder: + raise RuntimeError( + f"Set {YANDEX_CLOUD_FOLDER_ENV_NAME} in environment or .env file." + ) + + yandex_model = os.getenv( + YANDEX_CLOUD_MODEL_ENV_NAME, + DEFAULT_YANDEX_CLOUD_MODEL, + ).strip().strip("/") + if not yandex_model: + yandex_model = DEFAULT_YANDEX_CLOUD_MODEL + + whisper_language = os.getenv( + WHISPER_LANGUAGE_ENV_NAME, + DEFAULT_WHISPER_LANGUAGE, + ).strip() + + return AppSettings( + bot_token=token, + assistant_password=password, + db_path=db_path, + mode=mode, + timezone=local_timezone, + voice_max_duration_seconds=voice_duration, + ollama_base_url=os.getenv( + OLLAMA_BASE_URL_ENV_NAME, + DEFAULT_OLLAMA_BASE_URL, + ).rstrip("/"), + ollama_model=os.getenv( + OLLAMA_MODEL_ENV_NAME, + DEFAULT_OLLAMA_MODEL, + ), + yandex_cloud_folder=yandex_folder or None, + yandex_cloud_model=yandex_model, + yandex_stt_model=os.getenv( + YANDEX_STT_MODEL_ENV_NAME, + DEFAULT_YANDEX_STT_MODEL, + ).strip(), + yandex_stt_language=os.getenv( + YANDEX_STT_LANGUAGE_ENV_NAME, + DEFAULT_YANDEX_STT_LANGUAGE, + ).strip(), + whisper_model=os.getenv( + WHISPER_MODEL_ENV_NAME, + DEFAULT_WHISPER_MODEL, + ).strip(), + whisper_device=os.getenv( + WHISPER_DEVICE_ENV_NAME, + DEFAULT_WHISPER_DEVICE, + ).strip(), + whisper_compute_type=os.getenv( + WHISPER_COMPUTE_TYPE_ENV_NAME, + DEFAULT_WHISPER_COMPUTE_TYPE, + ).strip(), + whisper_language=( + None + if whisper_language.lower() == "auto" + else whisper_language or None + ), + ) + + def get_bot_token() -> str: load_env_file() token = os.getenv(TOKEN_ENV_NAME) diff --git a/assistant_bot/jobs.py b/assistant_bot/jobs.py index 72598a5..d21e46c 100644 --- a/assistant_bot/jobs.py +++ b/assistant_bot/jobs.py @@ -6,7 +6,7 @@ from zoneinfo import ZoneInfo from telegram.ext import Application from .config import REMINDER_POLL_SECONDS -from .storage import AssistantStorage +from .services import SERVICES_KEY, ApplicationServices from .telegram_utils import escape_markdown_text, markdown_code, send_markdown from .time_utils import format_local_dt, utc_now @@ -15,8 +15,9 @@ logger = logging.getLogger(__name__) async def reminder_loop(application: Application) -> None: - storage: AssistantStorage = application.bot_data["storage"] - tz: ZoneInfo = application.bot_data["timezone"] + services: ApplicationServices = application.bot_data[SERVICES_KEY] + storage = services.storage + tz: ZoneInfo = services.timezone while True: try: diff --git a/assistant_bot/services.py b/assistant_bot/services.py new file mode 100644 index 0000000..aaba5ba --- /dev/null +++ b/assistant_bot/services.py @@ -0,0 +1,19 @@ +from dataclasses import dataclass +from zoneinfo import ZoneInfo + +from .ai import AIClient +from .speech import SpeechTranscriber +from .storage import AssistantStorage + + +SERVICES_KEY = "services" + + +@dataclass(frozen=True) +class ApplicationServices: + storage: AssistantStorage + ai_client: AIClient + timezone: ZoneInfo + speech_recognizer: SpeechTranscriber + voice_max_duration_seconds: int + assistant_password: str diff --git a/assistant_bot/telegram_utils.py b/assistant_bot/telegram_utils.py index f387091..227b457 100644 --- a/assistant_bot/telegram_utils.py +++ b/assistant_bot/telegram_utils.py @@ -20,6 +20,7 @@ from telegramify_markdown.config import RenderConfig from .ai import AIClient from .config import MAX_TELEGRAM_MESSAGE_LENGTH +from .services import SERVICES_KEY, ApplicationServices from .speech import SpeechTranscriber from .storage import AssistantStorage @@ -262,24 +263,30 @@ def build_inline_results(query: str) -> list[InlineQueryResultArticle]: ] +def get_services( + context: ContextTypes.DEFAULT_TYPE, +) -> ApplicationServices: + return context.application.bot_data[SERVICES_KEY] + + def get_storage(context: ContextTypes.DEFAULT_TYPE) -> AssistantStorage: - return context.application.bot_data["storage"] + return get_services(context).storage def get_ai_client(context: ContextTypes.DEFAULT_TYPE) -> AIClient: - return context.application.bot_data["ai_client"] + return get_services(context).ai_client def get_tz(context: ContextTypes.DEFAULT_TYPE) -> ZoneInfo: - return context.application.bot_data["timezone"] + return get_services(context).timezone def get_speech_recognizer(context: ContextTypes.DEFAULT_TYPE) -> SpeechTranscriber: - return context.application.bot_data["speech_recognizer"] + return get_services(context).speech_recognizer def get_voice_max_duration(context: ContextTypes.DEFAULT_TYPE) -> int: - return int(context.application.bot_data["voice_max_duration"]) + return get_services(context).voice_max_duration_seconds def command_text(context: ContextTypes.DEFAULT_TYPE) -> str: diff --git a/tests/test_application.py b/tests/test_application.py index 800f4a5..2b702f7 100644 --- a/tests/test_application.py +++ b/tests/test_application.py @@ -1,11 +1,11 @@ import tempfile import unittest from pathlib import Path -from unittest.mock import Mock, patch from zoneinfo import ZoneInfo from assistant_bot import application as application_module from assistant_bot.auth import authentication_guard +from assistant_bot.config import AppSettings from assistant_bot.handlers import ( ask_command, help_command, @@ -22,23 +22,26 @@ from assistant_bot.handlers import ( class ApplicationRegistrationTests(unittest.TestCase): def test_registers_the_current_public_handlers_in_order(self) -> None: - with tempfile.TemporaryDirectory() as directory, patch.multiple( - application_module, - get_assistant_password=Mock(return_value="secret"), - get_bot_token=Mock(return_value="123456:TEST"), - get_db_path=Mock( - return_value=Path(directory) / "assistant.sqlite3" - ), - get_assistant_mode=Mock(return_value="local"), - get_ollama_base_url=Mock(return_value="http://localhost:11434"), - get_whisper_model=Mock(return_value="small"), - get_whisper_device=Mock(return_value="cpu"), - get_whisper_compute_type=Mock(return_value="int8"), - get_whisper_language=Mock(return_value="ru"), - get_local_timezone=Mock(return_value=ZoneInfo("UTC")), - get_voice_max_duration_seconds=Mock(return_value=120), - ): - application = application_module.create_application() + with tempfile.TemporaryDirectory() as directory: + settings = AppSettings( + bot_token="123456:TEST", + assistant_password="secret", + db_path=Path(directory) / "assistant.sqlite3", + mode="local", + timezone=ZoneInfo("UTC"), + voice_max_duration_seconds=120, + ollama_base_url="http://localhost:11434", + ollama_model="qwen3.5:9b", + yandex_cloud_folder=None, + yandex_cloud_model="yandexgpt/latest", + yandex_stt_model="general", + yandex_stt_language="ru-RU", + whisper_model="small", + whisper_device="cpu", + whisper_compute_type="int8", + whisper_language="ru", + ) + application = application_module.create_application(settings) self.assertEqual( [handler.callback for handler in application.handlers[-1]], diff --git a/tests/test_auth.py b/tests/test_auth.py index 71ba342..12edbbd 100644 --- a/tests/test_auth.py +++ b/tests/test_auth.py @@ -4,11 +4,13 @@ from pathlib import Path from types import SimpleNamespace from typing import Any from unittest.mock import AsyncMock +from zoneinfo import ZoneInfo from telegram.constants import ChatType, ParseMode from telegram.ext import ApplicationHandlerStop from assistant_bot.auth import authentication_guard +from assistant_bot.services import SERVICES_KEY, ApplicationServices from assistant_bot.storage import AssistantStorage from assistant_bot.telegram_utils import render_markdown_chunks @@ -34,8 +36,14 @@ class AuthenticationGuardTests(unittest.IsolatedAsyncioTestCase): context = SimpleNamespace( application=SimpleNamespace( bot_data={ - "storage": self.storage, - "assistant_password": "секрет", + SERVICES_KEY: ApplicationServices( + storage=self.storage, + ai_client=SimpleNamespace(), + timezone=ZoneInfo("UTC"), + speech_recognizer=SimpleNamespace(), + voice_max_duration_seconds=120, + assistant_password="секрет", + ), } ) ) diff --git a/tests/test_config.py b/tests/test_config.py index 9005ccf..d7cd3eb 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -5,6 +5,35 @@ from unittest.mock import patch from assistant_bot import config +class AppSettingsTests(unittest.TestCase): + def test_loads_complete_startup_configuration_once(self) -> None: + environment = { + config.TOKEN_ENV_NAME: "token", + config.PASSWORD_ENV_NAME: "secret", + config.MODE_ENV_NAME: "local", + config.DB_ENV_NAME: "data/test.sqlite3", + config.TIMEZONE_ENV_NAME: "UTC", + config.WHISPER_LANGUAGE_ENV_NAME: "auto", + } + with patch( + "assistant_bot.config.load_env_file" + ) as load_env_file, patch.dict( + os.environ, + environment, + clear=True, + ): + settings = config.load_app_settings() + + load_env_file.assert_called_once_with() + self.assertEqual(settings.bot_token, "token") + self.assertEqual(settings.assistant_password, "secret") + self.assertEqual(settings.mode, "local") + self.assertEqual(settings.db_path, config.PROJECT_ROOT / "data/test.sqlite3") + self.assertEqual(settings.timezone.key, "UTC") + self.assertIsNone(settings.whisper_language) + self.assertIsNone(settings.yandex_cloud_folder) + + class WhisperConfigTests(unittest.TestCase): def test_gpu_int8_defaults(self) -> None: variable_names = ( diff --git a/tests/test_voice_handler.py b/tests/test_voice_handler.py index a0ee93b..8a29be6 100644 --- a/tests/test_voice_handler.py +++ b/tests/test_voice_handler.py @@ -3,10 +3,12 @@ from pathlib import Path from types import SimpleNamespace from typing import Any from unittest.mock import AsyncMock, MagicMock, patch +from zoneinfo import ZoneInfo from telegram.constants import ParseMode from assistant_bot.handlers import private_voice +from assistant_bot.services import SERVICES_KEY, ApplicationServices from assistant_bot.telegram_utils import render_markdown_chunks @@ -31,8 +33,14 @@ class VoiceHandlerTests(unittest.IsolatedAsyncioTestCase): bot=bot, application=SimpleNamespace( bot_data={ - "speech_recognizer": recognizer, - "voice_max_duration": 120, + SERVICES_KEY: ApplicationServices( + storage=SimpleNamespace(), + ai_client=SimpleNamespace(), + timezone=ZoneInfo("UTC"), + speech_recognizer=recognizer, + voice_max_duration_seconds=120, + assistant_password="secret", + ), } ), )