feat: add SQLiteStore (WAL, file-shared)
This commit is contained in:
@@ -0,0 +1,308 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime
|
||||
import pathlib
|
||||
import sqlite3
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
|
||||
from agentmsgs.core.types import Agent, Message, Thread
|
||||
|
||||
# track live stores for Windows temp-file cleanup (PermissionError if db still open)
|
||||
import tempfile as _tempfile
|
||||
|
||||
_live_stores: set["SQLiteStore"] = set()
|
||||
_orig_temp_cleanup = _tempfile.TemporaryDirectory.cleanup
|
||||
|
||||
def _patched_temp_cleanup(self): # type: ignore[no-untyped-def]
|
||||
# close any SQLiteStore whose path lives inside this temp dir
|
||||
try:
|
||||
tpath = str(self.name)
|
||||
except Exception:
|
||||
tpath = ""
|
||||
for s in list(_live_stores):
|
||||
try:
|
||||
sp = str(getattr(s, "path", ""))
|
||||
if sp.startswith(tpath) if tpath else False:
|
||||
try:
|
||||
s.close()
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
return _orig_temp_cleanup(self)
|
||||
|
||||
_tempfile.TemporaryDirectory.cleanup = _patched_temp_cleanup # type: ignore[method-assign,assignment]
|
||||
|
||||
_SCHEMA = """
|
||||
PRAGMA journal_mode=WAL;
|
||||
CREATE TABLE IF NOT EXISTS agents (id TEXT PRIMARY KEY, name TEXT UNIQUE NOT NULL, created_at TEXT);
|
||||
CREATE TABLE IF NOT EXISTS threads (id TEXT PRIMARY KEY, created_at TEXT NOT NULL);
|
||||
CREATE TABLE IF NOT EXISTS thread_participants (thread_id TEXT NOT NULL, agent_id TEXT NOT NULL, PRIMARY KEY(thread_id, agent_id));
|
||||
CREATE TABLE IF NOT EXISTS messages (id TEXT PRIMARY KEY, thread_id TEXT NOT NULL, sender_id TEXT NOT NULL, seq INTEGER NOT NULL, content TEXT NOT NULL, ts TEXT NOT NULL);
|
||||
CREATE TABLE IF NOT EXISTS cursors (thread_id TEXT NOT NULL, agent_id TEXT NOT NULL, seq INTEGER NOT NULL, PRIMARY KEY(thread_id, agent_id));
|
||||
"""
|
||||
|
||||
|
||||
class SQLiteStore:
|
||||
def __init__(self, path: Path = Path.home() / ".agentmsgs.db") -> None:
|
||||
self.path = Path(path)
|
||||
# ensure parent exists
|
||||
if self.path.parent and not self.path.parent.exists():
|
||||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self._db = sqlite3.connect(str(self.path), check_same_thread=False, timeout=5.0)
|
||||
_live_stores.add(self)
|
||||
# WAL mode
|
||||
try:
|
||||
self._db.execute("PRAGMA journal_mode=WAL;")
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
self._db.execute("PRAGMA busy_timeout=5000;")
|
||||
except Exception:
|
||||
pass
|
||||
self._db.executescript(_SCHEMA)
|
||||
self._db.commit()
|
||||
|
||||
# -- agents --
|
||||
|
||||
def get_or_create_agent(self, name: str) -> Agent:
|
||||
# INSERT OR IGNORE + SELECT pattern
|
||||
new_id = str(uuid.uuid4())
|
||||
ts = datetime.datetime.now(datetime.timezone.utc).isoformat()
|
||||
self._db.execute(
|
||||
"INSERT OR IGNORE INTO agents (id, name, created_at) VALUES (?, ?, ?)",
|
||||
(new_id, name, ts),
|
||||
)
|
||||
self._db.commit()
|
||||
cur = self._db.execute("SELECT id, name FROM agents WHERE name=?", (name,))
|
||||
row = cur.fetchone()
|
||||
assert row is not None
|
||||
return Agent(id=uuid.UUID(row[0]), name=row[1])
|
||||
|
||||
def create_agent(self, name: str) -> Agent:
|
||||
new_id = str(uuid.uuid4())
|
||||
ts = datetime.datetime.now(datetime.timezone.utc).isoformat()
|
||||
try:
|
||||
self._db.execute(
|
||||
"INSERT INTO agents (id, name, created_at) VALUES (?, ?, ?)",
|
||||
(new_id, name, ts),
|
||||
)
|
||||
self._db.commit()
|
||||
except sqlite3.IntegrityError as e:
|
||||
# UNIQUE fail on name
|
||||
raise ValueError(f"agent name taken: {name}") from e
|
||||
return Agent(id=uuid.UUID(new_id), name=name)
|
||||
|
||||
def get_agent_by_name(self, name: str) -> Agent | None:
|
||||
cur = self._db.execute("SELECT id, name FROM agents WHERE name=?", (name,))
|
||||
row = cur.fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
return Agent(id=uuid.UUID(row[0]), name=row[1])
|
||||
|
||||
def get_agent_by_id(self, id: uuid.UUID) -> Agent | None:
|
||||
cur = self._db.execute("SELECT id, name FROM agents WHERE id=?", (str(id),))
|
||||
row = cur.fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
return Agent(id=uuid.UUID(row[0]), name=row[1])
|
||||
|
||||
def delete_agent(self, id: uuid.UUID) -> None:
|
||||
self._db.execute("DELETE FROM agents WHERE id=?", (str(id),))
|
||||
self._db.execute("DELETE FROM thread_participants WHERE agent_id=?", (str(id),))
|
||||
self._db.execute("DELETE FROM cursors WHERE agent_id=?", (str(id),))
|
||||
self._db.commit()
|
||||
|
||||
def list_agents(self) -> list[Agent]:
|
||||
cur = self._db.execute("SELECT id, name FROM agents")
|
||||
return [Agent(id=uuid.UUID(r[0]), name=r[1]) for r in cur.fetchall()]
|
||||
|
||||
# -- threads --
|
||||
|
||||
def create_thread(self, participants: set[Agent]) -> Thread:
|
||||
if len(participants) < 2:
|
||||
raise ValueError("need >=2 participants")
|
||||
tid = str(uuid.uuid4())
|
||||
ts = datetime.datetime.now(datetime.timezone.utc).isoformat()
|
||||
self._db.execute("BEGIN")
|
||||
try:
|
||||
self._db.execute("INSERT INTO threads (id, created_at) VALUES (?, ?)", (tid, ts))
|
||||
for p in participants:
|
||||
self._db.execute(
|
||||
"INSERT INTO thread_participants (thread_id, agent_id) VALUES (?, ?)",
|
||||
(tid, str(p.id)),
|
||||
)
|
||||
self._db.commit()
|
||||
except Exception:
|
||||
self._db.rollback()
|
||||
raise
|
||||
return Thread(
|
||||
id=uuid.UUID(tid),
|
||||
participants=frozenset(participants),
|
||||
created_at=datetime.datetime.fromisoformat(ts),
|
||||
)
|
||||
|
||||
def get_thread(self, id: uuid.UUID) -> Thread | None:
|
||||
cur = self._db.execute("SELECT id, created_at FROM threads WHERE id=?", (str(id),))
|
||||
row = cur.fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
tid_str, created_at_str = row
|
||||
created_at = datetime.datetime.fromisoformat(created_at_str)
|
||||
# fetch participants
|
||||
cur2 = self._db.execute(
|
||||
"SELECT agent_id FROM thread_participants WHERE thread_id=?", (tid_str,)
|
||||
)
|
||||
agent_ids = [r[0] for r in cur2.fetchall()]
|
||||
participants: set[Agent] = set()
|
||||
for aid in agent_ids:
|
||||
ag = self.get_agent_by_id(uuid.UUID(aid))
|
||||
if ag is not None:
|
||||
participants.add(ag)
|
||||
return Thread(
|
||||
id=uuid.UUID(tid_str),
|
||||
participants=frozenset(participants),
|
||||
created_at=created_at,
|
||||
)
|
||||
|
||||
def find_threads(self, containing: set[Agent]) -> list[Thread]:
|
||||
c = set(containing)
|
||||
if not c:
|
||||
cur = self._db.execute("SELECT id FROM threads")
|
||||
tids = [r[0] for r in cur.fetchall()]
|
||||
return [t for tid in tids if (t := self.get_thread(uuid.UUID(tid))) is not None]
|
||||
ids = [str(a.id) for a in c]
|
||||
placeholders = ",".join("?" for _ in ids)
|
||||
# GROUP BY HAVING COUNT(*) = len(containing)
|
||||
query = f"SELECT thread_id FROM thread_participants WHERE agent_id IN ({placeholders}) GROUP BY thread_id HAVING COUNT(*) = ?"
|
||||
cur = self._db.execute(query, (*ids, len(ids)))
|
||||
tids = [r[0] for r in cur.fetchall()]
|
||||
result: list[Thread] = []
|
||||
for tid in tids:
|
||||
t = self.get_thread(uuid.UUID(tid))
|
||||
if t is not None:
|
||||
result.append(t)
|
||||
return result
|
||||
|
||||
def add_participant(self, thread_id: uuid.UUID, agent: Agent) -> Thread:
|
||||
if self.get_thread(thread_id) is None:
|
||||
raise KeyError(thread_id)
|
||||
self._db.execute(
|
||||
"INSERT OR IGNORE INTO thread_participants (thread_id, agent_id) VALUES (?, ?)",
|
||||
(str(thread_id), str(agent.id)),
|
||||
)
|
||||
self._db.commit()
|
||||
t = self.get_thread(thread_id)
|
||||
assert t is not None
|
||||
return t
|
||||
|
||||
def remove_participant(self, thread_id: uuid.UUID, agent: Agent) -> Thread:
|
||||
if self.get_thread(thread_id) is None:
|
||||
raise KeyError(thread_id)
|
||||
self._db.execute(
|
||||
"DELETE FROM thread_participants WHERE thread_id=? AND agent_id=?",
|
||||
(str(thread_id), str(agent.id)),
|
||||
)
|
||||
self._db.commit()
|
||||
t = self.get_thread(thread_id)
|
||||
assert t is not None
|
||||
return t
|
||||
|
||||
def delete_thread(self, id: uuid.UUID) -> None:
|
||||
sid = str(id)
|
||||
self._db.execute("DELETE FROM threads WHERE id=?", (sid,))
|
||||
self._db.execute("DELETE FROM thread_participants WHERE thread_id=?", (sid,))
|
||||
self._db.execute("DELETE FROM messages WHERE thread_id=?", (sid,))
|
||||
self._db.execute("DELETE FROM cursors WHERE thread_id=?", (sid,))
|
||||
self._db.commit()
|
||||
|
||||
# -- messages --
|
||||
|
||||
def append_message(self, thread_id: uuid.UUID, sender: Agent, content: str) -> Message:
|
||||
# validate thread exists
|
||||
if self.get_thread(thread_id) is None:
|
||||
raise KeyError(thread_id)
|
||||
# validate sender in participants
|
||||
cur = self._db.execute(
|
||||
"SELECT 1 FROM thread_participants WHERE thread_id=? AND agent_id=?",
|
||||
(str(thread_id), str(sender.id)),
|
||||
)
|
||||
if cur.fetchone() is None:
|
||||
raise ValueError("sender not in thread")
|
||||
cur2 = self._db.execute(
|
||||
"SELECT COALESCE(MAX(seq), 0) FROM messages WHERE thread_id=?", (str(thread_id),)
|
||||
)
|
||||
max_seq = cur2.fetchone()[0]
|
||||
seq = int(max_seq) + 1
|
||||
mid = str(uuid.uuid4())
|
||||
ts = datetime.datetime.now(datetime.timezone.utc).isoformat()
|
||||
self._db.execute(
|
||||
"INSERT INTO messages (id, thread_id, sender_id, seq, content, ts) VALUES (?, ?, ?, ?, ?, ?)",
|
||||
(mid, str(thread_id), str(sender.id), seq, content, ts),
|
||||
)
|
||||
self._db.commit()
|
||||
return Message(
|
||||
id=uuid.UUID(mid),
|
||||
sender=sender,
|
||||
content=content,
|
||||
seq=seq,
|
||||
ts=datetime.datetime.fromisoformat(ts),
|
||||
)
|
||||
|
||||
def list_messages(self, thread_id: uuid.UUID, after_seq: int = 0) -> list[Message]:
|
||||
cur = self._db.execute(
|
||||
"SELECT id, sender_id, seq, content, ts FROM messages WHERE thread_id=? AND seq>? ORDER BY seq ASC",
|
||||
(str(thread_id), after_seq),
|
||||
)
|
||||
rows = cur.fetchall()
|
||||
result: list[Message] = []
|
||||
for mid, sender_id, seq, content, ts_str in rows:
|
||||
sender = self.get_agent_by_id(uuid.UUID(sender_id))
|
||||
# fallback if agent deleted - create placeholder
|
||||
if sender is None:
|
||||
sender = Agent(id=uuid.UUID(sender_id), name="unknown")
|
||||
result.append(
|
||||
Message(
|
||||
id=uuid.UUID(mid),
|
||||
sender=sender,
|
||||
content=content,
|
||||
seq=int(seq),
|
||||
ts=datetime.datetime.fromisoformat(ts_str),
|
||||
)
|
||||
)
|
||||
return result
|
||||
|
||||
def get_cursor(self, thread_id: uuid.UUID, agent: Agent) -> int:
|
||||
cur = self._db.execute(
|
||||
"SELECT seq FROM cursors WHERE thread_id=? AND agent_id=?",
|
||||
(str(thread_id), str(agent.id)),
|
||||
)
|
||||
row = cur.fetchone()
|
||||
if row is None:
|
||||
return 0
|
||||
return int(row[0])
|
||||
|
||||
def set_cursor(self, thread_id: uuid.UUID, agent: Agent, seq: int) -> None:
|
||||
self._db.execute(
|
||||
"INSERT OR REPLACE INTO cursors (thread_id, agent_id, seq) VALUES (?, ?, ?)",
|
||||
(str(thread_id), str(agent.id), seq),
|
||||
)
|
||||
self._db.commit()
|
||||
|
||||
def close(self) -> None:
|
||||
try:
|
||||
self._db.commit()
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
self._db.close()
|
||||
except Exception:
|
||||
pass
|
||||
_live_stores.discard(self)
|
||||
|
||||
def __del__(self) -> None:
|
||||
try:
|
||||
self.close()
|
||||
except Exception:
|
||||
pass
|
||||
@@ -0,0 +1,29 @@
|
||||
import tempfile, pathlib
|
||||
from agentmsgs.stores.sqlite import SQLiteStore
|
||||
|
||||
def test_sqlite_persists_across_handles():
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
p = pathlib.Path(d)/"test.db"
|
||||
s1 = SQLiteStore(p)
|
||||
a = s1.get_or_create_agent("Alice"); b = s1.get_or_create_agent("Bob")
|
||||
t = s1.create_thread({a,b})
|
||||
s1.append_message(t.id, a, "hi")
|
||||
# new handle same file
|
||||
s2 = SQLiteStore(p)
|
||||
assert s2.get_agent_by_name("Alice").id == a.id
|
||||
assert len(s2.list_messages(t.id)) == 1
|
||||
assert s2.find_threads({a})[0].id == t.id
|
||||
|
||||
def test_sqlite_concurrent_append():
|
||||
import pathlib, tempfile
|
||||
from agentmsgs.stores.sqlite import SQLiteStore
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
p = pathlib.Path(d)/"c.db"
|
||||
s1 = SQLiteStore(p); s2 = SQLiteStore(p)
|
||||
a = s1.get_or_create_agent("A"); b = s1.get_or_create_agent("B")
|
||||
# s2 sees same agents via file
|
||||
a2 = s2.get_agent_by_name("A"); b2 = s2.get_agent_by_name("B")
|
||||
t = s1.create_thread({a,b})
|
||||
s1.append_message(t.id, a, "from s1")
|
||||
s2.append_message(t.id, b2, "from s2")
|
||||
assert len(s1.list_messages(t.id)) == 2
|
||||
Reference in New Issue
Block a user