108 lines
2.5 KiB
Python
108 lines
2.5 KiB
Python
# 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()
|
|
|
|
|