Update Runtime.register()

Author: RussellLuo

coagent/core/runtime.py

@@ -35,7 +35,7 @@ async def stop(self) -> None:
         await self.deregister()
         await self._channel.close()
 
-    async def register_spec(self, spec: AgentSpec) -> None:
+    async def register(self, spec: AgentSpec) -> None:
         spec.register(self)
 
         if self._discovery:

coagent/core/types.py

@@ -299,15 +299,8 @@ async def start(self) -> None:
     async def stop(self) -> None:
         pass
 
-    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:
+    async def register(self, spec: AgentSpec) -> None:
         pass
 
     @abc.abstractmethod

examples/app-builder/auto_team.py

@@ -1,7 +1,7 @@
 import asyncio
 
-from coagent.core import idle_loop, new, set_stderr_logger
 from coagent.agents import StreamChatAgent, tool
+from coagent.core import AgentSpec, idle_loop, new, set_stderr_logger
 from coagent.runtimes import NATSRuntime
 
 
@@ -27,9 +27,12 @@ async def transfer_to_qa(self):
             yield chunk
 
 
+auto_team = AgentSpec("auto_team", new(AutoTeam))
+
+
 async def main():
     async with NATSRuntime.from_servers() as runtime:
-        await runtime.register("auto_team", new(AutoTeam))
+        await runtime.register(auto_team)
         await idle_loop()
 
 

examples/app-builder/dev.py

@@ -1,8 +1,10 @@
 import asyncio
 from typing import AsyncIterator
 
-from coagent.agents.chat_agent import ChatMessage, chat
+from coagent.agents.chat_agent import ChatMessage
+from coagent.agents.util import chat_stream
 from coagent.core import (
+    AgentSpec,
     BaseAgent,
     Context,
     handler,
@@ -48,17 +50,20 @@ async def handle(
         ]
 
         reply = ""
-        async for chunk in chat(msgs):
+        async for chunk in chat_stream(msgs):
             yield ChatMessage(role="assistant", content=chunk.content)
             reply += chunk.content
 
         pretty_trace_agent_output("DevEngineer", reply)
         msg.messages.append(ChatMessage(role="user", content=reply))
 
 
+dev = AgentSpec("dev", new(DevEngineer))
+
+
 async def main():
     async with NATSRuntime.from_servers() as runtime:
-        await runtime.register("dev", new(DevEngineer))
+        await runtime.register(dev)
         await idle_loop()
 
 

examples/app-builder/qa.py

@@ -1,8 +1,10 @@
 import asyncio
 from typing import AsyncIterator
 
-from coagent.agents.chat_agent import ChatMessage, chat
+from coagent.agents.chat_agent import ChatMessage
+from coagent.agents.util import chat_stream
 from coagent.core import (
+    AgentSpec,
     BaseAgent,
     Context,
     handler,
@@ -67,17 +69,20 @@ async def handle(
         ]
 
         reply = ""
-        async for chunk in chat(msgs):
+        async for chunk in chat_stream(msgs):
             yield ChatMessage(role="assistant", content=chunk.content)
             reply += chunk.content
 
         pretty_trace_agent_output("QaEngineer", reply)
         msg.messages.append(ChatMessage(role="user", content=reply))
 
 
+qa = AgentSpec("qa", new(QaEngineer))
+
+
 async def main():
     async with NATSRuntime.from_servers() as runtime:
-        await runtime.register("qa", new(QaEngineer))
+        await runtime.register(qa)
         await idle_loop()
 
 

examples/app-builder/team.py

@@ -1,14 +1,17 @@
 import asyncio
 
 from coagent.agents import Sequential
-from coagent.core import new, set_stderr_logger
+from coagent.core import AgentSpec, new, set_stderr_logger
 from coagent.core.util import idle_loop
 from coagent.runtimes import NATSRuntime
 
 
+team = AgentSpec("team", new(Sequential, "dev", "qa"))
+
+
 async def main():
     async with NATSRuntime.from_servers() as runtime:
-        await runtime.register("team", new(Sequential, "dev", "qa"))
+        await runtime.register(team)
         await idle_loop()
 
 

examples/autogen/autogen.py

@@ -3,6 +3,7 @@
 
 from coagent.agents.chat_agent import ChatHistory, ChatMessage
 from coagent.core import (
+    AgentSpec,
     BaseAgent,
     Context,
     handler,
@@ -60,9 +61,12 @@ async def handle(self, msg: ChatHistory, ctx: Context) -> ChatHistory:
         return msg
 
 
+autogen = AgentSpec("autogen", new(AutoGenWeatherAgent))
+
+
 async def main():
     async with NATSRuntime.from_servers() as runtime:
-        await runtime.register("autogen", new(AutoGenWeatherAgent))
+        await runtime.register(autogen)
         await idle_loop()
 
 

examples/discovery/server.py

@@ -2,6 +2,7 @@
 import asyncio
 
 from coagent.core import (
+    AgentSpec,
     BaseAgent,
     Context,
     GenericMessage,
@@ -27,8 +28,10 @@ async def main(name: str, description: str, server: str):
     else:
         raise ValueError(f"Unsupported server: {server}")
 
+    employee = AgentSpec(name, new(Employee), description=description)
+
     async with runtime:
-        await runtime.register(name, new(Employee), description=description)
+        await runtime.register(employee)
         await idle_loop()
 
 

examples/opencsg/csghub.py

@@ -4,7 +4,7 @@
 
 from coagent.agents import StreamChatAgent, confirm, RunContext, tool
 from coagent.agents.messages import ChatMessage
-from coagent.core import idle_loop, new, set_stderr_logger
+from coagent.core import AgentSpec, idle_loop, new, set_stderr_logger
 from coagent.runtimes import NATSRuntime
 
 import httpx
@@ -58,8 +58,9 @@ async def search_model(
 
 
 async def main(name: str):
+    csghub = AgentSpec(name, new(CSGHub))
     async with NATSRuntime.from_servers() as runtime:
-        await runtime.register(name, new(CSGHub))
+        await runtime.register(csghub)
         await idle_loop()
 
 

examples/opencsg/dataflow.py

@@ -2,7 +2,7 @@
 import asyncio
 
 from coagent.agents import StreamChatAgent, RunContext, submit, tool
-from coagent.core import idle_loop, new, set_stderr_logger
+from coagent.core import AgentSpec, idle_loop, new, set_stderr_logger
 from coagent.runtimes import NATSRuntime
 
 import httpx
@@ -34,8 +34,9 @@ async def search_dataset(
 
 
 async def main(name: str):
+    dataflow = AgentSpec(name, new(DataFlow))
     async with NATSRuntime.from_servers() as runtime:
-        await runtime.register(name, new(DataFlow))
+        await runtime.register(dataflow)
         await idle_loop()