feat: multicast a2a

Signed-off-by: Mauro Sardara <msardara@cisco.com>

Author: msardara

buf.gen.yaml

@@ -9,7 +9,7 @@ managed:
 plugins:
   # Generate python protobuf related code
   # Generates _pb2_slimrpc.py files
-  - local: protoc-gen-slimrpc-python
+  - local: /Users/msardara/repos/slim/data-plane/target/release/protoc-gen-slimrpc-python
     out: slima2a/types
     opt:
       - 'types_import=from a2a.grpc import a2a_pb2 as a2a__pb2'

dist/slim_bindings-1.2.0-py3-none-macosx_11_0_arm64.whl

@@ -1,6 +1,7 @@
 import argparse
 import asyncio
 import logging
+from typing import cast
 from uuid import uuid4
 
 import httpx
@@ -15,6 +16,7 @@
     Message,
     Part,
     Role,
+    TaskArtifactUpdateEvent,
     TextPart,
 )
 from a2a.utils.constants import (
@@ -30,6 +32,12 @@
 
 BASE_URL = "http://localhost:9999"
 
+# Two echo agent instances to multicast to
+AGENT_NAMES = [
+    "agntcy/demo/echo_agent1",
+    "agntcy/demo/echo_agent2",
+]
+
 logger = logging.getLogger(__name__)
 
 
@@ -66,7 +74,7 @@ async def main() -> None:
 
     client_config = ClientConfig(
         supported_transports=["JSONRPC", "slimrpc"],
-        streaming=args.stream,
+        streaming=True,
         httpx_client=httpx_client,
         slimrpc_channel_factory=slimrpc_channel_factory(slim_local_app, conn_id),
     )
@@ -78,7 +86,7 @@ async def main() -> None:
     agent_card: AgentCard
     match args.type:
         case "slimrpc":
-            agent_card = minimal_agent_card("agntcy/demo/echo_agent", ["slimrpc"])
+            agent_card = minimal_agent_card(",".join(AGENT_NAMES), ["slimrpc"])
         case "starlette":
             agent_card = await fetch_agent_card(
                 resolver=A2ACardResolver(
@@ -90,6 +98,19 @@ async def main() -> None:
             raise ValueError(f"Invalid client type: {args.type}")
 
     client = client_factory.create(card=agent_card)
+
+    if args.type == "slimrpc":
+        # Fetch agent cards from all servers in the group.
+        transport = cast(SRPCTransport, client._transport)  # type: ignore[attr-defined]
+        async for card in transport.get_all_cards():
+            logger.info(
+                f"agent card: {card.model_dump_json(indent=2, exclude_none=True)}"
+            )
+
+        # Fetch the real card from the first server so that
+        # client._card.capabilities.streaming=True before send_message.
+        await client.get_card()
+
     logger.info("A2AClient initialized.")
 
     response_text = await send_message(client, args.text)
@@ -106,12 +127,6 @@ def parse_arguments() -> argparse.Namespace:
         required=False,
         default="ERROR",
     )
-    parser.add_argument(
-        "--stream",
-        action="store_true",
-        required=False,
-        default=False,
-    )
     parser.add_argument(
         "--text",
         type=str,
@@ -144,24 +159,25 @@ async def send_message(
     )
     logger.info(f"associated request ({request_id}) with text: {text}")
 
-    output = ""
+    # Group output by task_id so responses from each server are kept separate.
+    outputs: dict[str, str] = {}
     try:
         async for event in client.send_message(request=request):
             if isinstance(event, Message):
+                outputs.setdefault("msg", "")
                 for part in event.parts:
                     if isinstance(part.root, TextPart):
-                        output += part.root.text
+                        outputs["msg"] += part.root.text
             else:
                 task, update = event
                 logger.info(f"task ({task.id}) status: {task.status.state}")
 
-                if task.status.state == "completed" and task.artifacts:
-                    for artifact in task.artifacts:
-                        for part in artifact.parts:
-                            if isinstance(part.root, TextPart):
-                                output += part.root.text
-
-                if update:
+                if isinstance(update, TaskArtifactUpdateEvent):
+                    outputs.setdefault(update.task_id, "")
+                    for part in update.artifact.parts:
+                        if isinstance(part.root, TextPart):
+                            outputs[update.task_id] += part.root.text
+                elif update:
                     logger.info(f"update: {update.model_dump(mode='json')}")
     except Exception as e:
         logger.error(
@@ -170,7 +186,7 @@ async def send_message(
         )
         raise RuntimeError("failed sending message or processing response") from e
 
-    return output
+    return "\n---\n".join(outputs.values())
 
 
 if __name__ == "__main__":

examples/echo_agent/client.py

@@ -59,7 +59,7 @@ async def main() -> None:
             service, local_app, local_name, conn_id = await setup_slim_client(
                 namespace="agntcy",
                 group="demo",
-                name="echo_agent",
+                name=args.instance,
             )
 
             # Create server
@@ -90,6 +90,13 @@ def parse_arguments() -> argparse.Namespace:
 
     parser.add_argument("--type", type=str, required=False, default="slimrpc")
 
+    parser.add_argument(
+        "--instance",
+        type=str,
+        required=False,
+        default="echo_agent",
+    )
+
     parser.add_argument(
         "--log-level",
         type=str,

examples/echo_agent/server.py

@@ -7,16 +7,15 @@
 logging.getLogger("asyncio").setLevel(logging.ERROR)
 
 # ruff: noqa: E402
+from typing import cast
+
 import httpx
-from a2a.client import (
-    Client,
-    ClientFactory,
-    minimal_agent_card,
-)
+from a2a.client import Client, ClientFactory, minimal_agent_card
 from a2a.types import (
     Message,
     Part,
     Role,
+    TaskArtifactUpdateEvent,
     TextPart,
 )
 
@@ -27,6 +26,14 @@
     slimrpc_channel_factory,
 )
 
+# Two travel planner instances to multicast to
+AGENT_NAMES = [
+    "agntcy/demo/travel_planner_agent1",
+    "agntcy/demo/travel_planner_agent2",
+]
+
+logger = logging.getLogger(__name__)
+
 
 def print_welcome_message() -> None:
     print("Welcome to the generic A2A client!")
@@ -51,26 +58,31 @@ async def interact_with_server(client: Client) -> None:
             parts=[Part(root=TextPart(text=user_input))],
         )
 
-        output = ""
+        # Group output by task_id so responses from each server are kept separate.
+        outputs: dict[str, str] = {}
         try:
             async for response in client.send_message(request=request):
                 if isinstance(response, Message):
+                    outputs.setdefault("msg", "")
                     for part in response.parts:
                         if isinstance(part.root, TextPart):
-                            output += part.root.text
+                            outputs["msg"] += part.root.text
                 else:
-                    task, _ = response
+                    _, update = response
 
-                    if task.status.state == "completed" and task.artifacts:
-                        for artifact in task.artifacts:
-                            for part in artifact.parts:
-                                if isinstance(part.root, TextPart):
-                                    output += part.root.text
+                    if isinstance(update, TaskArtifactUpdateEvent):
+                        outputs.setdefault(update.task_id, "")
+                        for part in update.artifact.parts:
+                            if isinstance(part.root, TextPart):
+                                outputs[update.task_id] += part.root.text
 
         except Exception as e:
             raise RuntimeError("failed sending message or processing response") from e
 
-        print(output, end="", flush=True)
+        for i, text in enumerate(outputs.values()):
+            if i > 0:
+                print("\n---")
+            print(text, end="", flush=True)
         await asyncio.sleep(0.1)
 
 
@@ -96,8 +108,19 @@ async def main() -> None:
 
     # mypy: the register API expects a different callable type; safe to ignore here.
     client_factory.register("slimrpc", SRPCTransport.create)  # type: ignore
-    agent_card = minimal_agent_card("agntcy/demo/travel_planner_agent", ["slimrpc"])
-    client = client_factory.create(card=agent_card)
+
+    client = client_factory.create(
+        card=minimal_agent_card(",".join(AGENT_NAMES), ["slimrpc"])
+    )
+
+    # Fetch agent cards from all servers in the group.
+    transport = cast(SRPCTransport, client._transport)  # type: ignore[attr-defined]
+    async for card in transport.get_all_cards():
+        logger.info(f"agent card: {card.model_dump_json(indent=2, exclude_none=True)}")
+
+    # Fetch the real card from the first server so that
+    # client._card.capabilities.streaming=True before send_message.
+    await client.get_card()
 
     await interact_with_server(client)
 

examples/travel_planner_agent/client.py

@@ -1,3 +1,4 @@
+import argparse
 import asyncio
 import logging
 
@@ -22,6 +23,8 @@
 
 
 async def main() -> None:
+    args = parse_arguments()
+
     skill = AgentSkill(
         id="travel_planner",
         name="travel planner agent",
@@ -52,7 +55,7 @@ async def main() -> None:
     service, local_app, local_name, conn_id = await setup_slim_client(
         namespace="agntcy",
         group="demo",
-        name="travel_planner_agent",
+        name=args.instance,
     )
 
     # Create server
@@ -67,5 +70,18 @@ async def main() -> None:
     await server.serve_async()
 
 
+def parse_arguments() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--instance",
+        type=str,
+        required=False,
+        default="travel_planner_agent",
+    )
+
+    return parser.parse_args()
+
+
 if __name__ == "__main__":
     asyncio.run(main())

examples/travel_planner_agent/server.py

@@ -9,7 +9,7 @@ readme = "README.md"
 requires-python = ">=3.10, <4.0"
 dependencies = [
     "a2a-sdk[telemetry]==0.3.20",
-    "slim-bindings~=1.1",
+    "slim-bindings",
 ]
 classifiers = [
     "Development Status :: 3 - Alpha",
@@ -44,6 +44,9 @@ examples = [
 [tool.uv]
 default-groups = ["linting", "testing", "examples"]
 
+[tool.uv.sources]
+slim-bindings = { path = "../slim/data-plane/target/wheels/slim_bindings-1.2.0-py3-none-macosx_11_0_arm64.whl" }
+
 [tool.ruff.format]
 exclude = ["*_pb2.py", "*_pb2.pyi", "*_pb2_slimrpc.py"]