2.1

6 files changed, 59 insertions, 32 deletions

diff --git a/shared/__init__.py b/shared/__init__.py
index f708bfe..edfd5c2 100644
--- a/shared/__init__.py
+++ b/shared/__init__.py
@@ -2,4 +2,5 @@
 from . import settings
 from . import instances
 from . import database
+from . import samples
 from . import commands
diff --git a/shared/commands.py b/shared/commands.py
index 684997c..a7ff002 100644
--- a/shared/commands.py
+++ b/shared/commands.py
@@ -6,19 +6,16 @@ from aiogram.types import BotCommandScopeAllPrivateChats as private
 commands = {
     group(): [
         cmd("gen", "Сгенерировать сообщение"),
-        cmd("del", "Удалить сообщение бота"),
         cmd("pin", "Создать опрос для закрепления сообщения"),
     ],
     admin(): [
         cmd("gen", "Сгенерировать сообщение"),
-        cmd("del", "Удалить сообщение бота"),
         cmd("pin", "Создать опрос для закрепления сообщения"),
         cmd("void", "Отчистить связи для генерации сообщений"),
         cmd("config", "Открыть настройки чата"),
     ],
     private(): [
         cmd("gen", "Сгенерировать сообщение"),
-        cmd("del", "Удалить сообщение бота"),
         cmd("void", "Отчистить связи для генерации сообщений"),
     ],
 }
diff --git a/shared/database.py b/shared/database.py
index 9dfe868..6ce5fa2 100644
--- a/shared/database.py
+++ b/shared/database.py
@@ -11,6 +11,7 @@ class Message(Base):
     __tablename__ = "messages"
     chat_id: Mapped[int] = mapped_column(primary_key=True)
     message_id: Mapped[int] = mapped_column(primary_key=True)
+    user_id: Mapped[int]
     message: Mapped[str]
 
 
diff --git a/shared/instances.py b/shared/instances.py
index fb511b3..5528206 100644
--- a/shared/instances.py
+++ b/shared/instances.py
@@ -7,10 +7,10 @@ from sqlalchemy.orm import Session, sessionmaker
 from shared.settings import Chats, Settings
 
 settings = Settings()
-config = Chats()
+chats = Chats("data/config.json")
 if not exists("data/config.json"):
-    config.save("data/config.json")
-config.load("data/config.json")
+    chats.save()
+chats.load()
 
 bot = Bot(token=settings.token)
 dp = Dispatcher(bot)
diff --git a/shared/samples.py b/shared/samples.py
new file mode 100644
index 0000000..740cff5
--- /dev/null
+++ b/shared/samples.py
@@ -0,0 +1,31 @@
+from tgen import TextGenerator
+
+from shared.database import Message
+from shared.instances import session
+
+
+class Samples:
+    samples: dict[int, TextGenerator]
+
+    def __init__(self) -> None:
+        self.samples = {}
+
+    def get(self, chat_id: int) -> TextGenerator:
+        if chat_id not in self.samples:
+            with session() as s:
+                samples = [
+                    m.tuple()[0]
+                    for m in s.query(Message.message)
+                    .filter(Message.chat_id == chat_id)
+                    .all()
+                ]
+
+            self.samples[chat_id] = TextGenerator.from_samples(samples)
+        return self.samples[chat_id]
+
+    def delete(self, chat_id: int) -> None:
+        if chat_id in self.samples:
+            self.samples.pop(chat_id)
+
+
+samples = Samples()
diff --git a/shared/settings.py b/shared/settings.py
index 560f8f8..bff7ff9 100644
--- a/shared/settings.py
+++ b/shared/settings.py
@@ -11,7 +11,7 @@ class GenConfig(BaseModel):
     chance: int = Field(
         10,
         description="Шанс с которым бот сгенерирует сообщение",
-        ge=1,
+        ge=0,
         le=100,
     )
     reply: bool = Field(
@@ -22,16 +22,6 @@ class GenConfig(BaseModel):
         True,
         description="Включить/Выключить удаление /gen команды",
     )
-    min_word_count: int | str | None = Field(
-        None,
-        description="Минимальное количество слов в сгенерированном предложении",
-        ge=1,
-    )
-    max_word_count: int | None = Field(
-        None,
-        description="Максимальное количество слов в сгенерированном предложении",
-        ge=1,
-    )
 
 
 class PollConfig(BaseModel):
@@ -60,22 +50,29 @@ class Config(BaseModel):
     )
 
 
-class Chats(BaseModel):
-    chats: dict[int, Config] = {}
+class Chats:
+    file_name: str
+    configs: dict[int, Config]
+
+    def __init__(self, file_name: str) -> None:
+        self.file_name = file_name
+        self.configs = {}
 
-    @classmethod
-    def load(cls, file_name: str) -> "Chats":
-        with open(file_name, "r") as file:
-            return cls.parse_obj(load(file))
+    def load(self) -> None:
+        with open(self.file_name, "r") as file:
+            self.configs = {
+                id_: Config.parse_obj(config) for id_, config in load(file).items()
+            }
 
-    def save(self, file_name: str) -> None:
-        with open(file_name, "w") as file:
-            dump(self.schema(), file)
+    def save(self) -> None:
+        with open(self.file_name, "w") as file:
+            dump({id_: config.dict() for id_, config in self.configs.items()}, file)
 
-    def get_config(self, chat_id: int) -> Config:
-        if chat_id not in self.chats:
-            self.chats[chat_id] = Config()
-        return self.chats[chat_id]
+    def get(self, chat_id: int) -> Config:
+        if chat_id not in self.configs:
+            self.configs[chat_id] = Config()
+        return self.configs[chat_id]
 
-    def set_config(self, chat_id: int, config: Config) -> None:
-        self.chats[chat_id] = config
+    def set(self, chat_id: int, config: Config) -> None:
+        self.configs[chat_id] = config
+        self.save()