100 lines
3.5 KiB
Python
100 lines
3.5 KiB
Python
"""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
|