autogen/python/examples/inner_outer.py
2024-06-20 15:19:56 -04:00

61 lines
2.1 KiB
Python

import argparse
import asyncio
import logging
from dataclasses import dataclass
from agnext.application import SingleThreadedAgentRuntime
from agnext.components import TypeRoutedAgent, message_handler
from agnext.core import AgentId, CancellationToken
@dataclass
class MessageType:
body: str
sender: str
class Inner(TypeRoutedAgent): # type: ignore
def __init__(self) -> None: # type: ignore
super().__init__("The inner agent")
@message_handler() # type: ignore
async def on_new_message(self, message: MessageType, cancellation_token: CancellationToken) -> MessageType: # type: ignore
return MessageType(body=f"Inner: {message.body}", sender=self.metadata["name"])
class Outer(TypeRoutedAgent): # type: ignore
def __init__(self, inner: AgentId) -> None: # type: ignore
super().__init__("The outer agent")
self._inner = inner
@message_handler() # type: ignore
async def on_new_message(self, message: MessageType, cancellation_token: CancellationToken) -> MessageType: # type: ignore
inner_response = self.send_message(message, self._inner)
inner_message = await inner_response
assert isinstance(inner_message, MessageType)
return MessageType(body=f"Outer: {inner_message.body}", sender=self.metadata["name"])
async def main() -> None:
runtime = SingleThreadedAgentRuntime()
inner = runtime.register_and_get("inner", Inner)
outer = runtime.register_and_get("outer", lambda: Outer(inner))
response = runtime.send_message(MessageType(body="Hello", sender="external"), outer)
while not response.done():
await runtime.process_next()
print(await response)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Inner-Outter agent example.")
parser.add_argument("--verbose", action="store_true", help="Enable verbose logging.")
args = parser.parse_args()
if args.verbose:
logging.basicConfig(level=logging.WARNING)
logging.getLogger("agnext").setLevel(logging.DEBUG)
handler = logging.FileHandler("inner_outter.log")
logging.getLogger("agnext").addHandler(handler)
asyncio.run(main())