chore: Support specifying agent_framework in Agent Engine creation and update.

PiperOrigin-RevId: 825734997

Author: shawn-yang-google

tests/unit/vertexai/genai/test_agent_engines.py

@@ -40,7 +40,8 @@
 import pytest
 
 
-_TEST_AGENT_FRAMEWORK = "test-agent-framework"
+_TEST_AGENT_FRAMEWORK = _genai_types.ReasoningEngineAgentFramework.GOOGLE_ADK
+_TEST_AGENT_FRAMEWORK_STR = "google-adk"
 GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY = (
     "GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY"
 )
@@ -976,9 +977,11 @@ def test_create_agent_engine_config_with_source_packages(
                 entrypoint_object="app",
                 requirements_file=requirements_file_path,
                 class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS,
+                agent_framework=_TEST_AGENT_FRAMEWORK,
             )
             assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME
             assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION
+            assert config["spec"]["agent_framework"] == _TEST_AGENT_FRAMEWORK_STR
             assert config["spec"]["source_code_spec"] == {
                 "inline_source": {"source_archive": "test_tarball"},
                 "python_spec": {
@@ -1500,6 +1503,7 @@ def test_create_agent_engine_with_env_vars_dict(
                 entrypoint_module=None,
                 entrypoint_object=None,
                 requirements_file=None,
+                agent_framework=None,
             )
             request_mock.assert_called_with(
                 "post",
@@ -1586,6 +1590,7 @@ def test_create_agent_engine_with_custom_service_account(
                 entrypoint_module=None,
                 entrypoint_object=None,
                 requirements_file=None,
+                agent_framework=None,
             )
             request_mock.assert_called_with(
                 "post",
@@ -1674,6 +1679,7 @@ def test_create_agent_engine_with_experimental_mode(
                 entrypoint_module=None,
                 entrypoint_object=None,
                 requirements_file=None,
+                agent_framework=None,
             )
             request_mock.assert_called_with(
                 "post",
@@ -1826,6 +1832,7 @@ def test_create_agent_engine_with_class_methods(
                 entrypoint_module=None,
                 entrypoint_object=None,
                 requirements_file=None,
+                agent_framework=None,
             )
             request_mock.assert_called_with(
                 "post",
@@ -1845,6 +1852,92 @@ def test_create_agent_engine_with_class_methods(
                 None,
             )
 
+    @mock.patch.object(agent_engines.AgentEngines, "_create_config")
+    @mock.patch.object(_agent_engines_utils, "_await_operation")
+    def test_create_agent_engine_with_agent_framework(
+        self,
+        mock_await_operation,
+        mock_create_config,
+    ):
+        mock_create_config.return_value = {
+            "display_name": _TEST_AGENT_ENGINE_DISPLAY_NAME,
+            "description": _TEST_AGENT_ENGINE_DESCRIPTION,
+            "spec": {
+                "package_spec": {
+                    "python_version": _TEST_PYTHON_VERSION,
+                    "pickle_object_gcs_uri": _TEST_AGENT_ENGINE_GCS_URI,
+                    "requirements_gcs_uri": _TEST_AGENT_ENGINE_REQUIREMENTS_GCS_URI,
+                },
+                "class_methods": [_TEST_AGENT_ENGINE_CLASS_METHOD_1],
+                "agent_framework": _TEST_AGENT_FRAMEWORK_STR,
+            },
+        }
+        mock_await_operation.return_value = _genai_types.AgentEngineOperation(
+            response=_genai_types.ReasoningEngine(
+                name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
+                spec=_TEST_AGENT_ENGINE_SPEC,
+            )
+        )
+        with mock.patch.object(
+            self.client.agent_engines._api_client, "request"
+        ) as request_mock:
+            request_mock.return_value = genai_types.HttpResponse(body="")
+            self.client.agent_engines.create(
+                agent=self.test_agent,
+                config=_genai_types.AgentEngineConfig(
+                    display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
+                    requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
+                    extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH],
+                    staging_bucket=_TEST_STAGING_BUCKET,
+                    agent_framework=_TEST_AGENT_FRAMEWORK,
+                ),
+            )
+            mock_create_config.assert_called_with(
+                mode="create",
+                agent=self.test_agent,
+                staging_bucket=_TEST_STAGING_BUCKET,
+                requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
+                display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
+                description=None,
+                gcs_dir_name=None,
+                extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH],
+                env_vars=None,
+                service_account=None,
+                context_spec=None,
+                psc_interface_config=None,
+                min_instances=None,
+                max_instances=None,
+                resource_limits=None,
+                container_concurrency=None,
+                encryption_spec=None,
+                labels=None,
+                agent_server_mode=None,
+                class_methods=None,
+                source_packages=None,
+                entrypoint_module=None,
+                entrypoint_object=None,
+                requirements_file=None,
+                agent_framework=_TEST_AGENT_FRAMEWORK,
+            )
+            request_mock.assert_called_with(
+                "post",
+                "reasoningEngines",
+                {
+                    "displayName": _TEST_AGENT_ENGINE_DISPLAY_NAME,
+                    "description": _TEST_AGENT_ENGINE_DESCRIPTION,
+                    "spec": {
+                        "agent_framework": _TEST_AGENT_FRAMEWORK_STR,
+                        "class_methods": [_TEST_AGENT_ENGINE_CLASS_METHOD_1],
+                        "package_spec": {
+                            "pickle_object_gcs_uri": _TEST_AGENT_ENGINE_GCS_URI,
+                            "python_version": _TEST_PYTHON_VERSION,
+                            "requirements_gcs_uri": _TEST_AGENT_ENGINE_REQUIREMENTS_GCS_URI,
+                        },
+                    },
+                },
+                None,
+            )
+
     @pytest.mark.usefixtures("caplog")
     @mock.patch.object(_agent_engines_utils, "_prepare")
     @mock.patch.object(_agent_engines_utils, "_await_operation")

vertexai/_genai/_agent_engines_utils.py

@@ -128,6 +128,26 @@
 _BASE_MODULES = set(_BUILTIN_MODULE_NAMES + tuple(_STDLIB_MODULE_NAMES))
 _BLOB_FILENAME = "agent_engine.pkl"
 _DEFAULT_AGENT_FRAMEWORK = "custom"
+_SUPPORTED_AGENT_FRAMEWORKS = frozenset(
+    [
+        "google-adk",
+        "langchain",
+        "langgraph",
+        "ag2",
+        "llama-index",
+        "custom",
+    ]
+)
+_AGENT_FRAMEWORK_TO_STR = types.MappingProxyType(
+    {
+        genai_types.ReasoningEngineAgentFramework.GOOGLE_ADK: "google-adk",
+        genai_types.ReasoningEngineAgentFramework.LANGCHAIN: "langchain",
+        genai_types.ReasoningEngineAgentFramework.LANGGRAPH: "langgraph",
+        genai_types.ReasoningEngineAgentFramework.AG2: "ag2",
+        genai_types.ReasoningEngineAgentFramework.LLAMA_INDEX: "llama-index",
+        genai_types.ReasoningEngineAgentFramework.CUSTOM: "custom",
+    }
+)
 _DEFAULT_ASYNC_METHOD_NAME = "async_query"
 _DEFAULT_ASYNC_METHOD_RETURN_TYPE = "Coroutine[Any]"
 _DEFAULT_ASYNC_STREAM_METHOD_NAME = "async_stream_query"
@@ -705,13 +725,42 @@ def _generate_schema(
     return schema
 
 
-def _get_agent_framework(*, agent: _AgentEngineInterface) -> str:
-    if (
-        hasattr(agent, _AGENT_FRAMEWORK_ATTR)
-        and getattr(agent, _AGENT_FRAMEWORK_ATTR) is not None
-        and isinstance(getattr(agent, _AGENT_FRAMEWORK_ATTR), str)
-    ):
-        return getattr(agent, _AGENT_FRAMEWORK_ATTR)
+def _get_agent_framework_str(
+    *,
+    agent_framework: genai_types.ReasoningEngineAgentFramework,
+    agent: _AgentEngineInterface,
+) -> str:
+    """Gets the agent framework to use.
+
+    It prioritizes the provided `agent_framework`. If not provided or not
+    supported, it checks the `_AGENT_FRAMEWORK_ATTR` attribute on the agent.
+    If neither is found, it defaults to "_DEFAULT_AGENT_FRAMEWORK".
+
+    Args:
+        agent_framework (genai_types.ReasoningEngineAgentFramework):
+            The agent framework provided by the user.
+        agent (_AgentEngineInterface):
+            The agent engine instance.
+
+    Returns:
+        str: The name of the agent framework to use.
+    """
+    if agent_framework is not None and agent_framework in _AGENT_FRAMEWORK_TO_STR:
+        logger.info(f"Using agent framework: {agent_framework}")
+        return _AGENT_FRAMEWORK_TO_STR[agent_framework]
+    if hasattr(agent, _AGENT_FRAMEWORK_ATTR):
+        agent_framework_attr = getattr(agent, _AGENT_FRAMEWORK_ATTR)
+        if (
+            agent_framework_attr is not None
+            and isinstance(agent_framework_attr, str)
+            and agent_framework_attr in _SUPPORTED_AGENT_FRAMEWORKS
+        ):
+            logger.info(f"Using agent framework: {agent_framework_attr}")
+            return agent_framework_attr
+    logger.info(
+        f"The provided agent framework {agent_framework} is not supported."
+        f" Defaulting to {_DEFAULT_AGENT_FRAMEWORK}."
+    )
     return _DEFAULT_AGENT_FRAMEWORK
 
 

vertexai/_genai/agent_engines.py

@@ -94,6 +94,9 @@ def _CreateAgentEngineConfig_to_vertex(
             getv(from_object, ["requirements_file"]),
         )
 
+    if getv(from_object, ["agent_framework"]) is not None:
+        setv(parent_object, ["agentFramework"], getv(from_object, ["agent_framework"]))
+
     return to_object
 
 
@@ -285,6 +288,9 @@ def _UpdateAgentEngineConfig_to_vertex(
             getv(from_object, ["requirements_file"]),
         )
 
+    if getv(from_object, ["agent_framework"]) is not None:
+        setv(parent_object, ["agentFramework"], getv(from_object, ["agent_framework"]))
+
     if getv(from_object, ["update_mask"]) is not None:
         setv(
             parent_object, ["_query", "updateMask"], getv(from_object, ["update_mask"])
@@ -923,6 +929,7 @@ def create(
             entrypoint_module=config.entrypoint_module,
             entrypoint_object=config.entrypoint_object,
             requirements_file=config.requirements_file,
+            agent_framework=config.agent_framework,
         )
         operation = self._create(config=api_config)
         # TODO: Use a more specific link.
@@ -986,6 +993,7 @@ def _create_config(
         entrypoint_module: Optional[str] = None,
         entrypoint_object: Optional[str] = None,
         requirements_file: Optional[str] = None,
+        agent_framework: Optional[types.ReasoningEngineAgentFramework] = None,
     ) -> types.UpdateAgentEngineConfigDict:
         import sys
 
@@ -1195,7 +1203,10 @@ def _create_config(
                 ] = agent_server_mode
 
             agent_engine_spec["agent_framework"] = (
-                _agent_engines_utils._get_agent_framework(agent=agent)
+                _agent_engines_utils._get_agent_framework_str(
+                    agent_framework=agent_framework,
+                    agent=agent,
+                )
             )
             update_masks.append("spec.agent_framework")
             config["spec"] = agent_engine_spec
@@ -1423,6 +1434,7 @@ def update(
             entrypoint_module=config.entrypoint_module,
             entrypoint_object=config.entrypoint_object,
             requirements_file=config.requirements_file,
+            agent_framework=config.agent_framework,
         )
         operation = self._update(name=name, config=api_config)
         logger.info(

vertexai/_genai/types/__init__.py

@@ -648,6 +648,7 @@
 from .common import RawOutputDict
 from .common import RawOutputOrDict
 from .common import ReasoningEngine
+from .common import ReasoningEngineAgentFramework
 from .common import ReasoningEngineContextSpec
 from .common import ReasoningEngineContextSpecDict
 from .common import ReasoningEngineContextSpecMemoryBankConfig
@@ -1802,6 +1803,7 @@
     "RubricContentType",
     "EvaluationRunState",
     "OptimizeTarget",
+    "ReasoningEngineAgentFramework",
     "GenerateMemoriesResponseGeneratedMemoryAction",
     "PromptOptimizerMethod",
     "PromptData",