Skip to content

Commit a69de35

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
feat: add the delete method for the Agent Engine Task Store Service
PiperOrigin-RevId: 893515947
1 parent 62656c2 commit a69de35

File tree

4 files changed

+273
-0
lines changed

4 files changed

+273
-0
lines changed
Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
#
15+
# pylint: disable=protected-access,bad-continuation,missing-function-docstring
16+
17+
from tests.unit.vertexai.genai.replays import pytest_helper
18+
from vertexai._genai import types
19+
from google.genai import errors
20+
import pytest
21+
22+
23+
def test_get_a2a_task(client):
24+
# Use the autopush environment.
25+
client._api_client._http_options.base_url = (
26+
"https://us-central1-autopush-aiplatform.sandbox.googleapis.com/"
27+
)
28+
agent_engine = client.agent_engines.create()
29+
assert isinstance(agent_engine, types.AgentEngine)
30+
assert isinstance(agent_engine.api_resource, types.ReasoningEngine)
31+
# Use the internal API version for internal API access.
32+
client._api_client._http_options.api_version = "internal"
33+
34+
created_task = client.agent_engines.a2a_tasks.create(
35+
name=agent_engine.api_resource.name,
36+
a2a_task_id="task123",
37+
config=types.CreateAgentEngineTaskConfig(context_id="context123"),
38+
)
39+
assert isinstance(created_task, types.A2aTask)
40+
41+
# Test validity check.
42+
task = client.agent_engines.a2a_tasks.get(
43+
name=created_task.name,
44+
)
45+
assert task.name == f"{agent_engine.api_resource.name}/a2aTasks/task123"
46+
47+
client.agent_engines.a2a_tasks.delete(name=created_task.name)
48+
49+
with pytest.raises(errors.ClientError, match="404 NOT_FOUND"):
50+
client.agent_engines.a2a_tasks.get(name=created_task.name)
51+
52+
# Clean up resources.
53+
client.agent_engines.delete(name=agent_engine.api_resource.name, force=True)
54+
55+
56+
pytestmark = pytest_helper.setup(
57+
file=__file__,
58+
globals_for_file=globals(),
59+
)
60+
61+
62+
pytest_plugins = ("pytest_asyncio",)
63+
64+
65+
@pytest.mark.asyncio
66+
async def test_get_a2a_task_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+
77+
created_task = await client.aio.agent_engines.a2a_tasks.create(
78+
name=agent_engine.api_resource.name,
79+
a2a_task_id="task123",
80+
config=types.CreateAgentEngineTaskConfig(context_id="context123"),
81+
)
82+
assert isinstance(created_task, types.A2aTask)
83+
84+
# Test validity check.
85+
task = await client.aio.agent_engines.a2a_tasks.get(
86+
name=created_task.name,
87+
)
88+
assert task.name == f"{agent_engine.api_resource.name}/a2aTasks/task123"
89+
90+
await client.aio.agent_engines.a2a_tasks.delete(name=created_task.name)
91+
92+
with pytest.raises(errors.ClientError, match="404 NOT_FOUND"):
93+
await client.aio.agent_engines.a2a_tasks.get(name=created_task.name)
94+
95+
# Clean up resources.
96+
client.agent_engines.delete(name=agent_engine.api_resource.name, force=True)

vertexai/_genai/a2a_tasks.py

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,17 @@ def _CreateAgentEngineTaskRequestParameters_to_vertex(
8080
return to_object
8181

8282

83+
def _DeleteAgentEngineTaskRequestParameters_to_vertex(
84+
from_object: Union[dict[str, Any], object],
85+
parent_object: Optional[dict[str, Any]] = None,
86+
) -> dict[str, Any]:
87+
to_object: dict[str, Any] = {}
88+
if getv(from_object, ["name"]) is not None:
89+
setv(to_object, ["_url", "name"], getv(from_object, ["name"]))
90+
91+
return to_object
92+
93+
8394
def _GetAgentEngineTaskRequestParameters_to_vertex(
8495
from_object: Union[dict[str, Any], object],
8596
parent_object: Optional[dict[str, Any]] = None,
@@ -128,6 +139,62 @@ def _ListAgentEngineTasksRequestParameters_to_vertex(
128139

129140
class A2aTasks(_api_module.BaseModule):
130141

142+
def delete(
143+
self,
144+
*,
145+
name: str,
146+
config: Optional[types.DeleteAgentEngineTaskConfigOrDict] = None,
147+
) -> None:
148+
"""
149+
Deletes an agent engine task.
150+
151+
Args:
152+
name (str): Required. The name of the Agent Engine task to delete. Format:
153+
`projects/{project}/locations/{location}/reasoningEngines/{resource_id}/a2aTasks/{task_id}`.
154+
config (DeleteAgentEngineTaskConfig):
155+
Optional. Additional configurations for deleting the Agent Engine task.
156+
157+
Returns:
158+
None
159+
160+
"""
161+
162+
parameter_model = types._DeleteAgentEngineTaskRequestParameters(
163+
name=name,
164+
config=config,
165+
)
166+
167+
request_url_dict: Optional[dict[str, str]]
168+
if not self._api_client.vertexai:
169+
raise ValueError("This method is only supported in the Vertex AI client.")
170+
else:
171+
request_dict = _DeleteAgentEngineTaskRequestParameters_to_vertex(
172+
parameter_model
173+
)
174+
request_url_dict = request_dict.get("_url")
175+
if request_url_dict:
176+
path = "{name}".format_map(request_url_dict)
177+
else:
178+
path = "{name}"
179+
180+
query_params = request_dict.get("_query")
181+
if query_params:
182+
path = f"{path}?{urlencode(query_params)}"
183+
# TODO: remove the hack that pops config.
184+
request_dict.pop("config", None)
185+
186+
http_options: Optional[types.HttpOptions] = None
187+
if (
188+
parameter_model.config is not None
189+
and parameter_model.config.http_options is not None
190+
):
191+
http_options = parameter_model.config.http_options
192+
193+
request_dict = _common.convert_to_dict(request_dict)
194+
request_dict = _common.encode_unserializable_types(request_dict)
195+
196+
self._api_client.request("delete", path, request_dict, http_options)
197+
131198
def get(
132199
self,
133200
*,
@@ -372,6 +439,62 @@ def events(self) -> "a2a_task_events_module.A2aTaskEvents":
372439

373440
class AsyncA2aTasks(_api_module.BaseModule):
374441

442+
async def delete(
443+
self,
444+
*,
445+
name: str,
446+
config: Optional[types.DeleteAgentEngineTaskConfigOrDict] = None,
447+
) -> None:
448+
"""
449+
Deletes an agent engine task.
450+
451+
Args:
452+
name (str): Required. The name of the Agent Engine task to delete. Format:
453+
`projects/{project}/locations/{location}/reasoningEngines/{resource_id}/a2aTasks/{task_id}`.
454+
config (DeleteAgentEngineTaskConfig):
455+
Optional. Additional configurations for deleting the Agent Engine task.
456+
457+
Returns:
458+
None
459+
460+
"""
461+
462+
parameter_model = types._DeleteAgentEngineTaskRequestParameters(
463+
name=name,
464+
config=config,
465+
)
466+
467+
request_url_dict: Optional[dict[str, str]]
468+
if not self._api_client.vertexai:
469+
raise ValueError("This method is only supported in the Vertex AI client.")
470+
else:
471+
request_dict = _DeleteAgentEngineTaskRequestParameters_to_vertex(
472+
parameter_model
473+
)
474+
request_url_dict = request_dict.get("_url")
475+
if request_url_dict:
476+
path = "{name}".format_map(request_url_dict)
477+
else:
478+
path = "{name}"
479+
480+
query_params = request_dict.get("_query")
481+
if query_params:
482+
path = f"{path}?{urlencode(query_params)}"
483+
# TODO: remove the hack that pops config.
484+
request_dict.pop("config", None)
485+
486+
http_options: Optional[types.HttpOptions] = None
487+
if (
488+
parameter_model.config is not None
489+
and parameter_model.config.http_options is not None
490+
):
491+
http_options = parameter_model.config.http_options
492+
493+
request_dict = _common.convert_to_dict(request_dict)
494+
request_dict = _common.encode_unserializable_types(request_dict)
495+
496+
await self._api_client.async_request("delete", path, request_dict, http_options)
497+
375498
async def get(
376499
self,
377500
*,

vertexai/_genai/types/__init__.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
from .common import _DeleteAgentEngineRequestParameters
4646
from .common import _DeleteAgentEngineSandboxRequestParameters
4747
from .common import _DeleteAgentEngineSessionRequestParameters
48+
from .common import _DeleteAgentEngineTaskRequestParameters
4849
from .common import _DeleteDatasetRequestParameters
4950
from .common import _DeleteMultimodalDatasetRequestParameters
5051
from .common import _DeletePromptVersionRequestParameters
@@ -301,6 +302,9 @@
301302
from .common import DeleteAgentEngineSessionOperation
302303
from .common import DeleteAgentEngineSessionOperationDict
303304
from .common import DeleteAgentEngineSessionOperationOrDict
305+
from .common import DeleteAgentEngineTaskConfig
306+
from .common import DeleteAgentEngineTaskConfigDict
307+
from .common import DeleteAgentEngineTaskConfigOrDict
304308
from .common import DeletePromptConfig
305309
from .common import DeletePromptConfigDict
306310
from .common import DeletePromptConfigOrDict
@@ -1148,6 +1152,9 @@
11481152
from .common import WorkerPoolSpecOrDict
11491153

11501154
__all__ = [
1155+
"DeleteAgentEngineTaskConfig",
1156+
"DeleteAgentEngineTaskConfigDict",
1157+
"DeleteAgentEngineTaskConfigOrDict",
11511158
"GetAgentEngineTaskConfig",
11521159
"GetAgentEngineTaskConfigDict",
11531160
"GetAgentEngineTaskConfigOrDict",
@@ -2172,6 +2179,7 @@
21722179
"MessageDict",
21732180
"Importance",
21742181
"ParsedResponseUnion",
2182+
"_DeleteAgentEngineTaskRequestParameters",
21752183
"_GetAgentEngineTaskRequestParameters",
21762184
"_ListAgentEngineTasksRequestParameters",
21772185
"_CreateAgentEngineTaskRequestParameters",

vertexai/_genai/types/common.py

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -422,6 +422,52 @@ class OptimizationMethod(_common.CaseInSensitiveEnum):
422422
"""The data driven prompt optimizer designer for prompts from Android core API."""
423423

424424

425+
class DeleteAgentEngineTaskConfig(_common.BaseModel):
426+
"""Config for deleting an Agent Engine Task."""
427+
428+
http_options: Optional[genai_types.HttpOptions] = Field(
429+
default=None, description="""Used to override HTTP request options."""
430+
)
431+
432+
433+
class DeleteAgentEngineTaskConfigDict(TypedDict, total=False):
434+
"""Config for deleting an Agent Engine Task."""
435+
436+
http_options: Optional[genai_types.HttpOptionsDict]
437+
"""Used to override HTTP request options."""
438+
439+
440+
DeleteAgentEngineTaskConfigOrDict = Union[
441+
DeleteAgentEngineTaskConfig, DeleteAgentEngineTaskConfigDict
442+
]
443+
444+
445+
class _DeleteAgentEngineTaskRequestParameters(_common.BaseModel):
446+
"""Parameters for deleting an agent engine task."""
447+
448+
name: Optional[str] = Field(
449+
default=None, description="""Name of the agent engine task."""
450+
)
451+
config: Optional[DeleteAgentEngineTaskConfig] = Field(
452+
default=None, description=""""""
453+
)
454+
455+
456+
class _DeleteAgentEngineTaskRequestParametersDict(TypedDict, total=False):
457+
"""Parameters for deleting an agent engine task."""
458+
459+
name: Optional[str]
460+
"""Name of the agent engine task."""
461+
462+
config: Optional[DeleteAgentEngineTaskConfigDict]
463+
""""""
464+
465+
466+
_DeleteAgentEngineTaskRequestParametersOrDict = Union[
467+
_DeleteAgentEngineTaskRequestParameters, _DeleteAgentEngineTaskRequestParametersDict
468+
]
469+
470+
425471
class GetAgentEngineTaskConfig(_common.BaseModel):
426472
"""Config for getting an Agent Engine Task."""
427473

0 commit comments

Comments
 (0)