ums.management.db

  1# Agenten Plattform
  2#
  3# (c) 2024 Magnus Bender
  4# 	Institute of Humanities-Centered Artificial Intelligence (CHAI)
  5# 	Universitaet Hamburg
  6# 	https://www.chai.uni-hamburg.de/~bender
  7#  
  8# source code released under the terms of GNU Public License Version 3
  9# https://www.gnu.org/licenses/gpl-3.0.txt
 10
 11import os
 12import sqlite3, atexit
 13
 14from datetime import datetime
 15from threading import Lock
 16from typing import Generator
 17
 18from pydantic import validate_call, BaseModel
 19
 20from ums.utils import PERSIST_PATH, AgentMessage
 21
 22class RowObject(BaseModel):
 23	"""
 24		Object representing a database row.
 25	"""
 26
 27	count : int
 28	"""
 29		The count (primary key) of the item.
 30	"""
 31
 32	sender : str
 33	"""
 34		The sender of the message.
 35	""" 
 36
 37	recipient : str
 38	"""
 39		The recipient of the message
 40	"""
 41
 42	time : int
 43	"""
 44		The time (unix timestamp) the message was received/ sent.
 45	"""
 46
 47	message : AgentMessage
 48	"""
 49		The message  received/ sent.
 50	"""
 51
 52	processed : bool
 53	"""
 54		Did the management process the message, i.e., did the tasks necessary for this message (mostly only relevant for received messages).
 55	"""
 56
 57class DB():
 58
 59	_DB_TIME_FORMAT = "%Y-%m-%d %H:%M:%S"
 60	
 61	def __init__(self):
 62		self.db = sqlite3.connect(
 63			os.path.join(PERSIST_PATH, 'messages.db'),
 64			check_same_thread=False
 65		) 
 66		self.db.row_factory = sqlite3.Row
 67
 68		self.dblock = Lock()
 69		atexit.register(lambda db : db.close(), self.db)
 70
 71		self._assure_tables()
 72
 73	def _assure_tables(self):
 74		self.dblock.acquire()
 75		with self.db:
 76			self.db.execute("""CREATE TABLE IF NOT EXISTS Messages ( 
 77				count INTEGER PRIMARY KEY AUTOINCREMENT,
 78				id TEXT, 
 79				sender TEXT,
 80			 	recipient TEXT,
 81				time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
 82				json BLOB,
 83				processed BOOL DEFAULT FALSE
 84			)""")
 85		self.dblock.release()
 86
 87	@validate_call
 88	def add_message(self, sender:str, recipient:str, message:AgentMessage, processed:bool=False) -> int:
 89		self.dblock.acquire()
 90		with self.db:
 91			self.db.execute(
 92				"""INSERT INTO Messages (
 93						id, sender, recipient, json, processed 
 94					) VALUES (
 95						:id, :sender, :recipient, :json, :processed 
 96				)""", {
 97					"id" : message.id,
 98					"sender" : sender,
 99					"recipient" : recipient,
100					"json" : message.model_dump_json(),
101					"processed" : processed
102				})
103			new_count = self.db.execute("SELECT LAST_INSERT_ROWID() as last").fetchone()
104		self.dblock.release()
105
106		return new_count['last']
107
108	@validate_call
109	def set_processed(self, count:int, processed:bool=True) -> bool:
110		self.dblock.acquire()
111		with self.db:
112			try:
113				self.db.execute("UPDATE Messages SET processed = ? WHERE count = ?", (processed, count))
114				return True
115			except:
116				return False
117		self.dblock.release()
118
119	def __iter__(self) -> Generator[RowObject, None, None]:
120		yield from self.iterate()
121
122	@validate_call
123	def iterate(self,
124		id:str|None=None, sender:str|None=None, recipient:str|None=None,
125		processed:bool|None=None,
126		time_after:int|None=None, time_before:int|None=None,
127		limit:int=20, offset:int=0
128	) -> Generator[RowObject, None, None]:
129		
130		where = []
131		params = {
132			"lim": limit,
133			"off": offset
134		}
135
136		for v,n in ((id,'id'), (sender,'sender'), (recipient,'recipient'), (processed,'processed')):
137			if not v is None:
138				where.append('{} = :{}'.format(n,n))
139				params[n] = v
140
141		if time_after:
142			where.append("time > :t_after")
143			params['t_after'] = datetime.fromtimestamp(time_after).strftime(self._DB_TIME_FORMAT)
144
145		if time_before:
146			where.append("time < :t_before")
147			params['t_before'] = datetime.fromtimestamp(time_before).strftime(self._DB_TIME_FORMAT)
148
149		if len(where) > 0:
150			where_clause = "WHERE " + (' AND '.join(where))
151		else:
152			where_clause = ""
153
154		with self.db:
155			for row in self.db.execute(
156				"SELECT * FROM Messages {} LIMIT :lim OFFSET :off".format(where_clause),
157				params
158			):
159				yield self._create_row_object(row)
160
161	def _create_row_object(self, row:sqlite3.Row) -> RowObject:
162		return RowObject(
163			count=row['count'],
164			sender=row['sender'],
165			recipient=row['recipient'],
166			time=int(datetime.strptime(row['time'], self._DB_TIME_FORMAT).timestamp()),
167			message=AgentMessage.model_construct(row['json']),
168			processed=row['processed']
169		)
170
171	def by_count(self, count:int) -> RowObject|None:
172		with self.db:
173			try:
174				return self._create_row_object(
175					self.db.execute("SELECT * FROM Messages WHERE count = ?", (count,)).fetchone()
176				)
177			except:
178				return None
179