diff --git a/src/agentmsgs/stores/__init__.py b/src/agentmsgs/stores/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/agentmsgs/stores/memory.py b/src/agentmsgs/stores/memory.py new file mode 100644 index 0000000..918c434 --- /dev/null +++ b/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 diff --git a/tests/test_store_memory.py b/tests/test_store_memory.py new file mode 100644 index 0000000..0aaf21e --- /dev/null +++ b/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]