Skip to content

Commit 0d97c89

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 9ed3759 commit 0d97c89

4 files changed

Lines changed: 1608 additions & 1310 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"

0 commit comments

Comments
 (0)