Skip to content

Commit 8e473c2

Browse files
yeesiancopybara-github
yeesian
authored and
copybara-github
committed
chore: GenAI SDK client - Remove dependencies on legacy SDK in client-based agent engine SDK.
PiperOrigin-RevId: 785569825
1 parent ff39b2a commit 8e473c2

3 files changed

Lines changed: 1328 additions & 80 deletions

File tree

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

Lines changed: 46 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -30,9 +30,8 @@
3030
import vertexai
3131
from google.cloud.aiplatform import initializer
3232
from vertexai._genai import agent_engines
33+
from vertexai._genai import _agent_engines_utils
3334
from vertexai._genai import types as _genai_types
34-
from vertexai.agent_engines import _agent_engines
35-
from vertexai.agent_engines import _utils
3635
from google.genai import client as genai_client
3736
from google.genai import types as genai_types
3837

@@ -316,25 +315,25 @@ def register_operations(self) -> Dict[str, List[str]]:
316315
_TEST_AGENT_ENGINE_DISPLAY_NAME = "Agent Engine Display Name"
317316
_TEST_AGENT_ENGINE_DESCRIPTION = "Agent Engine Description"
318317
_TEST_AGENT_ENGINE_LIST_FILTER = f'display_name="{_TEST_AGENT_ENGINE_DISPLAY_NAME}"'
319-
_TEST_GCS_DIR_NAME = _agent_engines._DEFAULT_GCS_DIR_NAME
320-
_TEST_BLOB_FILENAME = _agent_engines._BLOB_FILENAME
321-
_TEST_REQUIREMENTS_FILE = _agent_engines._REQUIREMENTS_FILE
322-
_TEST_EXTRA_PACKAGES_FILE = _agent_engines._EXTRA_PACKAGES_FILE
323-
_TEST_STANDARD_API_MODE = _agent_engines._STANDARD_API_MODE
324-
_TEST_ASYNC_API_MODE = _agent_engines._ASYNC_API_MODE
325-
_TEST_STREAM_API_MODE = _agent_engines._STREAM_API_MODE
326-
_TEST_ASYNC_STREAM_API_MODE = _agent_engines._ASYNC_STREAM_API_MODE
327-
_TEST_DEFAULT_METHOD_NAME = _agent_engines._DEFAULT_METHOD_NAME
328-
_TEST_DEFAULT_ASYNC_METHOD_NAME = _agent_engines._DEFAULT_ASYNC_METHOD_NAME
329-
_TEST_DEFAULT_STREAM_METHOD_NAME = _agent_engines._DEFAULT_STREAM_METHOD_NAME
318+
_TEST_GCS_DIR_NAME = _agent_engines_utils._DEFAULT_GCS_DIR_NAME
319+
_TEST_BLOB_FILENAME = _agent_engines_utils._BLOB_FILENAME
320+
_TEST_REQUIREMENTS_FILE = _agent_engines_utils._REQUIREMENTS_FILE
321+
_TEST_EXTRA_PACKAGES_FILE = _agent_engines_utils._EXTRA_PACKAGES_FILE
322+
_TEST_STANDARD_API_MODE = _agent_engines_utils._STANDARD_API_MODE
323+
_TEST_ASYNC_API_MODE = _agent_engines_utils._ASYNC_API_MODE
324+
_TEST_STREAM_API_MODE = _agent_engines_utils._STREAM_API_MODE
325+
_TEST_ASYNC_STREAM_API_MODE = _agent_engines_utils._ASYNC_STREAM_API_MODE
326+
_TEST_DEFAULT_METHOD_NAME = _agent_engines_utils._DEFAULT_METHOD_NAME
327+
_TEST_DEFAULT_ASYNC_METHOD_NAME = _agent_engines_utils._DEFAULT_ASYNC_METHOD_NAME
328+
_TEST_DEFAULT_STREAM_METHOD_NAME = _agent_engines_utils._DEFAULT_STREAM_METHOD_NAME
330329
_TEST_DEFAULT_ASYNC_STREAM_METHOD_NAME = (
331-
_agent_engines._DEFAULT_ASYNC_STREAM_METHOD_NAME
330+
_agent_engines_utils._DEFAULT_ASYNC_STREAM_METHOD_NAME
332331
)
333332
_TEST_CAPITALIZE_ENGINE_METHOD_DOCSTRING = "Runs the engine."
334333
_TEST_STREAM_METHOD_DOCSTRING = "Runs the stream engine."
335334
_TEST_ASYNC_STREAM_METHOD_DOCSTRING = "Runs the async stream engine."
336-
_TEST_MODE_KEY_IN_SCHEMA = _agent_engines._MODE_KEY_IN_SCHEMA
337-
_TEST_METHOD_NAME_KEY_IN_SCHEMA = _agent_engines._METHOD_NAME_KEY_IN_SCHEMA
335+
_TEST_MODE_KEY_IN_SCHEMA = _agent_engines_utils._MODE_KEY_IN_SCHEMA
336+
_TEST_METHOD_NAME_KEY_IN_SCHEMA = _agent_engines_utils._METHOD_NAME_KEY_IN_SCHEMA
338337
_TEST_CUSTOM_METHOD_NAME = "custom_method"
339338
_TEST_CUSTOM_ASYNC_METHOD_NAME = "custom_async_method"
340339
_TEST_CUSTOM_STREAM_METHOD_NAME = "custom_stream_method"
@@ -422,13 +421,13 @@ def register_operations(self) -> Dict[str, List[str]]:
422421
"lib",
423422
"main.py",
424423
]
425-
_TEST_AGENT_ENGINE_QUERY_SCHEMA = _utils.generate_schema(
424+
_TEST_AGENT_ENGINE_QUERY_SCHEMA = _agent_engines_utils._generate_schema(
426425
CapitalizeEngine().query,
427426
schema_name=_TEST_DEFAULT_METHOD_NAME,
428427
)
429428
_TEST_AGENT_ENGINE_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = _TEST_STANDARD_API_MODE
430429
_TEST_PYTHON_VERSION = f"{sys.version_info.major}.{sys.version_info.minor}"
431-
_TEST_AGENT_ENGINE_FRAMEWORK = _agent_engines._DEFAULT_AGENT_FRAMEWORK
430+
_TEST_AGENT_ENGINE_FRAMEWORK = _agent_engines_utils._DEFAULT_AGENT_FRAMEWORK
432431
_TEST_AGENT_ENGINE_CLASS_METHOD_1 = {
433432
"description": "Runs the engine.",
434433
"name": "query",
@@ -513,47 +512,49 @@ def register_operations(self) -> Dict[str, List[str]]:
513512
_TEST_AGENT_ENGINE_STREAM_QUERY_RESPONSE = [{"output": "hello"}, {"output": "world"}]
514513
_TEST_AGENT_ENGINE_OPERATION_SCHEMAS = []
515514
_TEST_AGENT_ENGINE_EXTRA_PACKAGE = "fake.py"
516-
_TEST_AGENT_ENGINE_ASYNC_METHOD_SCHEMA = _utils.generate_schema(
515+
_TEST_AGENT_ENGINE_ASYNC_METHOD_SCHEMA = _agent_engines_utils._generate_schema(
517516
AsyncQueryEngine().async_query,
518517
schema_name=_TEST_DEFAULT_ASYNC_METHOD_NAME,
519518
)
520519
_TEST_AGENT_ENGINE_ASYNC_METHOD_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = _TEST_ASYNC_API_MODE
521-
_TEST_AGENT_ENGINE_CUSTOM_METHOD_SCHEMA = _utils.generate_schema(
520+
_TEST_AGENT_ENGINE_CUSTOM_METHOD_SCHEMA = _agent_engines_utils._generate_schema(
522521
OperationRegistrableEngine().custom_method,
523522
schema_name=_TEST_CUSTOM_METHOD_NAME,
524523
)
525524
_TEST_AGENT_ENGINE_CUSTOM_METHOD_SCHEMA[
526525
_TEST_MODE_KEY_IN_SCHEMA
527526
] = _TEST_STANDARD_API_MODE
528-
_TEST_AGENT_ENGINE_ASYNC_CUSTOM_METHOD_SCHEMA = _utils.generate_schema(
527+
_TEST_AGENT_ENGINE_ASYNC_CUSTOM_METHOD_SCHEMA = _agent_engines_utils._generate_schema(
529528
OperationRegistrableEngine().custom_async_method,
530529
schema_name=_TEST_CUSTOM_ASYNC_METHOD_NAME,
531530
)
532531
_TEST_AGENT_ENGINE_ASYNC_CUSTOM_METHOD_SCHEMA[
533532
_TEST_MODE_KEY_IN_SCHEMA
534533
] = _TEST_ASYNC_API_MODE
535-
_TEST_AGENT_ENGINE_STREAM_QUERY_SCHEMA = _utils.generate_schema(
534+
_TEST_AGENT_ENGINE_STREAM_QUERY_SCHEMA = _agent_engines_utils._generate_schema(
536535
StreamQueryEngine().stream_query,
537536
schema_name=_TEST_DEFAULT_STREAM_METHOD_NAME,
538537
)
539538
_TEST_AGENT_ENGINE_STREAM_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = _TEST_STREAM_API_MODE
540-
_TEST_AGENT_ENGINE_CUSTOM_STREAM_QUERY_SCHEMA = _utils.generate_schema(
539+
_TEST_AGENT_ENGINE_CUSTOM_STREAM_QUERY_SCHEMA = _agent_engines_utils._generate_schema(
541540
OperationRegistrableEngine().custom_stream_method,
542541
schema_name=_TEST_CUSTOM_STREAM_METHOD_NAME,
543542
)
544543
_TEST_AGENT_ENGINE_CUSTOM_STREAM_QUERY_SCHEMA[
545544
_TEST_MODE_KEY_IN_SCHEMA
546545
] = _TEST_STREAM_API_MODE
547-
_TEST_AGENT_ENGINE_ASYNC_STREAM_QUERY_SCHEMA = _utils.generate_schema(
546+
_TEST_AGENT_ENGINE_ASYNC_STREAM_QUERY_SCHEMA = _agent_engines_utils._generate_schema(
548547
AsyncStreamQueryEngine().async_stream_query,
549548
schema_name=_TEST_DEFAULT_ASYNC_STREAM_METHOD_NAME,
550549
)
551550
_TEST_AGENT_ENGINE_ASYNC_STREAM_QUERY_SCHEMA[
552551
_TEST_MODE_KEY_IN_SCHEMA
553552
] = _TEST_ASYNC_STREAM_API_MODE
554-
_TEST_AGENT_ENGINE_CUSTOM_ASYNC_STREAM_QUERY_SCHEMA = _utils.generate_schema(
555-
OperationRegistrableEngine().custom_async_stream_method,
556-
schema_name=_TEST_CUSTOM_ASYNC_STREAM_METHOD_NAME,
553+
_TEST_AGENT_ENGINE_CUSTOM_ASYNC_STREAM_QUERY_SCHEMA = (
554+
_agent_engines_utils._generate_schema(
555+
OperationRegistrableEngine().custom_async_stream_method,
556+
schema_name=_TEST_CUSTOM_ASYNC_STREAM_METHOD_NAME,
557+
)
557558
)
558559
_TEST_AGENT_ENGINE_CUSTOM_ASYNC_STREAM_QUERY_SCHEMA[
559560
_TEST_MODE_KEY_IN_SCHEMA
@@ -578,7 +579,7 @@ def register_operations(self) -> Dict[str, List[str]]:
578579
_TEST_NO_OPERATION_REGISTRABLE_SCHEMAS = [
579580
_TEST_AGENT_ENGINE_QUERY_SCHEMA,
580581
]
581-
_TEST_METHOD_TO_BE_UNREGISTERED_SCHEMA = _utils.generate_schema(
582+
_TEST_METHOD_TO_BE_UNREGISTERED_SCHEMA = _agent_engines_utils._generate_schema(
582583
MethodToBeUnregisteredEngine().method_to_be_unregistered,
583584
schema_name=_TEST_METHOD_TO_BE_UNREGISTERED_NAME,
584585
)
@@ -747,7 +748,7 @@ def setup_method(self):
747748
def teardown_method(self):
748749
initializer.global_pool.shutdown(wait=True)
749750

750-
@mock.patch.object(_agent_engines, "_prepare")
751+
@mock.patch.object(_agent_engines_utils, "_prepare")
751752
def test_create_agent_engine_config_lightweight(self, mock_prepare):
752753
config = self.client.agent_engines._create_config(
753754
mode="create",
@@ -760,7 +761,7 @@ def test_create_agent_engine_config_lightweight(self, mock_prepare):
760761
"description": _TEST_AGENT_ENGINE_DESCRIPTION,
761762
}
762763

763-
@mock.patch.object(_agent_engines, "_prepare")
764+
@mock.patch.object(_agent_engines_utils, "_prepare")
764765
def test_create_agent_engine_config_full(self, mock_prepare):
765766
config = self.client.agent_engines._create_config(
766767
mode="create",
@@ -799,7 +800,7 @@ def test_create_agent_engine_config_full(self, mock_prepare):
799800
}
800801
assert config["spec"]["class_methods"] == [_TEST_AGENT_ENGINE_CLASS_METHOD_1]
801802

802-
@mock.patch.object(_agent_engines, "_prepare")
803+
@mock.patch.object(_agent_engines_utils, "_prepare")
803804
def test_update_agent_engine_config_full(self, mock_prepare):
804805
config = self.client.agent_engines._create_config(
805806
mode="update",
@@ -965,7 +966,7 @@ def test_list_agent_engine(self):
965966
)
966967

967968
@pytest.mark.usefixtures("caplog")
968-
@mock.patch.object(_agent_engines, "_prepare")
969+
@mock.patch.object(_agent_engines_utils, "_prepare")
969970
@mock.patch.object(agent_engines.AgentEngines, "_await_operation")
970971
def test_create_agent_engine(self, mock_await_operation, mock_prepare, caplog):
971972
mock_await_operation.return_value = _genai_types.AgentEngineOperation(
@@ -1132,7 +1133,7 @@ def test_create_agent_engine_with_env_vars_dict(
11321133
)
11331134

11341135
@pytest.mark.usefixtures("caplog")
1135-
@mock.patch.object(_agent_engines, "_prepare")
1136+
@mock.patch.object(_agent_engines_utils, "_prepare")
11361137
@mock.patch.object(agent_engines.AgentEngines, "_await_operation")
11371138
def test_update_agent_engine_requirements(
11381139
self, mock_await_operation, mock_prepare, caplog
@@ -1189,7 +1190,7 @@ def test_update_agent_engine_requirements(
11891190
f"'{_TEST_AGENT_ENGINE_RESOURCE_NAME}')" in caplog.text
11901191
)
11911192

1192-
@mock.patch.object(_agent_engines, "_prepare")
1193+
@mock.patch.object(_agent_engines_utils, "_prepare")
11931194
@mock.patch.object(agent_engines.AgentEngines, "_await_operation")
11941195
def test_update_agent_engine_extra_packages(
11951196
self, mock_await_operation, mock_prepare
@@ -1243,7 +1244,7 @@ def test_update_agent_engine_extra_packages(
12431244
None,
12441245
)
12451246

1246-
@mock.patch.object(_agent_engines, "_prepare")
1247+
@mock.patch.object(_agent_engines_utils, "_prepare")
12471248
@mock.patch.object(agent_engines.AgentEngines, "_await_operation")
12481249
def test_update_agent_engine_env_vars(
12491250
self, mock_await_operation, mock_prepare, caplog
@@ -1522,7 +1523,7 @@ async def consume():
15221523
_TEST_NO_OPERATION_REGISTRABLE_SCHEMAS,
15231524
[
15241525
(
1525-
_utils.generate_schema(
1526+
_agent_engines_utils._generate_schema(
15261527
CapitalizeEngine().query,
15271528
schema_name=_TEST_DEFAULT_METHOD_NAME,
15281529
),
@@ -1535,56 +1536,56 @@ async def consume():
15351536
_TEST_OPERATION_REGISTRABLE_SCHEMAS,
15361537
[
15371538
(
1538-
_utils.generate_schema(
1539+
_agent_engines_utils._generate_schema(
15391540
OperationRegistrableEngine().query,
15401541
schema_name=_TEST_DEFAULT_METHOD_NAME,
15411542
),
15421543
_TEST_STANDARD_API_MODE,
15431544
),
15441545
(
1545-
_utils.generate_schema(
1546+
_agent_engines_utils._generate_schema(
15461547
OperationRegistrableEngine().custom_method,
15471548
schema_name=_TEST_CUSTOM_METHOD_NAME,
15481549
),
15491550
_TEST_STANDARD_API_MODE,
15501551
),
15511552
(
1552-
_utils.generate_schema(
1553+
_agent_engines_utils._generate_schema(
15531554
OperationRegistrableEngine().async_query,
15541555
schema_name=_TEST_DEFAULT_ASYNC_METHOD_NAME,
15551556
),
15561557
_TEST_ASYNC_API_MODE,
15571558
),
15581559
(
1559-
_utils.generate_schema(
1560+
_agent_engines_utils._generate_schema(
15601561
OperationRegistrableEngine().custom_async_method,
15611562
schema_name=_TEST_CUSTOM_ASYNC_METHOD_NAME,
15621563
),
15631564
_TEST_ASYNC_API_MODE,
15641565
),
15651566
(
1566-
_utils.generate_schema(
1567+
_agent_engines_utils._generate_schema(
15671568
OperationRegistrableEngine().stream_query,
15681569
schema_name=_TEST_DEFAULT_STREAM_METHOD_NAME,
15691570
),
15701571
_TEST_STREAM_API_MODE,
15711572
),
15721573
(
1573-
_utils.generate_schema(
1574+
_agent_engines_utils._generate_schema(
15741575
OperationRegistrableEngine().custom_stream_method,
15751576
schema_name=_TEST_CUSTOM_STREAM_METHOD_NAME,
15761577
),
15771578
_TEST_STREAM_API_MODE,
15781579
),
15791580
(
1580-
_utils.generate_schema(
1581+
_agent_engines_utils._generate_schema(
15811582
OperationRegistrableEngine().async_stream_query,
15821583
schema_name=_TEST_DEFAULT_ASYNC_STREAM_METHOD_NAME,
15831584
),
15841585
_TEST_ASYNC_STREAM_API_MODE,
15851586
),
15861587
(
1587-
_utils.generate_schema(
1588+
_agent_engines_utils._generate_schema(
15881589
OperationRegistrableEngine().custom_async_stream_method,
15891590
schema_name=_TEST_CUSTOM_ASYNC_STREAM_METHOD_NAME,
15901591
),
@@ -1597,7 +1598,7 @@ async def consume():
15971598
_TEST_OPERATION_NOT_REGISTERED_SCHEMAS,
15981599
[
15991600
(
1600-
_utils.generate_schema(
1601+
_agent_engines_utils._generate_schema(
16011602
OperationNotRegisteredEngine().custom_method,
16021603
schema_name=_TEST_CUSTOM_METHOD_NAME,
16031604
),

0 commit comments

Comments
 (0)