Begin DB

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

ums/utils/const.py

@@ -16,4 +16,5 @@
 import os
 
 BASE_PATH = '/ums-agent'
-SHARE_PATH = os.path.join(BASE_PATH, 'share')
+SHARE_PATH = os.path.join(BASE_PATH, 'share')
+PERSIST_PATH = os.path.join(BASE_PATH, 'persist')