Skip to content

Commit 7d925a7

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 e2e81c9 commit 7d925a7

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_delete_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_delete_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
*,
@@ -423,6 +490,62 @@ def events(self) -> "a2a_task_events_module.A2aTaskEvents":
423490

424491
class AsyncA2aTasks(_api_module.BaseModule):
425492

493+
async def delete(
494+
self,
495+
*,
496+
name: str,
497+
config: Optional[types.DeleteAgentEngineTaskConfigOrDict] = None,
498+
) -> None:
499+
"""
500+
Deletes an agent engine task.
501+
502+
Args:
503+
name (str): Required. The name of the Agent Engine task to delete. Format:
504+
`projects/{project}/locations/{location}/reasoningEngines/{resource_id}/a2aTasks/{task_id}`.
505+
config (DeleteAgentEngineTaskConfig):
506+
Optional. Additional configurations for deleting the Agent Engine task.
507+
508+
Returns:
509+
None
510+
511+
"""
512+
513+
parameter_model = types._DeleteAgentEngineTaskRequestParameters(
514+
name=name,
515+
config=config,
516+
)
517+
518+
request_url_dict: Optional[dict[str, str]]
519+
if not self._api_client.vertexai:
520+
raise ValueError("This method is only supported in the Vertex AI client.")
521+
else:
522+
request_dict = _DeleteAgentEngineTaskRequestParameters_to_vertex(
523+
parameter_model
524+
)
525+
request_url_dict = request_dict.get("_url")
526+
if request_url_dict:
527+
path = "{name}".format_map(request_url_dict)
528+
else:
529+
path = "{name}"
530+
531+
query_params = request_dict.get("_query")
532+
if query_params:
533+
path = f"{path}?{urlencode(query_params)}"
534+
# TODO: remove the hack that pops config.
535+
request_dict.pop("config", None)
536+
537+
http_options: Optional[types.HttpOptions] = None
538+
if (
539+
parameter_model.config is not None
540+
and parameter_model.config.http_options is not None
541+
):
542+
http_options = parameter_model.config.http_options
543+
544+
request_dict = _common.convert_to_dict(request_dict)
545+
request_dict = _common.encode_unserializable_types(request_dict)
546+
547+
await self._api_client.async_request("delete", path, request_dict, http_options)
548+
426549
async def get(
427550
self,
428551
*,

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
@@ -302,6 +303,9 @@
302303
from .common import DeleteAgentEngineSessionOperation
303304
from .common import DeleteAgentEngineSessionOperationDict
304305
from .common import DeleteAgentEngineSessionOperationOrDict
306+
from .common import DeleteAgentEngineTaskConfig
307+
from .common import DeleteAgentEngineTaskConfigDict
308+
from .common import DeleteAgentEngineTaskConfigOrDict
305309
from .common import DeletePromptConfig
306310
from .common import DeletePromptConfigDict
307311
from .common import DeletePromptConfigOrDict
@@ -1218,6 +1222,9 @@
12181222
from .common import WorkerPoolSpecOrDict
12191223

12201224
__all__ = [
1225+
"DeleteAgentEngineTaskConfig",
1226+
"DeleteAgentEngineTaskConfigDict",
1227+
"DeleteAgentEngineTaskConfigOrDict",
12211228
"GetAgentEngineTaskConfig",
12221229
"GetAgentEngineTaskConfigDict",
12231230
"GetAgentEngineTaskConfigOrDict",
@@ -2299,6 +2306,7 @@
22992306
"MessageDict",
23002307
"Importance",
23012308
"ParsedResponseUnion",
2309+
"_DeleteAgentEngineTaskRequestParameters",
23022310
"_GetAgentEngineTaskRequestParameters",
23032311
"_ListAgentEngineTasksRequestParameters",
23042312
"_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)