Add plain text db

bot/config.py

@@ -18,7 +18,9 @@
 allowed_telegram_usernames = config_yaml["allowed_telegram_usernames"]
 new_dialog_timeout = config_yaml["new_dialog_timeout"]
 enable_message_streaming = config_yaml.get("enable_message_streaming", True)
-mongodb_uri = f"mongodb://mongo:{config_env['MONGODB_PORT']}"
+db_type = config_yaml.get("db_type", "mongodb")
+if db_type == "mongodb":
+    mongodb_uri = f"mongodb://mongo:{config_env['MONGODB_PORT']}"
 
 # chat_modes
 with open(config_dir / "chat_modes.yml", 'r') as f:

bot/database.py

@@ -1,13 +1,47 @@
+import abc
+import os
+from collections import defaultdict
 from typing import Optional, Any
 
 import pymongo
 import uuid
 from datetime import datetime
 
+import yaml
+
 import config
 
+class AbstractDB:
+    @abc.abstractmethod
+    def check_if_user_exists(self, user_id: int, raise_exception: bool = False):
+        pass
 
-class Database:
+    @abc.abstractmethod
+    def add_new_user(self, user: dict):
+        pass
+
+    @abc.abstractmethod
+    def add_dialog(self, dialog: dict):
+        pass
+
+    @abc.abstractmethod
+    def get_user(self, user_id: int):
+        pass
+
+    @abc.abstractmethod
+    def set_user_attribute(self, user_id: int, key: str, value: Any):
+        pass
+
+    @abc.abstractmethod
+    def get_dialog(self, dialog_id: int, user_id: int):
+        pass
+
+    @abc.abstractmethod
+    def set_dialog_messages(self, dialog_id: int, user_id: int, dialog_messages: list):
+        pass
+
+
+class MongoDB(AbstractDB):
     def __init__(self):
         self.client = pymongo.MongoClient(config.mongodb_uri)
         self.db = self.client["chatgpt_telegram_bot"]
@@ -24,13 +58,139 @@ def check_if_user_exists(self, user_id: int, raise_exception: bool = False):
             else:
                 return False
 
+    def add_new_user(self, user: dict):
+        self.user_collection.insert_one(user)
+
+    def add_dialog(self, dialog: dict):
+        self.dialog_collection.insert_one(dialog)
+
+        # update user's current dialog
+        self.user_collection.update_one(
+            {"_id": dialog["user_id"]},
+            {"$set": {"current_dialog_id": dialog["_id"]}}
+        )
+
+    def get_user(self, user_id: int):
+        return self.user_collection.find_one({"_id": user_id})
+
+    def set_user_attribute(self, user_id: int, key: str, value: Any):
+        self.user_collection.update_one({"_id": user_id}, {"$set": {key: value}})
+
+    def get_dialog(self, dialog_id: int, user_id: int):
+        return self.dialog_collection.find_one({"_id": dialog_id, "user_id": user_id})
+
+    def set_dialog_messages(self, dialog_id: int, user_id: int, dialog_messages: list):
+        self.dialog_collection.update_one(
+            {"_id": dialog_id, "user_id": user_id},
+            {"$set": {"messages": dialog_messages}}
+        )
+
+
+class PlainFileUser:
+
+    def __init__(self):
+        self.user = None
+        self.dialogs = {}
+
+    def _folder_name(self):
+        return self.user.get("username", str(self.user["_id"]))
+
+    def commit(self):
+        if not os.path.exists(os.path.join(PlainFileDB.DB_PATH, self._folder_name())):
+            os.makedirs(os.path.join(PlainFileDB.DB_PATH, self._folder_name()))
+        with open(os.path.join(PlainFileDB.DB_PATH, self._folder_name(), "user.yaml"), "w") as f:
+            yaml.safe_dump(self.user, f)
+        for dialog_id, dialog in self.dialogs.items():
+            with open(os.path.join(PlainFileDB.DB_PATH, self._folder_name(), f"dialog_{dialog_id}.yaml"), "w") as f:
+                yaml.safe_dump(dialog, f)
+
+
+
+class PlainFileDB(AbstractDB):
+    DB_PATH = "db"
+
+    def load_user(self, name):
+        path = os.path.join(self.DB_PATH, name)
+        for file in os.listdir(path):
+            if not file.endswith(".yaml"):
+                raise ValueError(f"Unknown file type {path}")
+            if file.startswith("user"):
+                with open(os.path.join(path, file), "r") as f:
+                    user = yaml.safe_load(f)
+                    self.user_collection[user["_id"]].user = user
+            elif file.startswith("dialog"):
+                with open(os.path.join(path, file), "r") as f:
+                    dialog = yaml.safe_load(f)
+                    self.user_collection[dialog["user_id"]].dialogs[dialog["_id"]] = dialog
+
+    def __init__(self):
+        self.user_collection = defaultdict(PlainFileUser)
+
+        if not os.path.exists(self.DB_PATH):
+            os.makedirs(self.DB_PATH)
+        for file in os.listdir(self.DB_PATH):
+            if not os.path.isdir(os.path.join(self.DB_PATH, file)):
+                raise ValueError(f"Unknown file type {file}")
+            self.load_user(file)
+
+    def check_if_user_exists(self, user_id: int, raise_exception: bool = False):
+        if self.user_collection[user_id].user is not None:
+            return True
+        else:
+            if raise_exception:
+                raise ValueError(f"User {user_id} does not exist")
+            else:
+                return False
+
+    def add_new_user(self, user: dict):
+        db_user = self.user_collection[user["_id"]]
+        db_user.user = user
+        db_user.commit()
+
+
+    def add_dialog(self, dialog: dict):
+        self.user_collection[dialog["user_id"]].dialogs[dialog["_id"]] = dialog
+
+        # update user's current dialog
+        self.user_collection[dialog["user_id"]].user["current_dialog_id"] = dialog["_id"]
+
+        self.user_collection[dialog["user_id"]].commit()
+
+    def get_user(self, user_id: int):
+        return self.user_collection[user_id].user
+
+    def set_user_attribute(self, user_id: int, key: str, value: Any):
+        self.user_collection[user_id].user[key] = value
+        self.user_collection[user_id].commit()
+
+    def get_dialog(self, dialog_id: int, user_id: int):
+        return self.user_collection[user_id].dialogs[dialog_id]
+
+    def set_dialog_messages(self, dialog_id: int, user_id: int, dialog_messages: list):
+        self.user_collection[user_id].dialogs[dialog_id]["messages"] = dialog_messages
+        self.user_collection[user_id].commit()
+
+
+class Database:
+    db: AbstractDB
+
+    def __init__(self):
+        if config.db_type == "mongodb":
+            self.db = MongoDB()
+        elif config.db_type == "plain_file":
+            self.db = PlainFileDB()
+        else:
+            raise ValueError(f"Unknown database type {config.db_type}")
+    def check_if_user_exists(self, user_id: int, raise_exception: bool = False):
+        return self.db.check_if_user_exists(user_id, raise_exception)
+
     def add_new_user(
-        self,
-        user_id: int,
-        chat_id: int,
-        username: str = "",
-        first_name: str = "",
-        last_name: str = "",
+            self,
+            user_id: int,
+            chat_id: int,
+            username: str = "",
+            first_name: str = "",
+            last_name: str = "",
     ):
         user_dict = {
             "_id": user_id,
@@ -52,11 +212,11 @@ def add_new_user(
             "n_transcribed_seconds": 0.0  # voice message transcription
         }
 
-        if not self.check_if_user_exists(user_id):
-            self.user_collection.insert_one(user_dict)
+        if not self.db.check_if_user_exists(user_id):
+            self.db.add_new_user(user_dict)
 
     def start_new_dialog(self, user_id: int):
-        self.check_if_user_exists(user_id, raise_exception=True)
+        self.db.check_if_user_exists(user_id, raise_exception=True)
 
         dialog_id = str(uuid.uuid4())
         dialog_dict = {
@@ -69,28 +229,22 @@ def start_new_dialog(self, user_id: int):
         }
 
         # add new dialog
-        self.dialog_collection.insert_one(dialog_dict)
-
-        # update user's current dialog
-        self.user_collection.update_one(
-            {"_id": user_id},
-            {"$set": {"current_dialog_id": dialog_id}}
-        )
+        self.db.add_dialog(dialog_dict)
 
         return dialog_id
 
     def get_user_attribute(self, user_id: int, key: str):
-        self.check_if_user_exists(user_id, raise_exception=True)
-        user_dict = self.user_collection.find_one({"_id": user_id})
+        self.db.check_if_user_exists(user_id, raise_exception=True)
+        user_dict = self.db.get_user(user_id)
 
         if key not in user_dict:
             return None
 
         return user_dict[key]
 
     def set_user_attribute(self, user_id: int, key: str, value: Any):
-        self.check_if_user_exists(user_id, raise_exception=True)
-        self.user_collection.update_one({"_id": user_id}, {"$set": {key: value}})
+        self.db.check_if_user_exists(user_id, raise_exception=True)
+        self.db.set_user_attribute(user_id, key, value)
 
     def update_n_used_tokens(self, user_id: int, model: str, n_input_tokens: int, n_output_tokens: int):
         n_used_tokens_dict = self.get_user_attribute(user_id, "n_used_tokens")
@@ -107,21 +261,18 @@ def update_n_used_tokens(self, user_id: int, model: str, n_input_tokens: int, n_
         self.set_user_attribute(user_id, "n_used_tokens", n_used_tokens_dict)
 
     def get_dialog_messages(self, user_id: int, dialog_id: Optional[str] = None):
-        self.check_if_user_exists(user_id, raise_exception=True)
+        self.db.check_if_user_exists(user_id, raise_exception=True)
 
         if dialog_id is None:
             dialog_id = self.get_user_attribute(user_id, "current_dialog_id")
 
-        dialog_dict = self.dialog_collection.find_one({"_id": dialog_id, "user_id": user_id})
+        dialog_dict = self.db.get_dialog(dialog_id, user_id)
         return dialog_dict["messages"]
 
     def set_dialog_messages(self, user_id: int, dialog_messages: list, dialog_id: Optional[str] = None):
-        self.check_if_user_exists(user_id, raise_exception=True)
+        self.db.check_if_user_exists(user_id, raise_exception=True)
 
         if dialog_id is None:
             dialog_id = self.get_user_attribute(user_id, "current_dialog_id")
 
-        self.dialog_collection.update_one(
-            {"_id": dialog_id, "user_id": user_id},
-            {"$set": {"messages": dialog_messages}}
-        )
+        self.db.set_dialog_messages(dialog_id, user_id, dialog_messages)