Skip to content

Commit 883c01a

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
feat: Adds cancel_query_job to SDK for agent engine long running async tasks.
PiperOrigin-RevId: 896911226
1 parent 1e1faa3 commit 883c01a

4 files changed

Lines changed: 300 additions & 0 deletions

File tree

‎tests/unit/vertexai/genai/test_agent_engines.py‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3183,6 +3183,28 @@ def test_query_agent_engine_async(self):
31833183
None,
31843184
)
31853185

3186+
def test_cancel_query_job_agent_engine(self):
3187+
with mock.patch.object(
3188+
self.client.agent_engines._api_client, "request"
3189+
) as request_mock:
3190+
request_mock.return_value = genai_types.HttpResponse(body="{}")
3191+
3192+
result = self.client.agent_engines.cancel_query_job(
3193+
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
3194+
config={"operation_name": _TEST_AGENT_ENGINE_OPERATION_NAME},
3195+
)
3196+
3197+
assert isinstance(result, _genai_types.CancelQueryJobResult)
3198+
request_mock.assert_called_with(
3199+
"post",
3200+
f"{_TEST_AGENT_ENGINE_RESOURCE_NAME}:cancelAsyncQuery",
3201+
{
3202+
"_url": {"name": _TEST_AGENT_ENGINE_RESOURCE_NAME},
3203+
"operationName": _TEST_AGENT_ENGINE_OPERATION_NAME,
3204+
},
3205+
None,
3206+
)
3207+
31863208
def test_check_query_job_agent_engine(self):
31873209
with mock.patch.object(
31883210
self.client.agent_engines._api_client, "request"

‎vertexai/_genai/agent_engines.py‎

Lines changed: 192 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,38 @@
4949
logger.setLevel(logging.INFO)
5050

5151

52+
def _CancelQueryJobAgentEngineConfig_to_vertex(
53+
from_object: Union[dict[str, Any], object],
54+
parent_object: Optional[dict[str, Any]] = None,
55+
) -> dict[str, Any]:
56+
to_object: dict[str, Any] = {}
57+
58+
if getv(from_object, ["operation_name"]) is not None:
59+
setv(parent_object, ["operationName"], getv(from_object, ["operation_name"]))
60+
61+
return to_object
62+
63+
64+
def _CancelQueryJobAgentEngineRequestParameters_to_vertex(
65+
from_object: Union[dict[str, Any], object],
66+
parent_object: Optional[dict[str, Any]] = None,
67+
) -> dict[str, Any]:
68+
to_object: dict[str, Any] = {}
69+
if getv(from_object, ["name"]) is not None:
70+
setv(to_object, ["_url", "name"], getv(from_object, ["name"]))
71+
72+
if getv(from_object, ["config"]) is not None:
73+
setv(
74+
to_object,
75+
["config"],
76+
_CancelQueryJobAgentEngineConfig_to_vertex(
77+
getv(from_object, ["config"]), to_object
78+
),
79+
)
80+
81+
return to_object
82+
83+
5284
def _CheckQueryJobAgentEngineConfig_to_vertex(
5385
from_object: Union[dict[str, Any], object],
5486
parent_object: Optional[dict[str, Any]] = None,
@@ -410,6 +442,85 @@ def _UpdateAgentEngineRequestParameters_to_vertex(
410442

411443
class AgentEngines(_api_module.BaseModule):
412444

445+
def cancel_query_job(
446+
self,
447+
*,
448+
name: str,
449+
config: Optional[types.CancelQueryJobAgentEngineConfigOrDict] = None,
450+
) -> types.CancelQueryJobResult:
451+
"""
452+
Cancels a long-running query job on an Agent Engine.
453+
454+
Args:
455+
name (str):
456+
Required. The reasoning engine resource name.
457+
config (CancelQueryJobAgentEngineConfigOrDict):
458+
Optional. The configuration for the cancel_query_job.
459+
460+
"""
461+
462+
parameter_model = types._CancelQueryJobAgentEngineRequestParameters(
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 = _CancelQueryJobAgentEngineRequestParameters_to_vertex(
472+
parameter_model
473+
)
474+
request_url_dict = request_dict.get("_url")
475+
if request_url_dict:
476+
path = "{name}:cancelAsyncQuery".format_map(request_url_dict)
477+
else:
478+
path = "{name}:cancelAsyncQuery"
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+
response = self._api_client.request("post", path, request_dict, http_options)
497+
498+
response_dict = {} if not response.body else json.loads(response.body)
499+
500+
return_value = types.CancelQueryJobResult._from_response(
501+
response=response_dict,
502+
kwargs=(
503+
{
504+
"config": {
505+
"response_schema": getattr(
506+
parameter_model.config, "response_schema", None
507+
),
508+
"response_json_schema": getattr(
509+
parameter_model.config, "response_json_schema", None
510+
),
511+
"include_all_fields": getattr(
512+
parameter_model.config, "include_all_fields", None
513+
),
514+
}
515+
}
516+
if getattr(parameter_model, "config", None)
517+
else {}
518+
),
519+
)
520+
521+
self._api_client._verify_response(return_value)
522+
return return_value
523+
413524
def _check_query_job(
414525
self,
415526
*,
@@ -2643,6 +2754,87 @@ def list_session_events(
26432754

26442755
class AsyncAgentEngines(_api_module.BaseModule):
26452756

2757+
async def cancel_query_job(
2758+
self,
2759+
*,
2760+
name: str,
2761+
config: Optional[types.CancelQueryJobAgentEngineConfigOrDict] = None,
2762+
) -> types.CancelQueryJobResult:
2763+
"""
2764+
Cancels a long-running query job on an Agent Engine.
2765+
2766+
Args:
2767+
name (str):
2768+
Required. The reasoning engine resource name.
2769+
config (CancelQueryJobAgentEngineConfigOrDict):
2770+
Optional. The configuration for the cancel_query_job.
2771+
2772+
"""
2773+
2774+
parameter_model = types._CancelQueryJobAgentEngineRequestParameters(
2775+
name=name,
2776+
config=config,
2777+
)
2778+
2779+
request_url_dict: Optional[dict[str, str]]
2780+
if not self._api_client.vertexai:
2781+
raise ValueError("This method is only supported in the Vertex AI client.")
2782+
else:
2783+
request_dict = _CancelQueryJobAgentEngineRequestParameters_to_vertex(
2784+
parameter_model
2785+
)
2786+
request_url_dict = request_dict.get("_url")
2787+
if request_url_dict:
2788+
path = "{name}:cancelAsyncQuery".format_map(request_url_dict)
2789+
else:
2790+
path = "{name}:cancelAsyncQuery"
2791+
2792+
query_params = request_dict.get("_query")
2793+
if query_params:
2794+
path = f"{path}?{urlencode(query_params)}"
2795+
# TODO: remove the hack that pops config.
2796+
request_dict.pop("config", None)
2797+
2798+
http_options: Optional[types.HttpOptions] = None
2799+
if (
2800+
parameter_model.config is not None
2801+
and parameter_model.config.http_options is not None
2802+
):
2803+
http_options = parameter_model.config.http_options
2804+
2805+
request_dict = _common.convert_to_dict(request_dict)
2806+
request_dict = _common.encode_unserializable_types(request_dict)
2807+
2808+
response = await self._api_client.async_request(
2809+
"post", path, request_dict, http_options
2810+
)
2811+
2812+
response_dict = {} if not response.body else json.loads(response.body)
2813+
2814+
return_value = types.CancelQueryJobResult._from_response(
2815+
response=response_dict,
2816+
kwargs=(
2817+
{
2818+
"config": {
2819+
"response_schema": getattr(
2820+
parameter_model.config, "response_schema", None
2821+
),
2822+
"response_json_schema": getattr(
2823+
parameter_model.config, "response_json_schema", None
2824+
),
2825+
"include_all_fields": getattr(
2826+
parameter_model.config, "include_all_fields", None
2827+
),
2828+
}
2829+
}
2830+
if getattr(parameter_model, "config", None)
2831+
else {}
2832+
),
2833+
)
2834+
2835+
self._api_client._verify_response(return_value)
2836+
return return_value
2837+
26462838
async def _check_query_job(
26472839
self,
26482840
*,

‎vertexai/_genai/types/__init__.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
from .common import _AppendAgentEngineTaskEventRequestParameters
2727
from .common import _AssembleDatasetParameters
2828
from .common import _AssessDatasetParameters
29+
from .common import _CancelQueryJobAgentEngineRequestParameters
2930
from .common import _CheckQueryJobAgentEngineRequestParameters
3031
from .common import _CreateAgentEngineMemoryRequestParameters
3132
from .common import _CreateAgentEngineRequestParameters
@@ -192,6 +193,12 @@
192193
from .common import BleuResults
193194
from .common import BleuResultsDict
194195
from .common import BleuResultsOrDict
196+
from .common import CancelQueryJobAgentEngineConfig
197+
from .common import CancelQueryJobAgentEngineConfigDict
198+
from .common import CancelQueryJobAgentEngineConfigOrDict
199+
from .common import CancelQueryJobResult
200+
from .common import CancelQueryJobResultDict
201+
from .common import CancelQueryJobResultOrDict
195202
from .common import CandidateResponse
196203
from .common import CandidateResponseDict
197204
from .common import CandidateResponseOrDict
@@ -1648,6 +1655,12 @@
16481655
"VertexBaseConfig",
16491656
"VertexBaseConfigDict",
16501657
"VertexBaseConfigOrDict",
1658+
"CancelQueryJobAgentEngineConfig",
1659+
"CancelQueryJobAgentEngineConfigDict",
1660+
"CancelQueryJobAgentEngineConfigOrDict",
1661+
"CancelQueryJobResult",
1662+
"CancelQueryJobResultDict",
1663+
"CancelQueryJobResultOrDict",
16511664
"CheckQueryJobAgentEngineConfig",
16521665
"CheckQueryJobAgentEngineConfigDict",
16531666
"CheckQueryJobAgentEngineConfigOrDict",
@@ -2328,6 +2341,7 @@
23282341
"_OptimizeRequestParameters",
23292342
"_CustomJobParameters",
23302343
"_GetCustomJobParameters",
2344+
"_CancelQueryJobAgentEngineRequestParameters",
23312345
"_CheckQueryJobAgentEngineRequestParameters",
23322346
"_RunQueryJobAgentEngineRequestParameters",
23332347
"_CreateAgentEngineRequestParameters",

‎vertexai/_genai/types/common.py‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6240,6 +6240,78 @@ class _GetCustomJobParametersDict(TypedDict, total=False):
62406240
]
62416241

62426242

6243+
class CancelQueryJobAgentEngineConfig(_common.BaseModel):
6244+
"""Config for canceling async querying agent engines."""
6245+
6246+
http_options: Optional[genai_types.HttpOptions] = Field(
6247+
default=None, description="""Used to override HTTP request options."""
6248+
)
6249+
operation_name: Optional[str] = Field(
6250+
default=None,
6251+
description="""Name of the longrunning operation returned from run_query_job.""",
6252+
)
6253+
6254+
6255+
class CancelQueryJobAgentEngineConfigDict(TypedDict, total=False):
6256+
"""Config for canceling async querying agent engines."""
6257+
6258+
http_options: Optional[genai_types.HttpOptionsDict]
6259+
"""Used to override HTTP request options."""
6260+
6261+
operation_name: Optional[str]
6262+
"""Name of the longrunning operation returned from run_query_job."""
6263+
6264+
6265+
CancelQueryJobAgentEngineConfigOrDict = Union[
6266+
CancelQueryJobAgentEngineConfig, CancelQueryJobAgentEngineConfigDict
6267+
]
6268+
6269+
6270+
class _CancelQueryJobAgentEngineRequestParameters(_common.BaseModel):
6271+
"""Parameters for canceling async querying agent engines."""
6272+
6273+
name: Optional[str] = Field(
6274+
default=None, description="""Name of the reasoning engine resource."""
6275+
)
6276+
config: Optional[CancelQueryJobAgentEngineConfig] = Field(
6277+
default=None, description=""""""
6278+
)
6279+
6280+
6281+
class _CancelQueryJobAgentEngineRequestParametersDict(TypedDict, total=False):
6282+
"""Parameters for canceling async querying agent engines."""
6283+
6284+
name: Optional[str]
6285+
"""Name of the reasoning engine resource."""
6286+
6287+
config: Optional[CancelQueryJobAgentEngineConfigDict]
6288+
""""""
6289+
6290+
6291+
_CancelQueryJobAgentEngineRequestParametersOrDict = Union[
6292+
_CancelQueryJobAgentEngineRequestParameters,
6293+
_CancelQueryJobAgentEngineRequestParametersDict,
6294+
]
6295+
6296+
6297+
class CancelQueryJobResult(_common.BaseModel):
6298+
"""Result of canceling a query job."""
6299+
6300+
http_options: Optional[genai_types.HttpOptions] = Field(
6301+
default=None, description="""Used to override HTTP request options."""
6302+
)
6303+
6304+
6305+
class CancelQueryJobResultDict(TypedDict, total=False):
6306+
"""Result of canceling a query job."""
6307+
6308+
http_options: Optional[genai_types.HttpOptionsDict]
6309+
"""Used to override HTTP request options."""
6310+
6311+
6312+
CancelQueryJobResultOrDict = Union[CancelQueryJobResult, CancelQueryJobResultDict]
6313+
6314+
62436315
class CheckQueryJobAgentEngineConfig(_common.BaseModel):
62446316
"""Config for async querying agent engines."""
62456317

0 commit comments

Comments
 (0)