Skip to content

Commit f5867e7

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 505aa8c commit f5867e7

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
@@ -76,6 +76,38 @@ def _AgentEngineOperation_from_vertex(
7676
return to_object
7777

7878

79+
def _CancelQueryJobAgentEngineConfig_to_vertex(
80+
from_object: Union[dict[str, Any], object],
81+
parent_object: Optional[dict[str, Any]] = None,
82+
) -> dict[str, Any]:
83+
to_object: dict[str, Any] = {}
84+
85+
if getv(from_object, ["operation_name"]) is not None:
86+
setv(parent_object, ["operationName"], getv(from_object, ["operation_name"]))
87+
88+
return to_object
89+
90+
91+
def _CancelQueryJobAgentEngineRequestParameters_to_vertex(
92+
from_object: Union[dict[str, Any], object],
93+
parent_object: Optional[dict[str, Any]] = None,
94+
) -> dict[str, Any]:
95+
to_object: dict[str, Any] = {}
96+
if getv(from_object, ["name"]) is not None:
97+
setv(to_object, ["_url", "name"], getv(from_object, ["name"]))
98+
99+
if getv(from_object, ["config"]) is not None:
100+
setv(
101+
to_object,
102+
["config"],
103+
_CancelQueryJobAgentEngineConfig_to_vertex(
104+
getv(from_object, ["config"]), to_object
105+
),
106+
)
107+
108+
return to_object
109+
110+
79111
def _CheckQueryJobAgentEngineConfig_to_vertex(
80112
from_object: Union[dict[str, Any], object],
81113
parent_object: Optional[dict[str, Any]] = None,
@@ -720,6 +752,85 @@ def _UpdateAgentEngineRequestParameters_to_vertex(
720752

721753
class AgentEngines(_api_module.BaseModule):
722754

755+
def cancel_query_job(
756+
self,
757+
*,
758+
name: str,
759+
config: Optional[types.CancelQueryJobAgentEngineConfigOrDict] = None,
760+
) -> types.CancelQueryJobResult:
761+
"""
762+
Cancels a long-running query job on an Agent Engine.
763+
764+
Args:
765+
name (str):
766+
Required. The reasoning engine resource name.
767+
config (CancelQueryJobAgentEngineConfigOrDict):
768+
Optional. The configuration for the cancel_query_job.
769+
770+
"""
771+
772+
parameter_model = types._CancelQueryJobAgentEngineRequestParameters(
773+
name=name,
774+
config=config,
775+
)
776+
777+
request_url_dict: Optional[dict[str, str]]
778+
if not self._api_client.vertexai:
779+
raise ValueError("This method is only supported in the Vertex AI client.")
780+
else:
781+
request_dict = _CancelQueryJobAgentEngineRequestParameters_to_vertex(
782+
parameter_model
783+
)
784+
request_url_dict = request_dict.get("_url")
785+
if request_url_dict:
786+
path = "{name}:cancelAsyncQuery".format_map(request_url_dict)
787+
else:
788+
path = "{name}:cancelAsyncQuery"
789+
790+
query_params = request_dict.get("_query")
791+
if query_params:
792+
path = f"{path}?{urlencode(query_params)}"
793+
# TODO: remove the hack that pops config.
794+
request_dict.pop("config", None)
795+
796+
http_options: Optional[types.HttpOptions] = None
797+
if (
798+
parameter_model.config is not None
799+
and parameter_model.config.http_options is not None
800+
):
801+
http_options = parameter_model.config.http_options
802+
803+
request_dict = _common.convert_to_dict(request_dict)
804+
request_dict = _common.encode_unserializable_types(request_dict)
805+
806+
response = self._api_client.request("post", path, request_dict, http_options)
807+
808+
response_dict = {} if not response.body else json.loads(response.body)
809+
810+
return_value = types.CancelQueryJobResult._from_response(
811+
response=response_dict,
812+
kwargs=(
813+
{
814+
"config": {
815+
"response_schema": getattr(
816+
parameter_model.config, "response_schema", None
817+
),
818+
"response_json_schema": getattr(
819+
parameter_model.config, "response_json_schema", None
820+
),
821+
"include_all_fields": getattr(
822+
parameter_model.config, "include_all_fields", None
823+
),
824+
}
825+
}
826+
if getattr(parameter_model, "config", None)
827+
else {}
828+
),
829+
)
830+
831+
self._api_client._verify_response(return_value)
832+
return return_value
833+
723834
def _check_query_job(
724835
self,
725836
*,
@@ -2971,6 +3082,87 @@ def list_session_events(
29713082

29723083
class AsyncAgentEngines(_api_module.BaseModule):
29733084

3085+
async def cancel_query_job(
3086+
self,
3087+
*,
3088+
name: str,
3089+
config: Optional[types.CancelQueryJobAgentEngineConfigOrDict] = None,
3090+
) -> types.CancelQueryJobResult:
3091+
"""
3092+
Cancels a long-running query job on an Agent Engine.
3093+
3094+
Args:
3095+
name (str):
3096+
Required. The reasoning engine resource name.
3097+
config (CancelQueryJobAgentEngineConfigOrDict):
3098+
Optional. The configuration for the cancel_query_job.
3099+
3100+
"""
3101+
3102+
parameter_model = types._CancelQueryJobAgentEngineRequestParameters(
3103+
name=name,
3104+
config=config,
3105+
)
3106+
3107+
request_url_dict: Optional[dict[str, str]]
3108+
if not self._api_client.vertexai:
3109+
raise ValueError("This method is only supported in the Vertex AI client.")
3110+
else:
3111+
request_dict = _CancelQueryJobAgentEngineRequestParameters_to_vertex(
3112+
parameter_model
3113+
)
3114+
request_url_dict = request_dict.get("_url")
3115+
if request_url_dict:
3116+
path = "{name}:cancelAsyncQuery".format_map(request_url_dict)
3117+
else:
3118+
path = "{name}:cancelAsyncQuery"
3119+
3120+
query_params = request_dict.get("_query")
3121+
if query_params:
3122+
path = f"{path}?{urlencode(query_params)}"
3123+
# TODO: remove the hack that pops config.
3124+
request_dict.pop("config", None)
3125+
3126+
http_options: Optional[types.HttpOptions] = None
3127+
if (
3128+
parameter_model.config is not None
3129+
and parameter_model.config.http_options is not None
3130+
):
3131+
http_options = parameter_model.config.http_options
3132+
3133+
request_dict = _common.convert_to_dict(request_dict)
3134+
request_dict = _common.encode_unserializable_types(request_dict)
3135+
3136+
response = await self._api_client.async_request(
3137+
"post", path, request_dict, http_options
3138+
)
3139+
3140+
response_dict = {} if not response.body else json.loads(response.body)
3141+
3142+
return_value = types.CancelQueryJobResult._from_response(
3143+
response=response_dict,
3144+
kwargs=(
3145+
{
3146+
"config": {
3147+
"response_schema": getattr(
3148+
parameter_model.config, "response_schema", None
3149+
),
3150+
"response_json_schema": getattr(
3151+
parameter_model.config, "response_json_schema", None
3152+
),
3153+
"include_all_fields": getattr(
3154+
parameter_model.config, "include_all_fields", None
3155+
),
3156+
}
3157+
}
3158+
if getattr(parameter_model, "config", None)
3159+
else {}
3160+
),
3161+
)
3162+
3163+
self._api_client._verify_response(return_value)
3164+
return return_value
3165+
29743166
async def _check_query_job(
29753167
self,
29763168
*,

‎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
@@ -193,6 +194,12 @@
193194
from .common import BleuResults
194195
from .common import BleuResultsDict
195196
from .common import BleuResultsOrDict
197+
from .common import CancelQueryJobAgentEngineConfig
198+
from .common import CancelQueryJobAgentEngineConfigDict
199+
from .common import CancelQueryJobAgentEngineConfigOrDict
200+
from .common import CancelQueryJobResult
201+
from .common import CancelQueryJobResultDict
202+
from .common import CancelQueryJobResultOrDict
196203
from .common import CandidateResponse
197204
from .common import CandidateResponseDict
198205
from .common import CandidateResponseOrDict
@@ -1668,6 +1675,12 @@
16681675
"VertexBaseConfig",
16691676
"VertexBaseConfigDict",
16701677
"VertexBaseConfigOrDict",
1678+
"CancelQueryJobAgentEngineConfig",
1679+
"CancelQueryJobAgentEngineConfigDict",
1680+
"CancelQueryJobAgentEngineConfigOrDict",
1681+
"CancelQueryJobResult",
1682+
"CancelQueryJobResultDict",
1683+
"CancelQueryJobResultOrDict",
16711684
"CheckQueryJobAgentEngineConfig",
16721685
"CheckQueryJobAgentEngineConfigDict",
16731686
"CheckQueryJobAgentEngineConfigOrDict",
@@ -2367,6 +2380,7 @@
23672380
"_OptimizeRequestParameters",
23682381
"_CustomJobParameters",
23692382
"_GetCustomJobParameters",
2383+
"_CancelQueryJobAgentEngineRequestParameters",
23702384
"_CheckQueryJobAgentEngineRequestParameters",
23712385
"_RunQueryJobAgentEngineRequestParameters",
23722386
"_CreateAgentEngineRequestParameters",

‎vertexai/_genai/types/common.py‎

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

62806280

6281+
class CancelQueryJobAgentEngineConfig(_common.BaseModel):
6282+
"""Config for canceling async querying agent engines."""
6283+
6284+
http_options: Optional[genai_types.HttpOptions] = Field(
6285+
default=None, description="""Used to override HTTP request options."""
6286+
)
6287+
operation_name: Optional[str] = Field(
6288+
default=None,
6289+
description="""Name of the longrunning operation returned from run_query_job.""",
6290+
)
6291+
6292+
6293+
class CancelQueryJobAgentEngineConfigDict(TypedDict, total=False):
6294+
"""Config for canceling async querying agent engines."""
6295+
6296+
http_options: Optional[genai_types.HttpOptionsDict]
6297+
"""Used to override HTTP request options."""
6298+
6299+
operation_name: Optional[str]
6300+
"""Name of the longrunning operation returned from run_query_job."""
6301+
6302+
6303+
CancelQueryJobAgentEngineConfigOrDict = Union[
6304+
CancelQueryJobAgentEngineConfig, CancelQueryJobAgentEngineConfigDict
6305+
]
6306+
6307+
6308+
class _CancelQueryJobAgentEngineRequestParameters(_common.BaseModel):
6309+
"""Parameters for canceling async querying agent engines."""
6310+
6311+
name: Optional[str] = Field(
6312+
default=None, description="""Name of the reasoning engine resource."""
6313+
)
6314+
config: Optional[CancelQueryJobAgentEngineConfig] = Field(
6315+
default=None, description=""""""
6316+
)
6317+
6318+
6319+
class _CancelQueryJobAgentEngineRequestParametersDict(TypedDict, total=False):
6320+
"""Parameters for canceling async querying agent engines."""
6321+
6322+
name: Optional[str]
6323+
"""Name of the reasoning engine resource."""
6324+
6325+
config: Optional[CancelQueryJobAgentEngineConfigDict]
6326+
""""""
6327+
6328+
6329+
_CancelQueryJobAgentEngineRequestParametersOrDict = Union[
6330+
_CancelQueryJobAgentEngineRequestParameters,
6331+
_CancelQueryJobAgentEngineRequestParametersDict,
6332+
]
6333+
6334+
6335+
class CancelQueryJobResult(_common.BaseModel):
6336+
"""Result of canceling a query job."""
6337+
6338+
http_options: Optional[genai_types.HttpOptions] = Field(
6339+
default=None, description="""Used to override HTTP request options."""
6340+
)
6341+
6342+
6343+
class CancelQueryJobResultDict(TypedDict, total=False):
6344+
"""Result of canceling a query job."""
6345+
6346+
http_options: Optional[genai_types.HttpOptionsDict]
6347+
"""Used to override HTTP request options."""
6348+
6349+
6350+
CancelQueryJobResultOrDict = Union[CancelQueryJobResult, CancelQueryJobResultDict]
6351+
6352+
62816353
class CheckQueryJobAgentEngineConfig(_common.BaseModel):
62826354
"""Config for async querying agent engines."""
62836355

0 commit comments

Comments
 (0)