a2a init

cookbook/pocketflow_a2a/common/client/__init__.py

@@ -0,0 +1,4 @@
+from .client import A2AClient
+from .card_resolver import A2ACardResolver
+
+__all__ = ["A2AClient", "A2ACardResolver"]

cookbook/pocketflow_a2a/common/client/card_resolver.py

@@ -0,0 +1,21 @@
+import httpx
+from common.types import (
+    AgentCard,
+    A2AClientJSONError,
+)
+import json
+
+
+class A2ACardResolver:
+    def __init__(self, base_url, agent_card_path="/.well-known/agent.json"):
+        self.base_url = base_url.rstrip("/")
+        self.agent_card_path = agent_card_path.lstrip("/")
+
+    def get_agent_card(self) -> AgentCard:
+        with httpx.Client() as client:
+            response = client.get(self.base_url + "/" + self.agent_card_path)
+            response.raise_for_status()
+            try:
+                return AgentCard(**response.json())
+            except json.JSONDecodeError as e:
+                raise A2AClientJSONError(str(e)) from e

cookbook/pocketflow_a2a/common/client/client.py

@@ -0,0 +1,86 @@
+import httpx
+from httpx_sse import connect_sse
+from typing import Any, AsyncIterable
+from common.types import (
+    AgentCard,
+    GetTaskRequest,
+    SendTaskRequest,
+    SendTaskResponse,
+    JSONRPCRequest,
+    GetTaskResponse,
+    CancelTaskResponse,
+    CancelTaskRequest,
+    SetTaskPushNotificationRequest,
+    SetTaskPushNotificationResponse,
+    GetTaskPushNotificationRequest,
+    GetTaskPushNotificationResponse,
+    A2AClientHTTPError,
+    A2AClientJSONError,
+    SendTaskStreamingRequest,
+    SendTaskStreamingResponse,
+)
+import json
+
+
+class A2AClient:
+    def __init__(self, agent_card: AgentCard = None, url: str = None):
+        if agent_card:
+            self.url = agent_card.url
+        elif url:
+            self.url = url
+        else:
+            raise ValueError("Must provide either agent_card or url")
+
+    async def send_task(self, payload: dict[str, Any]) -> SendTaskResponse:
+        request = SendTaskRequest(params=payload)
+        return SendTaskResponse(**await self._send_request(request))
+
+    async def send_task_streaming(
+        self, payload: dict[str, Any]
+    ) -> AsyncIterable[SendTaskStreamingResponse]:
+        request = SendTaskStreamingRequest(params=payload)
+        with httpx.Client(timeout=None) as client:
+            with connect_sse(
+                client, "POST", self.url, json=request.model_dump()
+            ) as event_source:
+                try:
+                    for sse in event_source.iter_sse():
+                        yield SendTaskStreamingResponse(**json.loads(sse.data))
+                except json.JSONDecodeError as e:
+                    raise A2AClientJSONError(str(e)) from e
+                except httpx.RequestError as e:
+                    raise A2AClientHTTPError(400, str(e)) from e
+
+    async def _send_request(self, request: JSONRPCRequest) -> dict[str, Any]:
+        async with httpx.AsyncClient() as client:
+            try:
+                # Image generation could take time, adding timeout
+                response = await client.post(
+                    self.url, json=request.model_dump(), timeout=30
+                )
+                response.raise_for_status()
+                return response.json()
+            except httpx.HTTPStatusError as e:
+                raise A2AClientHTTPError(e.response.status_code, str(e)) from e
+            except json.JSONDecodeError as e:
+                raise A2AClientJSONError(str(e)) from e
+
+    async def get_task(self, payload: dict[str, Any]) -> GetTaskResponse:
+        request = GetTaskRequest(params=payload)
+        return GetTaskResponse(**await self._send_request(request))
+
+    async def cancel_task(self, payload: dict[str, Any]) -> CancelTaskResponse:
+        request = CancelTaskRequest(params=payload)
+        return CancelTaskResponse(**await self._send_request(request))
+
+    async def set_task_callback(
+        self, payload: dict[str, Any]
+    ) -> SetTaskPushNotificationResponse:
+        request = SetTaskPushNotificationRequest(params=payload)
+        return SetTaskPushNotificationResponse(**await self._send_request(request))
+
+    async def get_task_callback(
+        self, payload: dict[str, Any]
+    ) -> GetTaskPushNotificationResponse:
+        request = GetTaskPushNotificationRequest(params=payload)
+        return GetTaskPushNotificationResponse(**await self._send_request(request))