diff --git a/src/ocab_core/logger.py b/src/ocab_core/logger.py index 94989ae..b562c01 100644 --- a/src/ocab_core/logger.py +++ b/src/ocab_core/logger.py @@ -21,16 +21,7 @@ def setup_logger(): ) -async def log(message): - """ - Функция для логирования сообщений - - Она асинхронная, хотя таковой на самом деле не является. - """ - log_new(message) - - -def log_new(message): +def log(message): if isinstance(message, Exception): error_message = f"Error: {str(message)}\n{traceback.format_exc()}" logging.error(error_message) diff --git a/src/ocab_core/main.py b/src/ocab_core/main.py index 2960a8a..ce311fd 100644 --- a/src/ocab_core/main.py +++ b/src/ocab_core/main.py @@ -3,7 +3,7 @@ import traceback from aiogram import Bot, Dispatcher -from ocab_core.logger import setup_logger +from ocab_core.logger import log, setup_logger from ocab_core.modules_system import ModulesManager from ocab_core.modules_system.loaders import FSLoader from ocab_core.modules_system.loaders.unsafe_fs_loader import UnsafeFSLoader @@ -14,6 +14,7 @@ from service import paths bot_modules = [ UnsafeFSLoader(f"{paths.modules_standard}/config"), UnsafeFSLoader(f"{paths.modules_standard}/database"), + UnsafeFSLoader(f"{paths.modules_standard}/fsm_database_storage"), UnsafeFSLoader(f"{paths.modules_standard}/roles"), UnsafeFSLoader(f"{paths.modules_external}/yandexgpt"), FSLoader(f"{paths.modules_standard}/command_helper"), @@ -24,6 +25,11 @@ bot_modules = [ FSLoader(f"{paths.modules_standard}/message_processing"), ] +# import logging +# logger = logging.getLogger('peewee') +# logger.addHandler(logging.StreamHandler()) +# logger.setLevel(logging.DEBUG) + async def main(): bot = None @@ -33,13 +39,20 @@ async def main(): try: app.bot = Bot(token=get_telegram_token()) - - app.dp = Dispatcher() app.modules_manager = ModulesManager() for module_loader in bot_modules: + info = module_loader.info() + log(f"Loading {info.name}({info.id}) module") await app.modules_manager.load(module_loader) + app.dp = Dispatcher(storage=app.storage["_fsm_storage"]) + + app.dp.include_routers(*app.storage["_routers"]) + + for middleware in app.storage["_outer_message_middlewares"]: + app.dp.message.outer_middleware.register(middleware) + await app.modules_manager.late_init() await app.dp.start_polling(app.bot) except Exception: diff --git a/src/ocab_core/modules_system/public_api/public_api.py b/src/ocab_core/modules_system/public_api/public_api.py index 97b9d5f..6db0a33 100644 --- a/src/ocab_core/modules_system/public_api/public_api.py +++ b/src/ocab_core/modules_system/public_api/public_api.py @@ -5,20 +5,18 @@ from aiogram import BaseMiddleware, Router from aiogram.fsm.context import FSMContext from aiogram.fsm.storage.base import StorageKey -from ocab_core.logger import log_new +from ocab_core.logger import log from ocab_core.singleton import Singleton -log = log_new - def register_router(router: Router): app = Singleton() - app.dp.include_router(router) + app.storage["_routers"].append(router) def register_outer_message_middleware(middleware: BaseMiddleware): app = Singleton() - app.dp.message.outer_middleware.register(middleware) + app.storage["_outer_message_middlewares"].append(middleware) async def set_my_commands(commands): @@ -40,6 +38,12 @@ async def get_fsm_context(chat_id: int, user_id: int) -> FSMContext: ) +def set_fsm(storage): + app = Singleton() + log(storage) + app.storage["_fsm_storage"] = storage + + def get_module( module_id: str, paths=None ) -> Union[types.ModuleType, Union[Any, None], Tuple[Union[Any, None], ...]]: diff --git a/src/ocab_core/singleton.py b/src/ocab_core/singleton.py index cf2222f..228092c 100644 --- a/src/ocab_core/singleton.py +++ b/src/ocab_core/singleton.py @@ -1,4 +1,5 @@ from aiogram import Bot, Dispatcher +from aiogram.fsm.storage.memory import MemoryStorage from ocab_core.modules_system import ModulesManager @@ -17,4 +18,8 @@ class Singleton(metaclass=SingletonMeta): bot: Bot dp: Dispatcher = None modules_manager: ModulesManager = None - storage = dict() + storage = { + "_fsm_storage": MemoryStorage(), + "_routers": [], + "_outer_message_middlewares": [], + } diff --git a/src/ocab_modules/external/yandexgpt/handlers.py b/src/ocab_modules/external/yandexgpt/handlers.py index 47b4255..3505630 100644 --- a/src/ocab_modules/external/yandexgpt/handlers.py +++ b/src/ocab_modules/external/yandexgpt/handlers.py @@ -1,6 +1,4 @@ # flake8: noqa -import asyncio - from aiogram import Bot from aiogram.types import Message @@ -15,7 +13,7 @@ from ocab_modules.standard.database.db_api import add_message async def answer_to_message(message: Message, bot: Bot): # print("answer_to_message") - await log("answer_to_message") + log("answer_to_message") yagpt = YandexGPT(get_yandexgpt_token(), get_yandexgpt_catalog_id()) text = message.text prompt = get_yandexgpt_prompt() diff --git a/src/ocab_modules/external/yandexgpt/info.json b/src/ocab_modules/external/yandexgpt/info.json index 562921f..6deed05 100644 --- a/src/ocab_modules/external/yandexgpt/info.json +++ b/src/ocab_modules/external/yandexgpt/info.json @@ -1,5 +1,5 @@ { -w "id": "external.yandexgpt", + "id": "external.yandexgpt", "name": "Yandex GPT", "description": "Модуль для работы с Yandex GPT", "author": "OCAB Team", diff --git a/src/ocab_modules/external/yandexgpt/yandexgpt.py b/src/ocab_modules/external/yandexgpt/yandexgpt.py index 28bc195..b409d5b 100644 --- a/src/ocab_modules/external/yandexgpt/yandexgpt.py +++ b/src/ocab_modules/external/yandexgpt/yandexgpt.py @@ -57,7 +57,7 @@ class YandexGPT: ) except Exception as e: # TODO: Переделать обработку ошибок # print(e) - await log(f"Error: {e}") + log(f"Error: {e}") continue if int(len(response["tokens"])) < (max_tokens - answer_token): @@ -262,7 +262,7 @@ class YandexGPT: ) elif type == "yandexgpt": # print("yandexgpt_request") - await log("yandexgpt_request") + log("yandexgpt_request") messages = await self.collect_messages(message_id, chat_id) return await self.async_yandexgpt( system_prompt=get_yandexgpt_prompt(), diff --git a/src/ocab_modules/standard/command_helper/main.py b/src/ocab_modules/standard/command_helper/main.py index a9f1835..190d102 100644 --- a/src/ocab_modules/standard/command_helper/main.py +++ b/src/ocab_modules/standard/command_helper/main.py @@ -5,7 +5,6 @@ from aiogram.types import BotCommand, Message, TelegramObject from ocab_core.modules_system.public_api import ( get_module, - log, register_outer_message_middleware, set_my_commands, ) @@ -94,5 +93,4 @@ async def set_user_commands(): async def module_late_init(): - await log("module_late_init") await set_user_commands() diff --git a/src/ocab_modules/standard/create_report_apps/create_report.py b/src/ocab_modules/standard/create_report_apps/create_report.py index e151def..3e7d9c5 100644 --- a/src/ocab_modules/standard/create_report_apps/create_report.py +++ b/src/ocab_modules/standard/create_report_apps/create_report.py @@ -30,7 +30,7 @@ async def start_report(chat_id: int, bot: Bot): @router.message(ReportState.input_kernel_info) async def kernel_version_entered(message: Message, state: FSMContext): await state.update_data(kernel=message.text) - await message.answer(text="В каком приложении " "возникла проблема?") + await message.answer(text="В каком приложении возникла проблема?") await state.set_state(ReportState.input_app_name) @@ -38,7 +38,7 @@ async def kernel_version_entered(message: Message, state: FSMContext): async def app_name_entered(message: Message, state: FSMContext): await state.update_data(app_name=message.text) await message.answer( - text="Опиши проблему пошагово, " "что ты делал, что происходило, что не так" + text="Опиши проблему пошагово, что ты делал, что происходило, что не так" ) await state.set_state(ReportState.input_problem_step_by_step) @@ -46,7 +46,7 @@ async def app_name_entered(message: Message, state: FSMContext): @router.message(ReportState.input_problem_step_by_step) async def problem_step_by_step_entered(message: Message, state: FSMContext): await state.update_data(problem_step_by_step=message.text) - await message.answer(text="Вот твой отчет сообщением, " "а также файлом:") + await message.answer(text="Вот твой отчет сообщением, а также файлом:") data = await state.get_data() report = f"""Стенд с ошибкой: diff --git a/src/ocab_modules/standard/database/__init__.py b/src/ocab_modules/standard/database/__init__.py index c221e62..b282cf9 100644 --- a/src/ocab_modules/standard/database/__init__.py +++ b/src/ocab_modules/standard/database/__init__.py @@ -1,4 +1,4 @@ -from . import db_api, models +from . import db_api, models, repositories async def module_init(): diff --git a/src/ocab_modules/standard/database/db_api.py b/src/ocab_modules/standard/database/db_api.py index fc76455..1b32ace 100644 --- a/src/ocab_modules/standard/database/db_api.py +++ b/src/ocab_modules/standard/database/db_api.py @@ -7,6 +7,7 @@ from .exceptions import MissingModuleName, NotExpectedModuleName from .models.chat_stats import ChatStats from .models.chats import Chats from .models.db import database_proxy +from .models.fsm_data import FSMData from .models.messages import Messages from .models.user_stats import UserStats from .models.users import Users @@ -32,7 +33,7 @@ def connect_database(is_test: bool = False, module: str | None = None): def create_tables(db: pw.SqliteDatabase): """Создание таблиц""" - for table in Chats, Messages, Users, UserStats, ChatStats: + for table in Chats, Messages, Users, UserStats, ChatStats, FSMData: if not table.table_exists(): db.create_tables([table]) diff --git a/src/ocab_modules/standard/database/models/__init__.py b/src/ocab_modules/standard/database/models/__init__.py index e69de29..19118f4 100644 --- a/src/ocab_modules/standard/database/models/__init__.py +++ b/src/ocab_modules/standard/database/models/__init__.py @@ -0,0 +1 @@ +from .fsm_data import FSMData diff --git a/src/ocab_modules/standard/database/models/fsm_data.py b/src/ocab_modules/standard/database/models/fsm_data.py new file mode 100644 index 0000000..1cbff67 --- /dev/null +++ b/src/ocab_modules/standard/database/models/fsm_data.py @@ -0,0 +1,12 @@ +import peewee as pw + +from .db import database_proxy + + +class FSMData(pw.Model): + class Meta: + database = database_proxy + + key = pw.CharField(primary_key=True) + state = pw.CharField(null=True) + data = pw.CharField(null=True) diff --git a/src/ocab_modules/standard/database/repositories/__init__.py b/src/ocab_modules/standard/database/repositories/__init__.py new file mode 100644 index 0000000..5d225ba --- /dev/null +++ b/src/ocab_modules/standard/database/repositories/__init__.py @@ -0,0 +1 @@ +from .fsm_data import FSMDataRepository diff --git a/src/ocab_modules/standard/database/repositories/fsm_data.py b/src/ocab_modules/standard/database/repositories/fsm_data.py new file mode 100644 index 0000000..f8526bc --- /dev/null +++ b/src/ocab_modules/standard/database/repositories/fsm_data.py @@ -0,0 +1,32 @@ +from peewee import fn + +from ..models import FSMData + + +class FSMDataRepository: + def get(self, key: str): + return FSMData.get_or_none(key=key) + + def set_state(self, key: str, state: str): + FSMData.insert( + key=key, + state=state, + data=fn.COALESCE( + FSMData.select(FSMData.data).where(FSMData.key == key), None + ), + ).on_conflict( + conflict_target=[FSMData.key], + update={FSMData.state: state}, + ).execute() + + def set_data(self, key: str, data: str): + FSMData.insert( + key=key, + data=data, + state=fn.COALESCE( + FSMData.select(FSMData.state).where(FSMData.key == key), None + ), + ).on_conflict( + conflict_target=[FSMData.key], + update={FSMData.data: data}, + ).execute() diff --git a/src/ocab_modules/standard/fsm_database_storage/__init__.py b/src/ocab_modules/standard/fsm_database_storage/__init__.py new file mode 100644 index 0000000..87a9ce8 --- /dev/null +++ b/src/ocab_modules/standard/fsm_database_storage/__init__.py @@ -0,0 +1 @@ +from .fsm import module_init diff --git a/src/ocab_modules/standard/fsm_database_storage/fsm.py b/src/ocab_modules/standard/fsm_database_storage/fsm.py new file mode 100644 index 0000000..04a23e6 --- /dev/null +++ b/src/ocab_modules/standard/fsm_database_storage/fsm.py @@ -0,0 +1,122 @@ +import json +from typing import Any, Dict, Optional + +from aiogram.fsm.state import State +from aiogram.fsm.storage.base import BaseStorage, StorageKey + +from ocab_core.modules_system.public_api import get_module, log +from ocab_core.modules_system.public_api.public_api import set_fsm + +FSMDataRepository = get_module("standard.database", "repositories.FSMDataRepository") + + +def serialize_key(key: StorageKey) -> str: + return f"{key.bot_id}:{key.chat_id}:{key.user_id}" + + +def serialize_object(obj: object) -> str | None: + try: + return json.dumps(obj) + except Exception as e: + log(f"Serializing error! {e}") + return None + + +def deserialize_object(obj): + try: + return json.loads(obj) + except Exception as e: + log(f"Deserializing error! {e}") + return None + + +class SQLStorage(BaseStorage): + def __init__(self): + super().__init__() + self.repo = FSMDataRepository() + + async def set_state(self, key: StorageKey, state: State | None = None) -> None: + """ + Set state for specified key + + :param key: storage key + :param state: new state + """ + s_key = serialize_key(key) + s_state = state.state if isinstance(state, State) else state + + try: + self.repo.set_state(s_key, s_state) + except Exception as e: + log(f"FSM Storage database error: {e}") + + async def get_state(self, key: StorageKey) -> Optional[str]: + """ + Get key state + + :param key: storage key + :return: current state + """ + s_key = serialize_key(key) + + try: + s_state = self.repo.get(s_key) + return s_state.state if s_state else None + except Exception as e: + log(f"FSM Storage database error: {e}") + return None + + async def set_data(self, key: StorageKey, data: Dict[str, Any]) -> None: + """ + Write data (replace) + + :param key: storage key + :param data: new data + """ + s_key = serialize_key(key) + s_data = serialize_object(data) + + try: + self.repo.set_data(s_key, s_data) + except Exception as e: + log(f"FSM Storage database error: {e}") + + async def get_data(self, key: StorageKey) -> Optional[Dict[str, Any]]: + """ + Get current data for key + + :param key: storage key + :return: current data + """ + s_key = serialize_key(key) + + try: + s_data = self.repo.get(s_key) + return deserialize_object(s_data.data) if s_data else None + except Exception as e: + log(f"FSM Storage database error: {e}") + return None + + async def update_data( + self, key: StorageKey, data: Dict[str, Any] + ) -> Dict[str, Any]: + """ + Update data in the storage for key (like dict.update) + + :param key: storage key + :param data: partial data + :return: new data + """ + current_data = await self.get_data(key=key) + if not current_data: + current_data = {} + current_data.update(data) + await self.set_data(key=key, data=current_data) + return current_data.copy() + + async def close(self) -> None: # pragma: no cover + pass + + +async def module_init(): + set_fsm(SQLStorage()) diff --git a/src/ocab_modules/standard/fsm_database_storage/info.json b/src/ocab_modules/standard/fsm_database_storage/info.json new file mode 100644 index 0000000..a592463 --- /dev/null +++ b/src/ocab_modules/standard/fsm_database_storage/info.json @@ -0,0 +1,11 @@ +{ + "id": "standard.fsm_database_storage", + "name": "FSM Database Storage", + "description": "Очень полезный модуль", + "author": "OCAB Team", + "version": "1.0.0", + "privileged": false, + "dependencies": { + "standard.database": "^1.0.0" + } +} diff --git a/src/ocab_modules/standard/info/handlers.py b/src/ocab_modules/standard/info/handlers.py index 40388fe..22a3027 100644 --- a/src/ocab_modules/standard/info/handlers.py +++ b/src/ocab_modules/standard/info/handlers.py @@ -32,7 +32,7 @@ async def get_info_answer_by_id(message: Message, bot: Bot, user_id: int): if user is None: await message.reply("Пользователь не найден") - await log(f"Пользователь не найден: {user_id}, {user}") + log(f"Пользователь не найден: {user_id}, {user}") return roles = Roles() @@ -69,7 +69,7 @@ async def get_user_info(message: Message, bot: Bot): " попробуйте запросить информацию о пользователе по его тегу или ответив на его сообщение" ) # print(e) - await log(e) + log(e) async def get_chat_info(message: Message, bot: Bot):