Python host and worker runtime for distributed agents. (#369)

* Python host runtime impl

* update

* ignore proto generated files

* move worker runtime to application

* Move example to samples

* Fix import

* fix

* update

* server client

* better shutdown

* fix doc conf

* add type

python/samples/marketing-agents/worker.py

@@ -1,8 +1,8 @@
 import asyncio
 import logging
 
+from agnext.application import WorkerAgentRuntime
 from agnext.core._serialization import MESSAGE_TYPE_REGISTRY
-from agnext.worker.worker_runtime import WorkerAgentRuntime
 from app import build_app
 from dotenv import load_dotenv
 from messages import ArticleCreated, AuditorAlert, AuditText, GraphicDesignCreated
@@ -18,7 +18,7 @@ async def main() -> None:
     MESSAGE_TYPE_REGISTRY.add_type(AuditText)
     MESSAGE_TYPE_REGISTRY.add_type(AuditorAlert)
     agnext_logger.info("1")
-    await runtime.setup_channel("localhost:5145")
+    await runtime.start("localhost:5145")
 
     agnext_logger.info("2")
 
@@ -30,7 +30,7 @@ async def main() -> None:
         await asyncio.sleep(1000000)
     except KeyboardInterrupt:
         pass
-    await runtime.close_channel()
+    await runtime.stop()
 
 
 if __name__ == "__main__":

python/samples/worker/run_host.py

@@ -0,0 +1,53 @@
+import asyncio
+import signal
+
+import grpc
+from agnext.application import HostRuntimeServicer
+from agnext.application.protos import agent_worker_pb2_grpc
+
+
+async def serve(server: grpc.aio.Server) -> None:  # type: ignore
+    await server.start()
+    print("Server started")
+    await server.wait_for_termination()
+
+
+async def main() -> None:
+    server = grpc.aio.server()
+    agent_worker_pb2_grpc.add_AgentRpcServicer_to_server(HostRuntimeServicer(), server)
+    server.add_insecure_port("[::]:50051")
+
+    # Set up signal handling for graceful shutdown
+    loop = asyncio.get_running_loop()
+
+    shutdown_event = asyncio.Event()
+
+    def signal_handler() -> None:
+        print("Received exit signal, shutting down gracefully...")
+        shutdown_event.set()
+
+    loop.add_signal_handler(signal.SIGINT, signal_handler)
+    loop.add_signal_handler(signal.SIGTERM, signal_handler)
+
+    # Start server in background task
+    serve_task = asyncio.create_task(serve(server))
+
+    # Wait for the signal to trigger the shutdown event
+    await shutdown_event.wait()
+
+    # Graceful shutdown
+    await server.stop(5)  # 5 second grace period
+    await serve_task
+    print("Server stopped")
+
+
+if __name__ == "__main__":
+    import logging
+
+    logging.basicConfig(level=logging.WARNING)
+    logging.getLogger("agnext").setLevel(logging.DEBUG)
+
+    try:
+        asyncio.run(main())
+    except KeyboardInterrupt:
+        print("Server shutdown interrupted.")

python/samples/worker/run_worker_pub_sub.py

@@ -0,0 +1,87 @@
+import asyncio
+import logging
+from dataclasses import dataclass
+
+from agnext.application import WorkerAgentRuntime
+from agnext.components import TypeRoutedAgent, message_handler
+from agnext.core import MESSAGE_TYPE_REGISTRY, MessageContext
+
+
+@dataclass
+class AskToGreet:
+    content: str
+
+
+@dataclass
+class Greeting:
+    content: str
+
+
+@dataclass
+class ReturnedGreeting:
+    content: str
+
+
+@dataclass
+class Feedback:
+    content: str
+
+
+@dataclass
+class ReturnedFeedback:
+    content: str
+
+
+class ReceiveAgent(TypeRoutedAgent):
+    def __init__(self) -> None:
+        super().__init__("Receive Agent")
+
+    @message_handler
+    async def on_greet(self, message: Greeting, ctx: MessageContext) -> None:
+        await self.publish_message(ReturnedGreeting(f"Returned greeting: {message.content}"))
+
+    @message_handler
+    async def on_feedback(self, message: Feedback, ctx: MessageContext) -> None:
+        await self.publish_message(ReturnedFeedback(f"Returned feedback: {message.content}"))
+
+
+class GreeterAgent(TypeRoutedAgent):
+    def __init__(self) -> None:
+        super().__init__("Greeter Agent")
+
+    @message_handler
+    async def on_ask(self, message: AskToGreet, ctx: MessageContext) -> None:
+        await self.publish_message(Greeting(f"Hello, {message.content}!"))
+
+    @message_handler
+    async def on_returned_greet(self, message: ReturnedGreeting, ctx: MessageContext) -> None:
+        await self.publish_message(Feedback(f"Feedback: {message.content}"))
+
+
+async def main() -> None:
+    runtime = WorkerAgentRuntime()
+    MESSAGE_TYPE_REGISTRY.add_type(Greeting)
+    MESSAGE_TYPE_REGISTRY.add_type(AskToGreet)
+    MESSAGE_TYPE_REGISTRY.add_type(Feedback)
+    MESSAGE_TYPE_REGISTRY.add_type(ReturnedGreeting)
+    MESSAGE_TYPE_REGISTRY.add_type(ReturnedFeedback)
+    await runtime.start(host_connection_string="localhost:50051")
+
+    await runtime.register("reciever", lambda: ReceiveAgent())
+    await runtime.register("greeter", lambda: GreeterAgent())
+
+    await runtime.publish_message(AskToGreet("Hello World!"), namespace="default")
+
+    # Just to keep the runtime running
+    try:
+        await asyncio.sleep(1000000)
+    except KeyboardInterrupt:
+        pass
+    await runtime.stop()
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.DEBUG)
+    logger = logging.getLogger("agnext")
+    logger.setLevel(logging.DEBUG)
+    asyncio.run(main())

python/samples/worker/run_worker_rpc.py

@@ -0,0 +1,74 @@
+import asyncio
+import logging
+from dataclasses import dataclass
+
+from agnext.application import WorkerAgentRuntime
+from agnext.components import TypeRoutedAgent, message_handler
+from agnext.core import MESSAGE_TYPE_REGISTRY, AgentId, MessageContext
+
+
+@dataclass
+class AskToGreet:
+    content: str
+
+
+@dataclass
+class Greeting:
+    content: str
+
+
+@dataclass
+class Feedback:
+    content: str
+
+
+class ReceiveAgent(TypeRoutedAgent):
+    def __init__(self) -> None:
+        super().__init__("Receive Agent")
+
+    @message_handler
+    async def on_greet(self, message: Greeting, ctx: MessageContext) -> Greeting:
+        return Greeting(content=f"Received: {message.content}")
+
+    @message_handler
+    async def on_feedback(self, message: Feedback, ctx: MessageContext) -> None:
+        print(f"Feedback received: {message.content}")
+
+
+class GreeterAgent(TypeRoutedAgent):
+    def __init__(self, receive_agent_id: AgentId) -> None:
+        super().__init__("Greeter Agent")
+        self._receive_agent_id = receive_agent_id
+
+    @message_handler
+    async def on_ask(self, message: AskToGreet, ctx: MessageContext) -> None:
+        response = await self.send_message(Greeting(f"Hello, {message.content}!"), recipient=self._receive_agent_id)
+        await self.publish_message(Feedback(f"Feedback: {response.content}"))
+
+
+async def main() -> None:
+    runtime = WorkerAgentRuntime()
+    MESSAGE_TYPE_REGISTRY.add_type(Greeting)
+    MESSAGE_TYPE_REGISTRY.add_type(AskToGreet)
+    MESSAGE_TYPE_REGISTRY.add_type(Feedback)
+    await runtime.start(host_connection_string="localhost:50051")
+
+    await runtime.register("reciever", lambda: ReceiveAgent())
+    reciever = await runtime.get("reciever")
+    await runtime.register("greeter", lambda: GreeterAgent(reciever))
+
+    await runtime.publish_message(AskToGreet("Hello World!"), namespace="default")
+
+    # Just to keep the runtime running
+    try:
+        await asyncio.sleep(1000000)
+    except KeyboardInterrupt:
+        pass
+    await runtime.stop()
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.DEBUG)
+    logger = logging.getLogger("agnext")
+    logger.setLevel(logging.DEBUG)
+    asyncio.run(main())