Add AgentSpec

Author: RussellLuo

coagent/core/__init__.py

@@ -11,6 +11,7 @@
 from .types import (
     Address,
     Agent,
+    AgentSpec,
     Constructor,
     Channel,
     MessageHeader,

coagent/core/runtime.py

@@ -8,8 +8,8 @@
 from .messages import StopIteration, Error
 from .factory import Factory, DeleteAgent
 from .types import (
+    AgentSpec,
     Channel,
-    Constructor,
     Runtime,
     Address,
     RawMessage,
@@ -35,19 +35,21 @@ async def stop(self) -> None:
         await self.deregister()
         await self._channel.close()
 
-    async def register(
-        self, name: str, constructor: Constructor, description: str = ""
-    ) -> None:
+    async def register_spec(self, spec: AgentSpec) -> None:
+        spec._runtime = self
+
         if self._discovery:
-            await self._discovery.register(name, constructor, description)
+            await self._discovery.register(
+                spec.name, spec.constructor, spec.description
+            )
 
-        if name in self._factories:
-            raise ValueError(f"Agent type {name} already registered")
+        if spec.name in self._factories:
+            raise ValueError(f"Agent type {spec.name} already registered")
 
-        factory = Factory(name, constructor)
+        factory = Factory(spec.name, spec.constructor)
         # We MUST set the channel and address manually.
-        factory.init(self._channel, Address(name=name))
-        self._factories[name] = factory
+        factory.init(self._channel, Address(name=spec.name))
+        self._factories[spec.name] = factory
 
         await factory.start()
 

coagent/core/types.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
 import abc
+from dataclasses import dataclass
 import enum
 from typing import Any, AsyncIterator, Awaitable, Callable, Type
+import uuid
 
 from pydantic import BaseModel, Field
 
@@ -238,6 +240,26 @@ async def new_reply_topic(self) -> str:
         pass
 
 
+@dataclass
+class AgentSpec:
+    """The specification of an agent."""
+
+    name: str
+    constructor: Constructor
+    description: str = ""
+
+    _runtime: Runtime | None = None
+
+    async def run_stream(self, msg: RawMessage) -> AsyncIterator[RawMessage]:
+        if self._runtime is None:
+            raise ValueError(f"AgentSpec {self.name} is not registered to a runtime.")
+
+        addr = Address(name=self.name, id=uuid.uuid4().hex)
+        result = self._runtime.channel.publish_multi(addr, msg)
+        async for chunk in result:
+            yield chunk
+
+
 class Runtime(abc.ABC):
     async def __aenter__(self):
         await self.start()
@@ -254,10 +276,15 @@ async def start(self) -> None:
     async def stop(self) -> None:
         pass
 
-    @abc.abstractmethod
     async def register(
         self, name: str, constructor: Constructor, description: str = ""
     ) -> None:
+        await self.register_spec(
+            AgentSpec(name=name, constructor=constructor, description=description)
+        )
+
+    @abc.abstractmethod
+    async def register_spec(self, spec: AgentSpec) -> None:
         pass
 
     @abc.abstractmethod