Add agent proxy (#84)

Author: jackgerrits

src/agnext/application/_single_threaded_agent_runtime.py

@@ -302,3 +302,9 @@ async def process_next(self) -> None:
 
     def agent_metadata(self, agent: Agent) -> AgentMetadata:
         return agent.metadata
+
+    def agent_save_state(self, agent: Agent) -> Mapping[str, Any]:
+        return agent.save_state()
+
+    def agent_load_state(self, agent: Agent, state: Mapping[str, Any]) -> None:
+        agent.load_state(state)

src/agnext/core/__init__.py

@@ -3,10 +3,11 @@
 """
 
 from ._agent import Agent
+from ._agent_id import AgentId
 from ._agent_metadata import AgentMetadata
 from ._agent_props import AgentChildren
 from ._agent_runtime import AgentRuntime
 from ._base_agent import BaseAgent
 from ._cancellation_token import CancellationToken
 
-__all__ = ["Agent", "AgentMetadata", "AgentRuntime", "BaseAgent", "CancellationToken", "AgentChildren"]
+__all__ = ["Agent", "AgentId", "AgentMetadata", "AgentRuntime", "BaseAgent", "CancellationToken", "AgentChildren"]

src/agnext/core/_agent_id.py

@@ -0,0 +1,26 @@
+from typing_extensions import Self
+
+
+class AgentId:
+    def __init__(self, name: str, namespace: str) -> None:
+        self._name = name
+        self._namespace = namespace
+
+    def __str__(self) -> str:
+        return f"{self._namespace}/{self._name}"
+
+    def __hash__(self) -> int:
+        return hash((self._namespace, self._name))
+
+    @classmethod
+    def from_str(cls, agent_id: str) -> Self:
+        namespace, name = agent_id.split("/")
+        return cls(name, namespace)
+
+    @property
+    def namespace(self) -> str:
+        return self._namespace
+
+    @property
+    def name(self) -> str:
+        return self._name

src/agnext/core/_agent_proxy.py

@@ -0,0 +1,50 @@
+from asyncio import Future
+from typing import Any, Mapping
+
+from ._agent import Agent
+from ._agent_id import AgentId
+from ._agent_metadata import AgentMetadata
+from ._agent_runtime import AgentRuntime
+from ._cancellation_token import CancellationToken
+
+
+class AgentProxy:
+    def __init__(self, agent: Agent, runtime: AgentRuntime):
+        self._agent = agent
+        self._runtime = runtime
+
+    @property
+    def id(self) -> AgentId:
+        """Target agent for this proxy"""
+        raise NotImplementedError
+
+    @property
+    def metadata(self) -> AgentMetadata:
+        """Metadata of the agent."""
+        return self._runtime.agent_metadata(self._agent)
+
+    def send_message(
+        self,
+        message: Any,
+        *,
+        sender: Agent,
+        cancellation_token: CancellationToken | None = None,
+    ) -> Future[Any]:
+        return self._runtime.send_message(
+            message,
+            recipient=self._agent,
+            sender=sender,
+            cancellation_token=cancellation_token,
+        )
+
+    def save_state(self) -> Mapping[str, Any]:
+        """Save the state of the agent. The result must be JSON serializable."""
+        return self._runtime.agent_save_state(self._agent)
+
+    def load_state(self, state: Mapping[str, Any]) -> None:
+        """Load in the state of the agent obtained from `save_state`.
+
+        Args:
+            state (Mapping[str, Any]): State of the agent. Must be JSON serializable.
+        """
+        self._runtime.agent_load_state(self._agent, state)

src/agnext/core/_agent_runtime.py

@@ -44,3 +44,7 @@ def save_state(self) -> Mapping[str, Any]: ...
     def load_state(self, state: Mapping[str, Any]) -> None: ...
 
     def agent_metadata(self, agent: Agent) -> AgentMetadata: ...
+
+    def agent_save_state(self, agent: Agent) -> Mapping[str, Any]: ...
+
+    def agent_load_state(self, agent: Agent, state: Mapping[str, Any]) -> None: ...