"""In-memory custom agent store for memory persistence mode.""" from __future__ import annotations from datetime import UTC, datetime from typing import Any from deerflow.persistence.agents.base import AgentStore def _now() -> str: return datetime.now(UTC).isoformat() def _copy(data: dict[str, Any]) -> dict[str, Any]: copied = dict(data) copied["builtin"] = copied.get("user_id") is None return copied class MemoryAgentStore(AgentStore): def __init__(self) -> None: self._agents: dict[str, dict[str, Any]] = {} @staticmethod def _is_visible(agent: dict[str, Any], user_id: str) -> bool: return agent.get("user_id") is None or agent.get("user_id") == user_id or bool(agent.get("published")) async def list_visible(self, user_id: str) -> list[dict[str, Any]]: rows = [_copy(agent) for agent in self._agents.values() if self._is_visible(agent, user_id)] rows.sort(key=lambda agent: agent.get("created_at") or "") return rows async def get_visible(self, agent_id: str, user_id: str) -> dict[str, Any] | None: agent = self._agents.get(agent_id) if agent is None or not self._is_visible(agent, user_id): return None return _copy(agent) async def get_any(self, agent_id: str) -> dict[str, Any] | None: agent = self._agents.get(agent_id) return _copy(agent) if agent is not None else None async def get_owned(self, agent_id: str, user_id: str) -> dict[str, Any] | None: agent = self._agents.get(agent_id) if agent is None or agent.get("user_id") != user_id: return None return _copy(agent) async def create(self, data: dict[str, Any]) -> dict[str, Any]: now = _now() row = { "id": data["id"], "name": data["name"], "description": data.get("description") or "", "user_id": data.get("user_id"), "published": bool(data.get("published", False)), "created_at": now, "updated_at": now, } self._agents[row["id"]] = row return _copy(row) async def ensure_builtin(self, data: dict[str, Any]) -> dict[str, Any]: row = self._agents.get(data["id"]) if row is None: row = { "id": data["id"], "name": data["name"], "description": data.get("description") or "", "user_id": None, "published": bool(data.get("published", True)), "created_at": _now(), "updated_at": _now(), } self._agents[row["id"]] = row elif row.get("user_id") is None: row["name"] = data["name"] row["description"] = data.get("description") or "" row["published"] = bool(data.get("published", True)) row["updated_at"] = _now() return _copy(row) async def update(self, agent_id: str, user_id: str, data: dict[str, Any]) -> dict[str, Any] | None: row = self._agents.get(agent_id) if row is None or row.get("user_id") != user_id: return None for key in ("name", "description", "published"): if key in data: row[key] = data[key] row["updated_at"] = _now() return _copy(row) async def delete(self, agent_id: str, user_id: str) -> bool: row = self._agents.get(agent_id) if row is None or row.get("user_id") != user_id: return False self._agents.pop(agent_id, None) return True