Improve the CoS server

Author: RussellLuo

coagent/core/util.py

@@ -91,10 +91,3 @@ def pretty_trace_tool_call(name: str, args: dict[str, any]):
     logger.opt(colors=True).trace(
         f"<green>{name}</green><magenta>({args_str})</magenta>"
     )
-
-
-def get_nats_servers():
-    NATS_SERVERS = os.environ.get("NATS_SERVERS", None)
-    if NATS_SERVERS and len(NATS_SERVERS) > 0:
-        NATS_SERVERS = NATS_SERVERS.split(",")
-    return NATS_SERVERS

coagent/cos/app.py

@@ -1,37 +1,88 @@
+import argparse
+import asyncio
+from hypercorn.asyncio import serve
+from hypercorn.config import Config
 from starlette.applications import Starlette
 from starlette.routing import Route
 
 from coagent.cos.runtime import CosRuntime
-from coagent.runtimes import NATSRuntime, LocalRuntime
-from coagent.core.util import get_nats_servers
+from coagent.core import set_stderr_logger
+from coagent.runtimes import NATSRuntime, HTTPRuntime
 
-runtime = CosRuntime(NATSRuntime.from_servers(servers=get_nats_servers()))
 
+class Application:
+    def __init__(self, server: str, auth: str):
+        if server.startswith("nats://"):
+            runtime = NATSRuntime.from_servers(server)
+        elif server.startswith(("http://", "https://")):
+            runtime = HTTPRuntime.from_server(server, auth)
+        else:
+            raise ValueError(f"Unsupported server: {server}")
+        self.runtime = CosRuntime(runtime)
 
-async def startup():
-    await runtime.start()
+    async def startup(self):
+        await self.runtime.start()
 
+    async def shutdown(self):
+        await self.runtime.stop()
 
-async def shutdown():
-    await runtime.stop()
-
-
-routes = [
-    Route("/runtime/register", runtime.register, methods=["POST"]),
-    Route("/runtime/channel/subscribe", runtime.subscribe, methods=["POST"]),
-    Route("/runtime/channel/publish", runtime.publish, methods=["POST"]),
-    Route("/runtime/channel/publish_multi", runtime.publish_multi, methods=["POST"]),
-]
-
-
-app = Starlette(debug=True, routes=routes, on_startup=[startup], on_shutdown=[shutdown])
+    @property
+    def starlette(self) -> Starlette:
+        routes = [
+            Route(
+                "/runtime/register",
+                self.runtime.register,
+                methods=["POST"],
+            ),
+            Route(
+                "/runtime/channel/subscribe",
+                self.runtime.subscribe,
+                methods=["POST"],
+            ),
+            Route(
+                "/runtime/channel/publish",
+                self.runtime.publish,
+                methods=["POST"],
+            ),
+            Route(
+                "/runtime/channel/publish_multi",
+                self.runtime.publish_multi,
+                methods=["POST"],
+            ),
+        ]
+        return Starlette(
+            debug=True,
+            routes=routes,
+            on_startup=[self.startup],
+            on_shutdown=[self.shutdown],
+        )
 
 
 if __name__ == "__main__":
-    import asyncio
-    from hypercorn.asyncio import serve
-    from hypercorn.config import Config
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--listen",
+        type=str,
+        default="127.0.0.1:8000",
+        help="The listen address of the CoS server. (Defaults to %(default)r)",
+    )
+    parser.add_argument(
+        "--server",
+        type=str,
+        default="nats://localhost:4222",
+        help="The runtime server address. (Defaults to %(default)r)",
+    )
+    parser.add_argument(
+        "--auth",
+        type=str,
+        default="",
+        help="The runtime server authentication token. (Defaults to %(default)r)",
+    )
+    args = parser.parse_args()
+
+    set_stderr_logger()
+    app = Application(args.server, args.auth).starlette
 
     config = Config()
-    config.bind = ["127.0.0.1:8000"]
+    config.bind = [args.listen]
     asyncio.run(serve(app, config))

examples/cos/README.md

@@ -23,13 +23,13 @@ node examples/cos/cos.js
 Finally, start a ping-pong client in the third terminal:
 
 ```bash
-python examples/ping-pong/client.py
+coagent server -H type:Ping
 ```
 
 or start a stream-ping-pong client:
 
 ```bash
-python examples/stream-ping-pong/client.py
+coagent stream_server -H type:Ping --chat
 ```
 
 
@@ -50,27 +50,33 @@ python examples/cos/cos.py
 Finally, start a ping-pong client in the third terminal:
 
 ```bash
-python examples/ping-pong/client.py
+coagent server -H type:Ping
 ```
 
 or start a stream-ping-pong client:
 
 ```bash
-python examples/stream-ping-pong/client.py
+coagent stream_server -H type:Ping --chat
 ```
 
 ### Run Go Agent
 
-Build Go agent:
+Start the CoS server in one terminal:
 
 ```bash
-cd examples/cos/goagent
-go build
+python coagent/cos/app.py
 ```
 
-Run a Go agent in one terminal:
+Then build and run the Go agent:
 
 ```bash
 cd examples/cos/goagent
+go build
 ./goagent
 ```
+
+Finally, start a ping-pong client in the third terminal:
+
+```bash
+coagent server -H type:Ping
+```