2.1

1 files changed, 31 insertions, 0 deletions

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