feat: add InMemoryStore (dict-backed Store)

src/agentmsgs/stores/memory.py

@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+import datetime
+import uuid
+from collections import defaultdict
+
+from agentmsgs.core.types import Agent, Message, Thread
+
+
+class InMemoryStore:
+    def __init__(self) -> None:
+        self._agents: dict[uuid.UUID, Agent] = {}
+        self._name_idx: dict[str, uuid.UUID] = {}
+        self._threads: dict[uuid.UUID, Thread] = {}
+        self._msgs: dict[uuid.UUID, list[Message]] = defaultdict(list)
+        self._cursors: dict[tuple[uuid.UUID, uuid.UUID], int] = {}
+
+    def get_or_create_agent(self, name: str) -> Agent:
+        if name in self._name_idx:
+            return self._agents[self._name_idx[name]]
+        return self.create_agent(name)
+
+    def create_agent(self, name: str) -> Agent:
+        if name in self._name_idx:
+            raise ValueError(f"agent name taken: {name}")
+        a = Agent(id=uuid.uuid4(), name=name)
+        self._agents[a.id] = a
+        self._name_idx[name] = a.id
+        return a
+
+    def get_agent_by_name(self, name: str) -> Agent | None:
+        uid = self._name_idx.get(name)
+        if uid is None:
+            return None
+        return self._agents.get(uid)
+
+    def get_agent_by_id(self, id: uuid.UUID) -> Agent | None:
+        return self._agents.get(id)
+
+    def delete_agent(self, id: uuid.UUID) -> None:
+        a = self._agents.pop(id, None)
+        if a is not None:
+            self._name_idx.pop(a.name, None)
+            for tid, t in list(self._threads.items()):
+                if a in t.participants:
+                    self._threads[tid] = Thread(
+                        id=t.id,
+                        participants=t.participants - {a},
+                        created_at=t.created_at,
+                    )
+
+    def list_agents(self) -> list[Agent]:
+        return list(self._agents.values())
+
+    def create_thread(self, participants: set[Agent]) -> Thread:
+        if len(participants) < 2:
+            raise ValueError("need >=2 participants")
+        t = Thread(
+            id=uuid.uuid4(),
+            participants=frozenset(participants),
+            created_at=datetime.datetime.now(datetime.timezone.utc),
+        )
+        self._threads[t.id] = t
+        return t
+
+    def get_thread(self, id: uuid.UUID) -> Thread | None:
+        return self._threads.get(id)
+
+    def find_threads(self, containing: set[Agent]) -> list[Thread]:
+        c = set(containing)
+        return [t for t in self._threads.values() if c <= t.participants]
+
+    def add_participant(self, thread_id: uuid.UUID, agent: Agent) -> Thread:
+        t = self._threads[thread_id]
+        nt = Thread(id=t.id, participants=t.participants | {agent}, created_at=t.created_at)
+        self._threads[thread_id] = nt
+        return nt
+
+    def remove_participant(self, thread_id: uuid.UUID, agent: Agent) -> Thread:
+        t = self._threads[thread_id]
+        nt = Thread(id=t.id, participants=t.participants - {agent}, created_at=t.created_at)
+        self._threads[thread_id] = nt
+        return nt
+
+    def delete_thread(self, id: uuid.UUID) -> None:
+        self._threads.pop(id, None)
+        self._msgs.pop(id, None)
+
+    def append_message(self, thread_id: uuid.UUID, sender: Agent, content: str) -> Message:
+        if thread_id not in self._threads:
+            raise KeyError(thread_id)
+        if sender not in self._threads[thread_id].participants:
+            raise ValueError("sender not in thread")
+        seq = len(self._msgs[thread_id]) + 1
+        m = Message(
+            id=uuid.uuid4(),
+            sender=sender,
+            content=content,
+            seq=seq,
+            ts=datetime.datetime.now(datetime.timezone.utc),
+        )
+        self._msgs[thread_id].append(m)
+        return m
+
+    def list_messages(self, thread_id: uuid.UUID, after_seq: int = 0) -> list[Message]:
+        return [m for m in self._msgs.get(thread_id, []) if m.seq > after_seq]
+
+    def get_cursor(self, thread_id: uuid.UUID, agent: Agent) -> int:
+        return self._cursors.get((thread_id, agent.id), 0)
+
+    def set_cursor(self, thread_id: uuid.UUID, agent: Agent, seq: int) -> None:
+        self._cursors[(thread_id, agent.id)] = seq

tests/test_store_memory.py

@@ -0,0 +1,33 @@
+from agentmsgs.stores.memory import InMemoryStore
+
+def test_memory_create_and_find_thread():
+    s = InMemoryStore()
+    alice = s.get_or_create_agent("Alice")
+    bob = s.get_or_create_agent("Bob")
+    t = s.create_thread({alice, bob})
+    assert s.get_thread(t.id) == t
+    assert t in s.find_threads({alice})
+    assert t in s.find_threads({alice, bob})
+    assert s.find_threads({alice, bob, s.get_or_create_agent("Charlie")}) == []
+
+def test_memory_soft_unique_name():
+    s = InMemoryStore()
+    a1 = s.get_or_create_agent("Alice")
+    a2 = s.get_or_create_agent("Alice")
+    assert a1.id == a2.id
+
+def test_memory_delete_account():
+    s = InMemoryStore()
+    a = s.get_or_create_agent("Alice")
+    s.delete_agent(a.id)
+    assert s.get_agent_by_id(a.id) is None
+    assert s.get_agent_by_name("Alice") is None
+
+def test_memory_append_and_list():
+    s = InMemoryStore()
+    a = s.get_or_create_agent("A"); b = s.get_or_create_agent("B")
+    t = s.create_thread({a,b})
+    m1 = s.append_message(t.id, a, "hi")
+    m2 = s.append_message(t.id, b, "yo")
+    assert m1.seq == 1 and m2.seq == 2
+    assert s.list_messages(t.id, after_seq=1) == [m2]