diff --git a/config/organism.yaml b/config/organism.yaml index 33b5e29..5c78c3d 100644 --- a/config/organism.yaml +++ b/config/organism.yaml @@ -15,10 +15,16 @@ max_concurrent_per_agent: 5 thread_scheduling: breadth-first listeners: - # The greeter listener responds to Greeting payloads + # The greeter receives Greeting, sends GreetingResponse to shouter - name: greeter payload_class: handlers.hello.Greeting handler: handlers.hello.handle_greeting - description: Responds with a greeting message - agent: false - broadcast: false + description: Receives greeting, forwards to shouter + agent: true + + # The shouter receives GreetingResponse, sends ShoutedResponse back to user + - name: shouter + payload_class: handlers.hello.GreetingResponse + handler: handlers.hello.handle_shout + description: Shouts the greeting in ALL CAPS + agent: true diff --git a/handlers/hello.py b/handlers/hello.py index dd7754e..562a56e 100644 --- a/handlers/hello.py +++ b/handlers/hello.py @@ -1,17 +1,26 @@ """ -hello.py — Hello World handler for testing the message pump. +hello.py — Multi-agent hello world handlers for testing the message pump. -This module provides: -- Greeting: payload class (what the handler receives) -- GreetingResponse: response payload (what the handler returns) -- handle_greeting: async handler function +This module demonstrates a multi-agent flow: + user -> greeter -> shouter -> user + +Payload classes: +- Greeting: Initial request with a name +- GreetingResponse: Greeter's response +- ShoutedResponse: Shouter's ALL CAPS version + +Handlers: +- handle_greeting: Receives Greeting, sends GreetingResponse to shouter +- handle_shout: Receives GreetingResponse, sends ShoutedResponse to original sender Usage in organism.yaml: listeners: - name: greeter payload_class: handlers.hello.Greeting handler: handlers.hello.handle_greeting - description: Responds with a greeting message + - name: shouter + payload_class: handlers.hello.GreetingResponse + handler: handlers.hello.handle_shout """ from dataclasses import dataclass @@ -35,44 +44,82 @@ class Greeting: @xmlify @dataclass class GreetingResponse: - """Outgoing greeting response.""" + """Greeter's response - will be forwarded to shouter.""" + message: str + original_sender: str # Track who started the conversation + + +@xmlify +@dataclass +class ShoutedResponse: + """Shouter's ALL CAPS response - sent back to original sender.""" message: str def wrap_in_envelope(payload_bytes: bytes, from_id: str, to_id: str, thread_id: str) -> bytes: - """Wrap a payload in a proper message envelope.""" + """Wrap a payload in a proper message envelope. + + Adds xmlns="" to payload to prevent it inheriting envelope namespace. + """ + payload_str = payload_bytes.decode('utf-8') + + # Add xmlns="" to payload root to keep it out of envelope namespace + if 'xmlns=' not in payload_str: + idx = payload_str.index('>') + payload_str = payload_str[:idx] + ' xmlns=""' + payload_str[idx:] + return f""" {from_id} {to_id} {thread_id} - {payload_bytes.decode('utf-8')} + {payload_str} """.encode('utf-8') async def handle_greeting(payload: Greeting, metadata: HandlerMetadata) -> bytes: """ - Handle an incoming Greeting and respond with a GreetingResponse. + Handle an incoming Greeting and forward GreetingResponse to shouter. - Args: - payload: The deserialized Greeting instance - metadata: Contains thread_id, from_id, own_name - - Returns: - XML bytes of the response envelope + Flow: user -> greeter -> shouter """ - # Create response - response = GreetingResponse(message=f"Hello, {payload.name}!") + # Create response, tracking original sender for later + response = GreetingResponse( + message=f"Hello, {payload.name}!", + original_sender=metadata.from_id, + ) # Serialize to XML response_tree = response.xml_value("GreetingResponse") payload_bytes = etree.tostring(response_tree, encoding='utf-8') - # Wrap in envelope - respond back to sender + # Forward to shouter (not back to sender) return wrap_in_envelope( payload_bytes=payload_bytes, from_id=metadata.own_name or "greeter", - to_id=metadata.from_id, # Send back to whoever sent the greeting + to_id="shouter", # Forward to shouter agent + thread_id=metadata.thread_id, + ) + + +async def handle_shout(payload: GreetingResponse, metadata: HandlerMetadata) -> bytes: + """ + Handle GreetingResponse by shouting it back to original sender. + + Flow: greeter -> shouter -> user + """ + # Create ALL CAPS response + response = ShoutedResponse(message=payload.message.upper()) + + # Serialize to XML + response_tree = response.xml_value("ShoutedResponse") + payload_bytes = etree.tostring(response_tree, encoding='utf-8') + + # Send back to original sender (tracked in payload) + return wrap_in_envelope( + payload_bytes=payload_bytes, + from_id=metadata.own_name or "shouter", + to_id=payload.original_sender, # Back to whoever started the conversation thread_id=metadata.thread_id, ) diff --git a/tests/test_pump_integration.py b/tests/test_pump_integration.py index 81b654a..0252ebb 100644 --- a/tests/test_pump_integration.py +++ b/tests/test_pump_integration.py @@ -14,7 +14,7 @@ from unittest.mock import AsyncMock, patch from agentserver.message_bus import StreamPump, bootstrap, MessageState from agentserver.message_bus.stream_pump import ConfigLoader, ListenerConfig, OrganismConfig, Listener -from handlers.hello import Greeting, GreetingResponse, handle_greeting, ENVELOPE_NS +from handlers.hello import Greeting, GreetingResponse, handle_greeting, handle_shout, ENVELOPE_NS def make_envelope(payload_xml: str, from_id: str, to_id: str, thread_id: str) -> bytes: @@ -50,11 +50,18 @@ class TestPumpBootstrap: config = ConfigLoader.load('config/organism.yaml') assert config.name == "hello-world" - assert len(config.listeners) == 1 + assert len(config.listeners) == 2 + + # Greeter listener assert config.listeners[0].name == "greeter" assert config.listeners[0].payload_class == Greeting assert config.listeners[0].handler == handle_greeting + # Shouter listener + assert config.listeners[1].name == "shouter" + assert config.listeners[1].payload_class == GreetingResponse + assert config.listeners[1].handler == handle_shout + @pytest.mark.asyncio async def test_bootstrap_creates_pump(self): """bootstrap() should create a configured pump.""" @@ -62,7 +69,9 @@ class TestPumpBootstrap: assert pump.config.name == "hello-world" assert "greeter.greeting" in pump.routing_table + assert "shouter.greetingresponse" in pump.routing_table assert pump.listeners["greeter"].payload_class == Greeting + assert pump.listeners["shouter"].payload_class == GreetingResponse @pytest.mark.asyncio async def test_bootstrap_generates_xsd(self):