diff --git a/ums/management/db.py b/ums/management/db.py new file mode 100644 index 0000000..19a7bca --- /dev/null +++ b/ums/management/db.py @@ -0,0 +1,108 @@ +# Agenten Plattform +# +# (c) 2024 Magnus Bender +# Institute of Humanities-Centered Artificial Intelligence (CHAI) +# Universitaet Hamburg +# https://www.chai.uni-hamburg.de/~bender +# +# source code released under the terms of GNU Public License Version 3 +# https://www.gnu.org/licenses/gpl-3.0.txt + +import os +from datetime import datetime + +import sqlite3 +import atexit + +from threading import Lock + +from ums.utils import PERSIST_PATH, AgentMessage + +class DB(): + + def __init__(self): + self.db = sqlite3.connect( + os.path.join(PERSIST_PATH, 'messages.db'), + check_same_thread=False + ) + self.db.row_factory = sqlite3.Row + + self.dblock = Lock() + atexit.register(lambda db : db.close(), self.db) + + self._assure_tables() + + def _assure_tables(self): + self.dblock.acquire() + with self.db: + self.db.execute("""CREATE TABLE IF NOT EXISTS Messages ( + count INTEGER PRIMARY KEY AUTOINCREMENT, + id TEXT, + fr TEXT, + to TEXT, + time TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + json BLOB, + processed BOOL DEFAULT FALSE + )""") + self.dblock.release() + + def add_message(self, fr:str, to:str, message:AgentMessage, processed:bool=False) -> int: + self.dblock.acquire() + with self.db: + self.db.execute( + """INSERT INTO Messages ( + id, fr, to, json, processed + ) VALUES ( + :id, :fr, :to, :json, :processed + )""", { + "id" : message.id, + "fr" : fr, + "to" : to, + "json" : message.model_dump_json(), + "processed" : processed + }) + new_count = self.db.execute("LAST_INSERT_ROWID()").fetchone() + self.dblock.release() + + return new_count + + def set_processed(self, count:int, processed:bool=True) -> bool: + self.dblock.acquire() + with self.db: + self.db.execute("UPDATE Messages SET processed = ? WHERE count = ?", (processed, count)) + self.dblock.release() + + def iterate(self, + id:str|None=None, fr:str|None=None, to:str|None=None, + processed:bool|None=None, + time_after:int|None=0, time_before:int|None=None, + limit:int=20, offset:int=0 + ): + + where = [] + + if time_after: + time_after = datetime.fromtimestamp(time_after).strftime("%Y-%m-%d %H:%M:%S") + if time_before: + time_before = datetime.fromtimestamp(time_before).strftime("%Y-%m-%d %H:%M:%S") + + with self.db: + self.db.execute( + "SELECT * FROM Messages WHERE {} LIMIT :lim OFFSET :off".format(', '.join(where)), + { + + }) + + + + + def by_count(self, count:int): + with self.db: + self.db.execute("SELECT * FROM Messages WHERE count = ?").fetchone(count) + # error + # create object + + def __iter__(self): + yield from self.iterate() + + \ No newline at end of file diff --git a/ums/utils/const.py b/ums/utils/const.py index 185dcc2..8b5391c 100644 --- a/ums/utils/const.py +++ b/ums/utils/const.py @@ -16,4 +16,5 @@ import os BASE_PATH = '/ums-agent' -SHARE_PATH = os.path.join(BASE_PATH, 'share') \ No newline at end of file +SHARE_PATH = os.path.join(BASE_PATH, 'share') +PERSIST_PATH = os.path.join(BASE_PATH, 'persist') \ No newline at end of file