wip

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)

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:

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], ...]]:

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": [],
+    }

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

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

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(),

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

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"""Стенд с ошибкой:

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():

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

src/ocab_modules/standard/database/models/__init__.py

@@ -0,0 +1 @@
+from .fsm_data import FSMData

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)

src/ocab_modules/standard/database/repositories/__init__.py

@@ -0,0 +1 @@
+from .fsm_data import FSMDataRepository

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

src/ocab_modules/standard/fsm_database_storage/__init__.py

@@ -0,0 +1 @@
+from .fsm import module_init

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

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"
+    }
+}

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