From c53f6025aea0d63fc95a515556a78418ee41a977 Mon Sep 17 00:00:00 2001
From: Maxim Slipenko <maxim@slipenko.com>
Date: Sat, 13 Jul 2024 22:33:10 +0300
Subject: [PATCH] wip
---
src/ocab_core/logger.py | 11 +-
src/ocab_core/main.py | 19 ++-
.../modules_system/public_api/public_api.py | 14 +-
src/ocab_core/singleton.py | 7 +-
.../external/yandexgpt/handlers.py | 4 +-
src/ocab_modules/external/yandexgpt/info.json | 2 +-
.../external/yandexgpt/yandexgpt.py | 4 +-
.../standard/command_helper/main.py | 2 -
.../create_report_apps/create_report.py | 6 +-
.../standard/database/__init__.py | 2 +-
src/ocab_modules/standard/database/db_api.py | 3 +-
.../standard/database/models/__init__.py | 1 +
.../standard/database/models/fsm_data.py | 12 ++
.../database/repositories/__init__.py | 1 +
.../database/repositories/fsm_data.py | 32 +++++
.../standard/fsm_database_storage/__init__.py | 1 +
.../standard/fsm_database_storage/fsm.py | 122 ++++++++++++++++++
.../standard/fsm_database_storage/info.json | 11 ++
src/ocab_modules/standard/info/handlers.py | 4 +-
19 files changed, 224 insertions(+), 34 deletions(-)
create mode 100644 src/ocab_modules/standard/database/models/fsm_data.py
create mode 100644 src/ocab_modules/standard/database/repositories/__init__.py
create mode 100644 src/ocab_modules/standard/database/repositories/fsm_data.py
create mode 100644 src/ocab_modules/standard/fsm_database_storage/__init__.py
create mode 100644 src/ocab_modules/standard/fsm_database_storage/fsm.py
create mode 100644 src/ocab_modules/standard/fsm_database_storage/info.json
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):