Skip to content

Commit 5c2005d

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
feat: Adds cancel_query_job to SDK for agent engine long running async tasks.
PiperOrigin-RevId: 896911226
1 parent 4c28b34 commit 5c2005d

File tree

4 files changed

+1179
-943
lines changed

4 files changed

+1179
-943
lines changed

tests/unit/vertexai/genai/test_agent_engines.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3162,6 +3162,28 @@ def test_query_agent_engine_async(self):
31623162
None,
31633163
)
31643164

3165+
def test_cancel_query_job_agent_engine(self):
3166+
with mock.patch.object(
3167+
self.client.agent_engines._api_client, "request"
3168+
) as request_mock:
3169+
request_mock.return_value = genai_types.HttpResponse(body="{}")
3170+
3171+
result = self.client.agent_engines.cancel_query_job(
3172+
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
3173+
operation_name=_TEST_AGENT_ENGINE_OPERATION_NAME,
3174+
)
3175+
3176+
assert isinstance(result, _genai_types.CancelQueryJobResult)
3177+
request_mock.assert_called_with(
3178+
"post",
3179+
f"{_TEST_AGENT_ENGINE_RESOURCE_NAME}:cancelAsyncQuery",
3180+
{
3181+
"_url": {"name": _TEST_AGENT_ENGINE_RESOURCE_NAME},
3182+
"operationName": _TEST_AGENT_ENGINE_OPERATION_NAME,
3183+
},
3184+
None,
3185+
)
3186+
31653187
def test_check_query_job_agent_engine(self):
31663188
with mock.patch.object(
31673189
self.client.agent_engines._api_client, "request"

0 commit comments

Comments
 (0)