Skip to content

Commit 2a89939

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
feat: add async wrappers in the GenAI client for the Agent Engine Task Store Service
PiperOrigin-RevId: 877855591
1 parent 429a182 commit 2a89939

8 files changed

Lines changed: 424 additions & 68 deletions

‎tests/unit/vertexai/genai/replays/test_append_agent_engine_a2a_task_events.py‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
from tests.unit.vertexai.genai.replays import pytest_helper
1818
from vertexai._genai import types
19+
import pytest
1920

2021

2122
def test_append_simple_a2a_task_events(client):
@@ -57,3 +58,42 @@ def test_append_simple_a2a_task_events(client):
5758
file=__file__,
5859
globals_for_file=globals(),
5960
)
61+
62+
pytest_plugins = ("pytest_asyncio",)
63+
64+
65+
@pytest.mark.asyncio
66+
async def test_append_simple_a2a_task_events_async(client):
67+
# Use the autopush environment.
68+
client.aio._api_client._http_options.base_url = (
69+
"https://us-central1-autopush-aiplatform.sandbox.googleapis.com/"
70+
)
71+
agent_engine = client.agent_engines.create()
72+
assert isinstance(agent_engine, types.AgentEngine)
73+
assert isinstance(agent_engine.api_resource, types.ReasoningEngine)
74+
# Use the internal API version for internal API access.
75+
client.aio._api_client._http_options.api_version = "internal"
76+
task = await client.aio.agent_engines.a2a_tasks.create(
77+
name=agent_engine.api_resource.name,
78+
a2a_task_id="task123",
79+
config=types.CreateAgentEngineTaskConfig(context_id="context123"),
80+
)
81+
assert isinstance(task, types.A2aTask)
82+
83+
await client.aio.agent_engines.a2a_tasks.events.append(
84+
name=task.name,
85+
task_events=[
86+
types.TaskEvent(
87+
event_data=types.TaskEventData(
88+
metadata_change=types.TaskMetadataChange(
89+
new_metadata={"key1": "value1"}
90+
)
91+
),
92+
event_sequence_number=1,
93+
)
94+
],
95+
)
96+
97+
# Clean up resources.
98+
client.agent_engines.delete(name=agent_engine.api_resource.name, force=True)
99+

‎tests/unit/vertexai/genai/replays/test_create_agent_engine_a2a_task.py‎

Lines changed: 144 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -16,81 +16,159 @@
1616

1717
from tests.unit.vertexai.genai.replays import pytest_helper
1818
from vertexai._genai import types
19+
import pytest
1920

2021

2122
def test_create_simple_a2a_task(client):
22-
# Use the autopush environment.
23-
client._api_client._http_options.base_url = (
24-
"https://us-central1-autopush-aiplatform.sandbox.googleapis.com/"
25-
)
26-
agent_engine = client.agent_engines.create()
27-
assert isinstance(agent_engine, types.AgentEngine)
28-
assert isinstance(agent_engine.api_resource, types.ReasoningEngine)
29-
# Use the internal API version for internal API access.
30-
client._api_client._http_options.api_version = "internal"
23+
# Use the autopush environment.
24+
client._api_client._http_options.base_url = (
25+
"https://us-central1-autopush-aiplatform.sandbox.googleapis.com/"
26+
)
27+
agent_engine = client.agent_engines.create()
28+
assert isinstance(agent_engine, types.AgentEngine)
29+
assert isinstance(agent_engine.api_resource, types.ReasoningEngine)
30+
# Use the internal API version for internal API access.
31+
client._api_client._http_options.api_version = "internal"
3132

32-
task = client.agent_engines.a2a_tasks.create(
33-
name=agent_engine.api_resource.name,
34-
a2a_task_id="task123",
35-
config=types.CreateAgentEngineTaskConfig(
36-
context_id="context123",
37-
metadata={
38-
"key": "value",
39-
"key2": [{"key3": "value3", "key4": "value4"}],
40-
},
41-
status_details=types.TaskStatusDetails(
42-
task_message=types.TaskMessage(
43-
role="user",
44-
message_id="message123",
45-
parts=[
46-
types.Part(
47-
text="hello123",
48-
)
49-
],
50-
metadata={
51-
"key42": "value42",
52-
},
53-
),
54-
),
55-
output=types.TaskOutput(
56-
artifacts=[
57-
types.TaskArtifact(
58-
artifact_id="artifact123",
59-
display_name="display_name123",
60-
description="description123",
61-
parts=[
62-
types.Part(
63-
text="hello456",
64-
)
65-
],
66-
)
67-
],
68-
),
69-
),
70-
)
33+
task = client.agent_engines.a2a_tasks.create(
34+
name=agent_engine.api_resource.name,
35+
a2a_task_id="task123",
36+
config=types.CreateAgentEngineTaskConfig(
37+
context_id="context123",
38+
metadata={
39+
"key": "value",
40+
"key2": [{"key3": "value3", "key4": "value4"}],
41+
},
42+
status_details=types.TaskStatusDetails(
43+
task_message=types.TaskMessage(
44+
role="user",
45+
message_id="message123",
46+
parts=[
47+
types.Part(
48+
text="hello123",
49+
)
50+
],
51+
metadata={
52+
"key42": "value42",
53+
},
54+
),
55+
),
56+
output=types.TaskOutput(
57+
artifacts=[
58+
types.TaskArtifact(
59+
artifact_id="artifact123",
60+
display_name="display_name123",
61+
description="description123",
62+
parts=[
63+
types.Part(
64+
text="hello456",
65+
)
66+
],
67+
)
68+
],
69+
),
70+
),
71+
)
7172

72-
assert isinstance(task, types.A2aTask)
73-
assert task.name == f"{agent_engine.api_resource.name}/a2aTasks/task123"
74-
assert task.context_id == "context123"
75-
assert task.state == types.State.SUBMITTED
76-
assert task.status_details.task_message.role == "user"
77-
assert task.status_details.task_message.message_id == "message123"
78-
assert task.status_details.task_message.parts[0].text == "hello123"
79-
assert task.status_details.task_message.metadata["key42"] == "value42"
80-
assert task.output.artifacts[0].artifact_id == "artifact123"
81-
assert task.output.artifacts[0].display_name == "display_name123"
82-
assert task.output.artifacts[0].description == "description123"
83-
assert task.output.artifacts[0].parts[0].text == "hello456"
84-
assert task.metadata == {
85-
"key": "value",
86-
"key2": [{"key3": "value3", "key4": "value4"}],
87-
}
73+
assert isinstance(task, types.A2aTask)
74+
assert task.name == f"{agent_engine.api_resource.name}/a2aTasks/task123"
75+
assert task.context_id == "context123"
76+
assert task.state == types.State.SUBMITTED
77+
assert task.status_details.task_message.role == "user"
78+
assert task.status_details.task_message.message_id == "message123"
79+
assert task.status_details.task_message.parts[0].text == "hello123"
80+
assert task.status_details.task_message.metadata["key42"] == "value42"
81+
assert task.output.artifacts[0].artifact_id == "artifact123"
82+
assert task.output.artifacts[0].display_name == "display_name123"
83+
assert task.output.artifacts[0].description == "description123"
84+
assert task.output.artifacts[0].parts[0].text == "hello456"
85+
assert task.metadata == {
86+
"key": "value",
87+
"key2": [{"key3": "value3", "key4": "value4"}],
88+
}
8889

89-
# Clean up resources.
90-
client.agent_engines.delete(name=agent_engine.api_resource.name, force=True)
90+
# Clean up resources.
91+
client.agent_engines.delete(name=agent_engine.api_resource.name, force=True)
9192

9293

9394
pytestmark = pytest_helper.setup(
9495
file=__file__,
9596
globals_for_file=globals(),
9697
)
98+
99+
pytest_plugins = ("pytest_asyncio",)
100+
101+
102+
@pytest.mark.asyncio
103+
async def test_create_simple_a2a_task_async(client):
104+
# Use the autopush environment.
105+
client.aio._api_client._http_options.base_url = (
106+
"https://us-central1-autopush-aiplatform.sandbox.googleapis.com/"
107+
)
108+
agent_engine = client.agent_engines.create()
109+
assert isinstance(agent_engine, types.AgentEngine)
110+
assert isinstance(agent_engine.api_resource, types.ReasoningEngine)
111+
# Use the internal API version for internal API access.
112+
client.aio._api_client._http_options.api_version = "internal"
113+
114+
task = await client.aio.agent_engines.a2a_tasks.create(
115+
name=agent_engine.api_resource.name,
116+
a2a_task_id="task123",
117+
config=types.CreateAgentEngineTaskConfig(
118+
context_id="context123",
119+
metadata={
120+
"key": "value",
121+
"key2": [{"key3": "value3", "key4": "value4"}],
122+
},
123+
status_details=types.TaskStatusDetails(
124+
task_message=types.TaskMessage(
125+
role="user",
126+
message_id="message123",
127+
parts=[
128+
types.Part(
129+
text="hello123",
130+
)
131+
],
132+
metadata={
133+
"key42": "value42",
134+
},
135+
),
136+
),
137+
output=types.TaskOutput(
138+
artifacts=[
139+
types.TaskArtifact(
140+
artifact_id="artifact123",
141+
display_name="display_name123",
142+
description="description123",
143+
parts=[
144+
types.Part(
145+
text="hello456",
146+
)
147+
],
148+
)
149+
],
150+
),
151+
),
152+
)
153+
154+
assert isinstance(task, types.A2aTask)
155+
assert task.name == f"{agent_engine.api_resource.name}/a2aTasks/task123"
156+
assert task.context_id == "context123"
157+
assert task.state == types.State.SUBMITTED
158+
assert task.status_details.task_message.role == "user"
159+
assert task.status_details.task_message.message_id == "message123"
160+
assert task.status_details.task_message.parts[0].text == "hello123"
161+
assert task.status_details.task_message.metadata["key42"] == "value42"
162+
assert task.output.artifacts[0].artifact_id == "artifact123"
163+
assert task.output.artifacts[0].display_name == "display_name123"
164+
assert task.output.artifacts[0].description == "description123"
165+
assert task.output.artifacts[0].parts[0].text == "hello456"
166+
assert task.metadata == {
167+
"key": "value",
168+
"key2": [{"key3": "value3", "key4": "value4"}],
169+
}
170+
171+
# Clean up resources.
172+
client.agent_engines.delete(
173+
name=agent_engine.api_resource.name, force=True
174+
)

‎tests/unit/vertexai/genai/replays/test_get_agent_engine_a2a_task.py‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
from tests.unit.vertexai.genai.replays import pytest_helper
1818
from vertexai._genai import types
19+
import pytest
1920

2021

2122
def test_get_a2a_task(client):
@@ -52,3 +53,37 @@ def test_get_a2a_task(client):
5253
file=__file__,
5354
globals_for_file=globals(),
5455
)
56+
57+
58+
pytest_plugins = ("pytest_asyncio",)
59+
60+
61+
@pytest.mark.asyncio
62+
async def test_get_a2a_task_async(client):
63+
# Use the autopush environment.
64+
client.aio._api_client._http_options.base_url = (
65+
"https://us-central1-autopush-aiplatform.sandbox.googleapis.com/"
66+
)
67+
agent_engine = client.agent_engines.create()
68+
assert isinstance(agent_engine, types.AgentEngine)
69+
assert isinstance(agent_engine.api_resource, types.ReasoningEngine)
70+
# Use the internal API version for internal API access.
71+
client.aio._api_client._http_options.api_version = "internal"
72+
73+
created_task = await client.aio.agent_engines.a2a_tasks.create(
74+
name=agent_engine.api_resource.name,
75+
a2a_task_id="task123",
76+
config=types.CreateAgentEngineTaskConfig(context_id="context123"),
77+
)
78+
assert isinstance(created_task, types.A2aTask)
79+
80+
task = await client.aio.agent_engines.a2a_tasks.get(
81+
name=created_task.name,
82+
)
83+
assert isinstance(task, types.A2aTask)
84+
assert task.name == f"{agent_engine.api_resource.name}/a2aTasks/task123"
85+
assert task.context_id == "context123"
86+
assert task.state == types.State.SUBMITTED
87+
88+
# Clean up resources.
89+
client.agent_engines.delete(name=agent_engine.api_resource.name, force=True)

0 commit comments

Comments
 (0)