Skip to content

Commit eb696dd

Browse files
kthota-gkthota-g
committed
Refactor client
Change-Id: I92698288604fcb84dbf716880ddf64f5dda2ab75
1 parent ecf2ca5 commit eb696dd

File tree

2 files changed

+99
-60
lines changed

2 files changed

+99
-60
lines changed

‎examples/demo/test_client.py‎

Lines changed: 29 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,33 +1,37 @@
11
from a2a.client import A2AClient
22
from typing import Any
3+
import httpx
34

45

56
async def main() -> None:
6-
client = await A2AClient.get_client_from_agent_card_url(
7-
'http://localhost:9999'
8-
)
9-
send_task_payload: dict[str, Any] = {
10-
'id': '133',
11-
'message': {
12-
'role': 'user',
13-
'parts': [{'type': 'text', 'text': 'how much is 10 USD in INR?'}],
14-
},
15-
}
16-
17-
response = await client.send_task(payload=send_task_payload)
18-
print(response)
19-
20-
get_task_payload = {'id': '133'}
21-
get_response = await client.get_task(payload=get_task_payload)
22-
print(get_response)
23-
24-
stream_response = client.send_task_streaming(payload=send_task_payload)
25-
async for chunk in stream_response:
26-
print(chunk)
27-
28-
cancel_task_payload = {'id': '133'}
29-
cancel_response = await client.cancel_task(payload=cancel_task_payload)
30-
print(cancel_response)
7+
async with httpx.AsyncClient() as httpx_client:
8+
client = await A2AClient.get_client_from_agent_card_url(
9+
httpx_client, 'http://localhost:9999'
10+
)
11+
send_task_payload: dict[str, Any] = {
12+
'id': '133',
13+
'message': {
14+
'role': 'user',
15+
'parts': [
16+
{'type': 'text', 'text': 'how much is 10 USD in INR?'}
17+
],
18+
},
19+
}
20+
21+
response = await client.send_task(payload=send_task_payload)
22+
print(response)
23+
24+
get_task_payload = {'id': '133'}
25+
get_response = await client.get_task(payload=get_task_payload)
26+
print(get_response)
27+
28+
stream_response = client.send_task_streaming(payload=send_task_payload)
29+
async for chunk in stream_response:
30+
print(chunk)
31+
32+
cancel_task_payload = {'id': '133'}
33+
cancel_response = await client.cancel_task(payload=cancel_task_payload)
34+
print(cancel_response)
3135

3236

3337
if __name__ == '__main__':

‎examples/langgraph/test_client.py‎

Lines changed: 70 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,21 @@
11
from a2a.client import A2AClient
22
from typing import Any
33
from uuid import uuid4
4-
from a2a.types import SendTaskResponse, GetTaskResponse, SendTaskSuccessResponse, Task, TaskState
4+
from a2a.types import (
5+
SendTaskResponse,
6+
GetTaskResponse,
7+
SendTaskSuccessResponse,
8+
Task,
9+
TaskState,
10+
)
11+
import httpx
512

613
AGENT_URL = 'http://localhost:10000'
714

8-
def create_send_task_payload(task_id: str, text: str, session_id: str | None = None) -> dict[str, Any]:
15+
16+
def create_send_task_payload(
17+
task_id: str, text: str, session_id: str | None = None
18+
) -> dict[str, Any]:
919
"""Helper function to create the payload for sending a task."""
1020
payload: dict[str, Any] = {
1121
'id': task_id,
@@ -18,81 +28,106 @@ def create_send_task_payload(task_id: str, text: str, session_id: str | None = N
1828
payload['sessionId'] = session_id
1929
return payload
2030

31+
2132
def print_json_response(response: Any, description: str) -> None:
2233
"""Helper function to print the JSON representation of a response."""
23-
print(f"--- {description} ---")
34+
print(f'--- {description} ---')
2435
if hasattr(response, 'root'):
25-
print(f"{response.root.model_dump_json()}\n")
36+
print(f'{response.root.model_dump_json()}\n')
2637
else:
27-
print(f"{response.model_dump()}\n")
38+
print(f'{response.model_dump()}\n')
39+
2840

2941
async def run_single_turn_test(client: A2AClient) -> str:
3042
"""Runs a single-turn non-streaming test."""
3143
task_id: str = uuid4().hex
32-
send_payload = create_send_task_payload(task_id, 'how much is 10 USD in CAD?')
44+
send_payload = create_send_task_payload(
45+
task_id, 'how much is 10 USD in CAD?'
46+
)
3347
# Send Task
34-
send_response: SendTaskResponse = await client.send_task(payload=send_payload)
35-
print_json_response(send_response, "Single Turn Request Response")
48+
send_response: SendTaskResponse = await client.send_task(
49+
payload=send_payload
50+
)
51+
print_json_response(send_response, 'Single Turn Request Response')
3652

37-
print("---Query Task---")
53+
print('---Query Task---')
3854
# query the task
3955
task_id_payload = {'id': task_id}
40-
get_response: GetTaskResponse = await client.get_task(payload=task_id_payload)
41-
print_json_response(get_response, "Query Task Response")
42-
return task_id # Return task_id in case it's needed, though not used here
56+
get_response: GetTaskResponse = await client.get_task(
57+
payload=task_id_payload
58+
)
59+
print_json_response(get_response, 'Query Task Response')
60+
return task_id # Return task_id in case it's needed, though not used here
4361

4462

4563
async def run_streaming_test(client: A2AClient) -> None:
4664
"""Runs a single-turn streaming test."""
4765
task_id: str = uuid4().hex
48-
send_payload = create_send_task_payload(task_id, 'how much is 50 EUR in JPY?')
66+
send_payload = create_send_task_payload(
67+
task_id, 'how much is 50 EUR in JPY?'
68+
)
4969

50-
print("--- Single Turn Streaming Request ---")
70+
print('--- Single Turn Streaming Request ---')
5171
stream_response = client.send_task_streaming(payload=send_payload)
5272
async for chunk in stream_response:
53-
print_json_response(chunk, "Streaming Chunk")
73+
print_json_response(chunk, 'Streaming Chunk')
74+
5475

5576
async def run_multi_turn_test(client: A2AClient) -> None:
5677
"""Runs a multi-turn non-streaming test."""
57-
print("--- Multi-Turn Request ---")
78+
print('--- Multi-Turn Request ---')
5879
# --- First Turn ---
5980
task_id: str = uuid4().hex
60-
first_turn_payload = create_send_task_payload(task_id, 'how much is 100 USD?')
61-
first_turn_response: SendTaskResponse = await client.send_task(payload=first_turn_payload)
62-
print_json_response(first_turn_response, "Multi-Turn: First Turn Response")
63-
81+
first_turn_payload = create_send_task_payload(
82+
task_id, 'how much is 100 USD?'
83+
)
84+
first_turn_response: SendTaskResponse = await client.send_task(
85+
payload=first_turn_payload
86+
)
87+
print_json_response(first_turn_response, 'Multi-Turn: First Turn Response')
6488

6589
session_id: str | None = None
6690
if isinstance(first_turn_response.root, SendTaskSuccessResponse):
6791
task: Task = first_turn_response.root.result
68-
session_id = task.sessionId # Capture session ID
92+
session_id = task.sessionId # Capture session ID
6993

7094
# --- Second Turn (if input required) ---
7195
if task.status.state == TaskState.input_required and session_id:
72-
print("--- Multi-Turn: Second Turn (Input Required) ---")
73-
second_turn_payload = create_send_task_payload(task_id, 'in GBP', session_id)
74-
second_turn_response = await client.send_task(payload=second_turn_payload)
75-
print_json_response(second_turn_response, "Multi-Turn: Second Turn Response")
96+
print('--- Multi-Turn: Second Turn (Input Required) ---')
97+
second_turn_payload = create_send_task_payload(
98+
task_id, 'in GBP', session_id
99+
)
100+
second_turn_response = await client.send_task(
101+
payload=second_turn_payload
102+
)
103+
print_json_response(
104+
second_turn_response, 'Multi-Turn: Second Turn Response'
105+
)
76106
elif not session_id:
77-
print("Warning: Could not get session ID from first turn response.")
107+
print('Warning: Could not get session ID from first turn response.')
78108
else:
79-
print("First turn completed, no further input required for this test case.")
109+
print(
110+
'First turn completed, no further input required for this test case.'
111+
)
80112

81113

82114
async def main() -> None:
83115
"""Main function to run the tests."""
84-
print(f"Connecting to agent at {AGENT_URL}...")
116+
print(f'Connecting to agent at {AGENT_URL}...')
85117
try:
86-
client = await A2AClient.get_client_from_agent_card_url(AGENT_URL)
87-
print("Connection successful.")
118+
async with httpx.AsyncClient() as httpx_client:
119+
client = await A2AClient.get_client_from_agent_card_url(
120+
httpx_client, AGENT_URL
121+
)
122+
print('Connection successful.')
88123

89-
await run_single_turn_test(client)
90-
await run_streaming_test(client)
91-
await run_multi_turn_test(client)
124+
await run_single_turn_test(client)
125+
await run_streaming_test(client)
126+
await run_multi_turn_test(client)
92127

93128
except Exception as e:
94-
print(f"An error occurred: {e}")
95-
print("Ensure the agent server is running.")
129+
print(f'An error occurred: {e}')
130+
print('Ensure the agent server is running.')
96131

97132

98133
if __name__ == '__main__':

0 commit comments

Comments
 (0)