Move python code to subdir (#98)

python/src/agnext/core/__init__.py

@@ -0,0 +1,24 @@
+"""
+The :mod:`agnext.core` module provides the foundational generic interfaces upon which all else is built. This module must not depend on any other module.
+"""
+
+from ._agent import Agent
+from ._agent_id import AgentId
+from ._agent_metadata import AgentMetadata
+from ._agent_props import AgentChildren
+from ._agent_proxy import AgentProxy
+from ._agent_runtime import AgentRuntime, AllNamespaces
+from ._base_agent import BaseAgent
+from ._cancellation_token import CancellationToken
+
+__all__ = [
+    "Agent",
+    "AgentId",
+    "AgentProxy",
+    "AgentMetadata",
+    "AgentRuntime",
+    "AllNamespaces",
+    "BaseAgent",
+    "CancellationToken",
+    "AgentChildren",
+]

python/src/agnext/core/_agent.py

@@ -0,0 +1,46 @@
+from typing import Any, Mapping, Protocol, runtime_checkable
+
+from ._agent_id import AgentId
+from ._agent_metadata import AgentMetadata
+from ._cancellation_token import CancellationToken
+
+
+@runtime_checkable
+class Agent(Protocol):
+    @property
+    def metadata(self) -> AgentMetadata:
+        """Metadata of the agent."""
+        ...
+
+    @property
+    def id(self) -> AgentId:
+        """ID of the agent."""
+        ...
+
+    async def on_message(self, message: Any, cancellation_token: CancellationToken) -> Any:
+        """Message handler for the agent. This should only be called by the runtime, not by other agents.
+
+        Args:
+            message (Any): Received message. Type is one of the types in `subscriptions`.
+            cancellation_token (CancellationToken): Cancellation token for the message.
+
+        Returns:
+            Any: Response to the message. Can be None.
+
+        Notes:
+            If there was a cancellation, this function should raise a `CancelledError`.
+        """
+        ...
+
+    def save_state(self) -> Mapping[str, Any]:
+        """Save the state of the agent. The result must be JSON serializable."""
+        ...
+
+    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.
+        """
+
+        ...

python/src/agnext/core/_agent_id.py

@@ -0,0 +1,31 @@
+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))
+
+    def __eq__(self, value: object) -> bool:
+        if not isinstance(value, AgentId):
+            return False
+        return self._name == value.name and self._namespace == value.namespace
+
+    @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

python/src/agnext/core/_agent_metadata.py

@@ -0,0 +1,8 @@
+from typing import Sequence, TypedDict
+
+
+class AgentMetadata(TypedDict):
+    name: str
+    namespace: str
+    description: str
+    subscriptions: Sequence[type]

python/src/agnext/core/_agent_props.py

@@ -0,0 +1,11 @@
+from typing import Protocol, Sequence, runtime_checkable
+
+from ._agent_id import AgentId
+
+
+@runtime_checkable
+class AgentChildren(Protocol):
+    @property
+    def children(self) -> Sequence[AgentId]:
+        """Ids of the children of the agent."""
+        ...

python/src/agnext/core/_agent_proxy.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from asyncio import Future
+from typing import TYPE_CHECKING, Any, Mapping
+
+from ._agent_id import AgentId
+from ._agent_metadata import AgentMetadata
+from ._cancellation_token import CancellationToken
+
+if TYPE_CHECKING:
+    from ._agent_runtime import AgentRuntime
+
+
+class AgentProxy:
+    def __init__(self, agent: AgentId, runtime: AgentRuntime):
+        self._agent = agent
+        self._runtime = runtime
+
+    @property
+    def id(self) -> AgentId:
+        """Target agent for this proxy"""
+        return self._agent
+
+    @property
+    def metadata(self) -> AgentMetadata:
+        """Metadata of the agent."""
+        return self._runtime.agent_metadata(self._agent)
+
+    def send_message(
+        self,
+        message: Any,
+        *,
+        sender: AgentId,
+        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)

python/src/agnext/core/_agent_runtime.py

@@ -0,0 +1,162 @@
+from __future__ import annotations
+
+from asyncio import Future
+from typing import Any, Callable, Mapping, Protocol, Sequence, Type, TypeVar, overload, runtime_checkable
+
+from ._agent import Agent
+from ._agent_id import AgentId
+from ._agent_metadata import AgentMetadata
+from ._agent_proxy import AgentProxy
+from ._cancellation_token import CancellationToken
+
+# Undeliverable - error
+
+T = TypeVar("T", bound=Agent)
+
+
+class AllNamespaces:
+    pass
+
+
+@runtime_checkable
+class AgentRuntime(Protocol):
+    # Returns the response of the message
+    def send_message(
+        self,
+        message: Any,
+        recipient: AgentId,
+        *,
+        sender: AgentId | None = None,
+        cancellation_token: CancellationToken | None = None,
+    ) -> Future[Any]: ...
+
+    # No responses from publishing
+    def publish_message(
+        self,
+        message: Any,
+        *,
+        namespace: str | None = None,
+        sender: AgentId | None = None,
+        cancellation_token: CancellationToken | None = None,
+    ) -> Future[None]: ...
+
+    @overload
+    def register(
+        self, name: str, agent_factory: Callable[[], T], *, valid_namespaces: Sequence[str] | Type[AllNamespaces] = ...
+    ) -> None: ...
+
+    @overload
+    def register(
+        self,
+        name: str,
+        agent_factory: Callable[[AgentRuntime, AgentId], T],
+        *,
+        valid_namespaces: Sequence[str] | Type[AllNamespaces] = ...,
+    ) -> None: ...
+
+    def register(
+        self,
+        name: str,
+        agent_factory: Callable[[], T] | Callable[[AgentRuntime, AgentId], T],
+        *,
+        valid_namespaces: Sequence[str] | Type[AllNamespaces] = AllNamespaces,
+    ) -> None:
+        """Register an agent factory with the runtime associated with a specific name. The name must be unique.
+
+        Args:
+            name (str): The name of the type agent this factory creates.
+            agent_factory (Callable[[], T] | Callable[[AgentRuntime, AgentId], T]): The factory that creates the agent.
+            valid_namespaces (Sequence[str] | Type[AllNamespaces], optional): Valid namespaces for this type. Defaults to AllNamespaces.
+
+
+        Example:
+            .. code-block:: python
+
+                runtime.register(
+                    "chat_agent",
+                    lambda: ChatCompletionAgent(
+                        description="A generic chat agent.",
+                        system_messages=[SystemMessage("You are a helpful assistant")],
+                        model_client=OpenAI(model="gpt-4o"),
+                        memory=BufferedChatMemory(buffer_size=10),
+                    ),
+                )
+
+        """
+
+        ...
+
+    def get(self, name: str, *, namespace: str = "default") -> AgentId: ...
+    def get_proxy(self, name: str, *, namespace: str = "default") -> AgentProxy: ...
+
+    @overload
+    def register_and_get(
+        self,
+        name: str,
+        agent_factory: Callable[[], T],
+        *,
+        namespace: str = "default",
+        valid_namespaces: Sequence[str] | Type[AllNamespaces] = ...,
+    ) -> AgentId: ...
+
+    @overload
+    def register_and_get(
+        self,
+        name: str,
+        agent_factory: Callable[[AgentRuntime, AgentId], T],
+        *,
+        namespace: str = "default",
+        valid_namespaces: Sequence[str] | Type[AllNamespaces] = ...,
+    ) -> AgentId: ...
+
+    def register_and_get(
+        self,
+        name: str,
+        agent_factory: Callable[[], T] | Callable[[AgentRuntime, AgentId], T],
+        *,
+        namespace: str = "default",
+        valid_namespaces: Sequence[str] | Type[AllNamespaces] = AllNamespaces,
+    ) -> AgentId:
+        self.register(name, agent_factory)
+        return self.get(name, namespace=namespace)
+
+    @overload
+    def register_and_get_proxy(
+        self,
+        name: str,
+        agent_factory: Callable[[], T],
+        *,
+        namespace: str = "default",
+        valid_namespaces: Sequence[str] | Type[AllNamespaces] = ...,
+    ) -> AgentProxy: ...
+
+    @overload
+    def register_and_get_proxy(
+        self,
+        name: str,
+        agent_factory: Callable[[AgentRuntime, AgentId], T],
+        *,
+        namespace: str = "default",
+        valid_namespaces: Sequence[str] | Type[AllNamespaces] = ...,
+    ) -> AgentProxy: ...
+
+    def register_and_get_proxy(
+        self,
+        name: str,
+        agent_factory: Callable[[], T] | Callable[[AgentRuntime, AgentId], T],
+        *,
+        namespace: str = "default",
+        valid_namespaces: Sequence[str] | Type[AllNamespaces] = AllNamespaces,
+    ) -> AgentProxy:
+        self.register(name, agent_factory)
+        return self.get_proxy(name, namespace=namespace)
+
+    def save_state(self) -> Mapping[str, Any]: ...
+
+    def load_state(self, state: Mapping[str, Any]) -> None: ...
+
+    def agent_metadata(self, agent: AgentId) -> AgentMetadata: ...
+
+    def agent_save_state(self, agent: AgentId) -> Mapping[str, Any]: ...
+
+    def agent_load_state(self, agent: AgentId, state: Mapping[str, Any]) -> None: ...

python/src/agnext/core/_base_agent.py

@@ -0,0 +1,106 @@
+import warnings
+from abc import ABC, abstractmethod
+from asyncio import Future
+from typing import Any, Mapping, Sequence
+
+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 BaseAgent(ABC, Agent):
+    @property
+    def metadata(self) -> AgentMetadata:
+        assert self._id is not None
+        return AgentMetadata(
+            namespace=self._id.namespace,
+            name=self._id.name,
+            description=self._description,
+            subscriptions=self._subscriptions,
+        )
+
+    def __init__(self, description: str, subscriptions: Sequence[type]) -> None:
+        self._runtime: AgentRuntime | None = None
+        self._id: AgentId | None = None
+        self._description = description
+        self._subscriptions = subscriptions
+
+    def bind_runtime(self, runtime: AgentRuntime) -> None:
+        if self._runtime is not None:
+            raise RuntimeError("Agent has already been bound to a runtime.")
+
+        self._runtime = runtime
+
+    def bind_id(self, agent_id: AgentId) -> None:
+        if self._id is not None:
+            raise RuntimeError("Agent has already been bound to an id.")
+        self._id = agent_id
+
+    @property
+    def name(self) -> str:
+        return self.id.name
+
+    @property
+    def id(self) -> AgentId:
+        if self._id is None:
+            raise RuntimeError("Agent has not been bound to an id.")
+
+        return self._id
+
+    @property
+    def runtime(self) -> AgentRuntime:
+        if self._runtime is None:
+            raise RuntimeError("Agent has not been bound to a runtime.")
+
+        return self._runtime
+
+    @abstractmethod
+    async def on_message(self, message: Any, cancellation_token: CancellationToken) -> Any: ...
+
+    # Returns the response of the message
+    def send_message(
+        self,
+        message: Any,
+        recipient: AgentId,
+        *,
+        cancellation_token: CancellationToken | None = None,
+    ) -> Future[Any]:
+        if self._runtime is None:
+            raise RuntimeError("Agent has not been bound to a runtime.")
+
+        if cancellation_token is None:
+            cancellation_token = CancellationToken()
+
+        future = self._runtime.send_message(
+            message,
+            sender=self.id,
+            recipient=recipient,
+            cancellation_token=cancellation_token,
+        )
+        cancellation_token.link_future(future)
+        return future
+
+    def publish_message(
+        self,
+        message: Any,
+        *,
+        cancellation_token: CancellationToken | None = None,
+    ) -> Future[None]:
+        if self._runtime is None:
+            raise RuntimeError("Agent has not been bound to a runtime.")
+
+        if cancellation_token is None:
+            cancellation_token = CancellationToken()
+
+        future = self._runtime.publish_message(message, sender=self.id, cancellation_token=cancellation_token)
+        return future
+
+    def save_state(self) -> Mapping[str, Any]:
+        warnings.warn("save_state not implemented", stacklevel=2)
+        return {}
+
+    def load_state(self, state: Mapping[str, Any]) -> None:
+        warnings.warn("load_state not implemented", stacklevel=2)
+        pass

python/src/agnext/core/_cancellation_token.py

@@ -0,0 +1,39 @@
+import threading
+from asyncio import Future
+from typing import Any, Callable, List
+
+
+class CancellationToken:
+    def __init__(self) -> None:
+        self._cancelled: bool = False
+        self._lock: threading.Lock = threading.Lock()
+        self._callbacks: List[Callable[[], None]] = []
+
+    def cancel(self) -> None:
+        with self._lock:
+            if not self._cancelled:
+                self._cancelled = True
+                for callback in self._callbacks:
+                    callback()
+
+    def is_cancelled(self) -> bool:
+        with self._lock:
+            return self._cancelled
+
+    def add_callback(self, callback: Callable[[], None]) -> None:
+        with self._lock:
+            if self._cancelled:
+                callback()
+            else:
+                self._callbacks.append(callback)
+
+    def link_future(self, future: Future[Any]) -> None:
+        with self._lock:
+            if self._cancelled:
+                future.cancel()
+            else:
+
+                def _cancel() -> None:
+                    future.cancel()
+
+                self._callbacks.append(_cancel)

python/src/agnext/core/exceptions.py

@@ -0,0 +1,17 @@
+__all__ = [
+    "CantHandleException",
+    "UndeliverableException",
+    "MessageDroppedException",
+]
+
+
+class CantHandleException(Exception):
+    """Raised when a handler can't handle the exception."""
+
+
+class UndeliverableException(Exception):
+    """Raised when a message can't be delivered."""
+
+
+class MessageDroppedException(Exception):
+    """Raised when a message is dropped."""

python/src/agnext/core/intervention.py

@@ -0,0 +1,36 @@
+from typing import Any, Awaitable, Callable, Protocol, final
+
+from agnext.core import AgentId
+
+__all__ = [
+    "DropMessage",
+    "InterventionFunction",
+    "InterventionHandler",
+    "DefaultInterventionHandler",
+]
+
+
+@final
+class DropMessage: ...
+
+
+InterventionFunction = Callable[[Any], Any | Awaitable[type[DropMessage]]]
+
+
+class InterventionHandler(Protocol):
+    async def on_send(self, message: Any, *, sender: AgentId | None, recipient: AgentId) -> Any | type[DropMessage]: ...
+    async def on_publish(self, message: Any, *, sender: AgentId | None) -> Any | type[DropMessage]: ...
+    async def on_response(
+        self, message: Any, *, sender: AgentId, recipient: AgentId | None
+    ) -> Any | type[DropMessage]: ...
+
+
+class DefaultInterventionHandler(InterventionHandler):
+    async def on_send(self, message: Any, *, sender: AgentId | None, recipient: AgentId) -> Any | type[DropMessage]:
+        return message
+
+    async def on_publish(self, message: Any, *, sender: AgentId | None) -> Any | type[DropMessage]:
+        return message
+
+    async def on_response(self, message: Any, *, sender: AgentId, recipient: AgentId | None) -> Any | type[DropMessage]:
+        return message