Refactor client

Change-Id: I92698288604fcb84dbf716880ddf64f5dda2ab75

Author: kthota-g

examples/demo/test_client.py

@@ -1,33 +1,37 @@
 from a2a.client import A2AClient
 from typing import Any
+import httpx
 
 
 async def main() -> None:
-    client = await A2AClient.get_client_from_agent_card_url(
-        'http://localhost:9999'
-    )
-    send_task_payload: dict[str, Any] = {
-        'id': '133',
-        'message': {
-            'role': 'user',
-            'parts': [{'type': 'text', 'text': 'how much is 10 USD in INR?'}],
-        },
-    }
-
-    response = await client.send_task(payload=send_task_payload)
-    print(response)
-
-    get_task_payload = {'id': '133'}
-    get_response = await client.get_task(payload=get_task_payload)
-    print(get_response)
-
-    stream_response = client.send_task_streaming(payload=send_task_payload)
-    async for chunk in stream_response:
-        print(chunk)
-
-    cancel_task_payload = {'id': '133'}
-    cancel_response = await client.cancel_task(payload=cancel_task_payload)
-    print(cancel_response)
+    async with httpx.AsyncClient() as httpx_client:
+        client = await A2AClient.get_client_from_agent_card_url(
+            httpx_client, 'http://localhost:9999'
+        )
+        send_task_payload: dict[str, Any] = {
+            'id': '133',
+            'message': {
+                'role': 'user',
+                'parts': [
+                    {'type': 'text', 'text': 'how much is 10 USD in INR?'}
+                ],
+            },
+        }
+
+        response = await client.send_task(payload=send_task_payload)
+        print(response)
+
+        get_task_payload = {'id': '133'}
+        get_response = await client.get_task(payload=get_task_payload)
+        print(get_response)
+
+        stream_response = client.send_task_streaming(payload=send_task_payload)
+        async for chunk in stream_response:
+            print(chunk)
+
+        cancel_task_payload = {'id': '133'}
+        cancel_response = await client.cancel_task(payload=cancel_task_payload)
+        print(cancel_response)
 
 
 if __name__ == '__main__':

examples/langgraph/test_client.py

@@ -1,11 +1,21 @@
 from a2a.client import A2AClient
 from typing import Any
 from uuid import uuid4
-from a2a.types import SendTaskResponse, GetTaskResponse, SendTaskSuccessResponse, Task, TaskState
+from a2a.types import (
+    SendTaskResponse,
+    GetTaskResponse,
+    SendTaskSuccessResponse,
+    Task,
+    TaskState,
+)
+import httpx
 
 AGENT_URL = 'http://localhost:10000'
 
-def create_send_task_payload(task_id: str, text: str, session_id: str | None = None) -> dict[str, Any]:
+
+def create_send_task_payload(
+    task_id: str, text: str, session_id: str | None = None
+) -> dict[str, Any]:
     """Helper function to create the payload for sending a task."""
     payload: dict[str, Any] = {
         'id': task_id,
@@ -18,81 +28,106 @@ def create_send_task_payload(task_id: str, text: str, session_id: str | None = N
         payload['sessionId'] = session_id
     return payload
 
+
 def print_json_response(response: Any, description: str) -> None:
     """Helper function to print the JSON representation of a response."""
-    print(f"--- {description} ---")
+    print(f'--- {description} ---')
     if hasattr(response, 'root'):
-        print(f"{response.root.model_dump_json()}\n")
+        print(f'{response.root.model_dump_json()}\n')
     else:
-        print(f"{response.model_dump()}\n")
+        print(f'{response.model_dump()}\n')
+
 
 async def run_single_turn_test(client: A2AClient) -> str:
     """Runs a single-turn non-streaming test."""
     task_id: str = uuid4().hex
-    send_payload = create_send_task_payload(task_id, 'how much is 10 USD in CAD?')
+    send_payload = create_send_task_payload(
+        task_id, 'how much is 10 USD in CAD?'
+    )
     # Send Task
-    send_response: SendTaskResponse = await client.send_task(payload=send_payload)
-    print_json_response(send_response, "Single Turn Request Response")
+    send_response: SendTaskResponse = await client.send_task(
+        payload=send_payload
+    )
+    print_json_response(send_response, 'Single Turn Request Response')
 
-    print("---Query Task---")
+    print('---Query Task---')
     # query the task
     task_id_payload = {'id': task_id}
-    get_response: GetTaskResponse = await client.get_task(payload=task_id_payload)
-    print_json_response(get_response, "Query Task Response")
-    return task_id # Return task_id in case it's needed, though not used here
+    get_response: GetTaskResponse = await client.get_task(
+        payload=task_id_payload
+    )
+    print_json_response(get_response, 'Query Task Response')
+    return task_id  # Return task_id in case it's needed, though not used here
 
 
 async def run_streaming_test(client: A2AClient) -> None:
     """Runs a single-turn streaming test."""
     task_id: str = uuid4().hex
-    send_payload = create_send_task_payload(task_id, 'how much is 50 EUR in JPY?')
+    send_payload = create_send_task_payload(
+        task_id, 'how much is 50 EUR in JPY?'
+    )
 
-    print("--- Single Turn Streaming Request ---")
+    print('--- Single Turn Streaming Request ---')
     stream_response = client.send_task_streaming(payload=send_payload)
     async for chunk in stream_response:
-        print_json_response(chunk, "Streaming Chunk")
+        print_json_response(chunk, 'Streaming Chunk')
+
 
 async def run_multi_turn_test(client: A2AClient) -> None:
     """Runs a multi-turn non-streaming test."""
-    print("--- Multi-Turn Request ---")
+    print('--- Multi-Turn Request ---')
     # --- First Turn ---
     task_id: str = uuid4().hex
-    first_turn_payload = create_send_task_payload(task_id, 'how much is 100 USD?')
-    first_turn_response: SendTaskResponse = await client.send_task(payload=first_turn_payload)
-    print_json_response(first_turn_response, "Multi-Turn: First Turn Response")
-
+    first_turn_payload = create_send_task_payload(
+        task_id, 'how much is 100 USD?'
+    )
+    first_turn_response: SendTaskResponse = await client.send_task(
+        payload=first_turn_payload
+    )
+    print_json_response(first_turn_response, 'Multi-Turn: First Turn Response')
 
     session_id: str | None = None
     if isinstance(first_turn_response.root, SendTaskSuccessResponse):
         task: Task = first_turn_response.root.result
-        session_id = task.sessionId # Capture session ID
+        session_id = task.sessionId  # Capture session ID
 
         # --- Second Turn (if input required) ---
         if task.status.state == TaskState.input_required and session_id:
-            print("--- Multi-Turn: Second Turn (Input Required) ---")
-            second_turn_payload = create_send_task_payload(task_id, 'in GBP', session_id)
-            second_turn_response = await client.send_task(payload=second_turn_payload)
-            print_json_response(second_turn_response, "Multi-Turn: Second Turn Response")
+            print('--- Multi-Turn: Second Turn (Input Required) ---')
+            second_turn_payload = create_send_task_payload(
+                task_id, 'in GBP', session_id
+            )
+            second_turn_response = await client.send_task(
+                payload=second_turn_payload
+            )
+            print_json_response(
+                second_turn_response, 'Multi-Turn: Second Turn Response'
+            )
         elif not session_id:
-             print("Warning: Could not get session ID from first turn response.")
+            print('Warning: Could not get session ID from first turn response.')
         else:
-            print("First turn completed, no further input required for this test case.")
+            print(
+                'First turn completed, no further input required for this test case.'
+            )
 
 
 async def main() -> None:
     """Main function to run the tests."""
-    print(f"Connecting to agent at {AGENT_URL}...")
+    print(f'Connecting to agent at {AGENT_URL}...')
     try:
-        client = await A2AClient.get_client_from_agent_card_url(AGENT_URL)
-        print("Connection successful.")
+        async with httpx.AsyncClient() as httpx_client:
+            client = await A2AClient.get_client_from_agent_card_url(
+                httpx_client, AGENT_URL
+            )
+            print('Connection successful.')
 
-        await run_single_turn_test(client)
-        await run_streaming_test(client)
-        await run_multi_turn_test(client)
+            await run_single_turn_test(client)
+            await run_streaming_test(client)
+            await run_multi_turn_test(client)
 
     except Exception as e:
-        print(f"An error occurred: {e}")
-        print("Ensure the agent server is running.")
+        print(f'An error occurred: {e}')
+        print('Ensure the agent server is running.')
 
 
 if __name__ == '__main__':

src/a2a/client/client.py

@@ -6,7 +6,7 @@
 
 import httpx
 
-from httpx_sse import connect_sse
+from httpx_sse import SSEError, aconnect_sse
 
 from a2a.client.errors import A2AClientHTTPError, A2AClientJSONError
 from a2a.types import (
@@ -35,30 +35,40 @@ class A2ACardResolver:
     """Agent Card resolver."""
 
     def __init__(
-        self, base_url: str, agent_card_path: str = '/.well-known/agent.json'
+        self,
+        httpx_client: httpx.AsyncClient,
+        base_url: str,
+        agent_card_path: str = '/.well-known/agent.json',
     ):
         self.base_url = base_url.rstrip('/')
         self.agent_card_path = agent_card_path.lstrip('/')
+        self.httpx_client = httpx_client
 
     async def get_agent_card(self) -> AgentCard:
-        async with httpx.AsyncClient() as client:
-            try:
-                response = await client.get(
-                    self.base_url + '/' + self.agent_card_path
-                )
-                response.raise_for_status()
-                return AgentCard.model_validate(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
+        try:
+            response = await self.httpx_client.get(
+                f'{self.base_url}/{self.agent_card_path}'
+            )
+            response.raise_for_status()
+            return AgentCard.model_validate(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
+        except httpx.RequestError as e:
+            raise A2AClientHTTPError(
+                503, f'Network communication error: {e}'
+            ) from e
 
 
 class A2AClient:
     """A2A Client."""
 
     def __init__(
-        self, agent_card: AgentCard | None = None, url: str | None = None
+        self,
+        httpx_client: httpx.AsyncClient,
+        agent_card: AgentCard | None = None,
+        url: str | None = None,
     ):
         if agent_card:
             self.url = agent_card.url
@@ -67,14 +77,20 @@ def __init__(
         else:
             raise ValueError('Must provide either agent_card or url')
 
+        self.httpx_client = httpx_client
+
     @classmethod
     async def get_client_from_agent_card_url(
-        cls, base_url: str, agent_card_path: str = '/.well-known/agent.json'
+        cls,
+        httpx_client: httpx.AsyncClient,
+        base_url: str,
+        agent_card_path: str = '/.well-known/agent.json',
     ) -> 'A2AClient':
+        """Get a A2A client for provided agent card URL."""
         agent_card: AgentCard = await A2ACardResolver(
-            base_url=base_url, agent_card_path=agent_card_path
+            httpx_client, base_url=base_url, agent_card_path=agent_card_path
         ).get_agent_card()
-        return A2AClient(agent_card=agent_card)
+        return A2AClient(httpx_client=httpx_client, agent_card=agent_card)
 
     async def send_task(
         self, payload: dict[str, Any], request_id: str | int = uuid4().hex
@@ -92,35 +108,38 @@ async def send_task_streaming(
         request = SendTaskStreamingRequest(
             id=request_id, params=TaskSendParams.model_validate(payload)
         )
-        with (
-            httpx.Client(timeout=None) as client,
-            connect_sse(
-                client, 'POST', self.url, json=request.model_dump(mode='json')
-            ) as event_source,
-        ):
+        async with aconnect_sse(
+            self.httpx_client,
+            'POST',
+            self.url,
+            json=request.model_dump(mode='json'),
+            timeout=None,
+        ) as event_source:
             try:
-                for sse in event_source.iter_sse():
+                async for sse in event_source.aiter_sse():
                     yield SendTaskStreamingResponse(**json.loads(sse.data))
+            except SSEError as e:
+                raise A2AClientHTTPError(
+                    400, f'Invalid SSE response or protocol error: {e}'
+                ) from e
             except json.JSONDecodeError as e:
                 raise A2AClientJSONError(str(e)) from e
             except httpx.RequestError as e:
-                raise A2AClientHTTPError(400, str(e)) from e
+                raise A2AClientHTTPError(
+                    503, f'Network communication error: {e}'
+                ) from e
 
     async def _send_request(self, request: A2ARequest) -> 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.root.model_dump(mode='json'),
-                    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
+        try:
+            response = await self.httpx_client.post(
+                self.url, json=request.root.model_dump(mode='json')
+            )
+            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], request_id: str | int = uuid4().hex