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