feat: Support Developer Connect in AE

PiperOrigin-RevId: 832480435

Author: shawn-yang-google

tests/unit/vertexai/genai/test_agent_engines.py

@@ -1062,6 +1062,47 @@ def test_create_agent_engine_config_with_source_packages(
                 == _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT
             )
 
+    def test_create_agent_engine_config_with_developer_connect_source(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            requirements_file_path = os.path.join(tmpdir, "requirements.txt")
+            with open(requirements_file_path, "w") as f:
+                f.write("requests==2.0.0")
+            developer_connect_source = {
+                "git_repository_link": "projects/test-project/locations/us-central1/connections/test-connection/gitRepositoryLinks/test-repo",
+                "revision": "main",
+                "dir": "agent",
+            }
+            config = self.client.agent_engines._create_config(
+                mode="create",
+                display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
+                description=_TEST_AGENT_ENGINE_DESCRIPTION,
+                developer_connect_source=developer_connect_source,
+                entrypoint_module="main",
+                entrypoint_object="app",
+                requirements_file=requirements_file_path,
+                class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS,
+                agent_framework=_TEST_AGENT_FRAMEWORK,
+                identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT,
+                python_version=_TEST_PYTHON_VERSION_OVERRIDE,
+            )
+            assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME
+            assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION
+            assert config["spec"]["agent_framework"] == _TEST_AGENT_FRAMEWORK
+            assert config["spec"]["source_code_spec"] == {
+                "developer_connect_source": developer_connect_source,
+                "python_spec": {
+                    "version": _TEST_PYTHON_VERSION_OVERRIDE,
+                    "entrypoint_module": "main",
+                    "entrypoint_object": "app",
+                    "requirements_file": requirements_file_path,
+                },
+            }
+            assert config["spec"]["class_methods"] == _TEST_AGENT_ENGINE_CLASS_METHODS
+            assert (
+                config["spec"]["identity_type"]
+                == _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT
+            )
+
     @mock.patch.object(
         _agent_engines_utils,
         "_create_base64_encoded_tarball",
@@ -1718,6 +1759,7 @@ def test_create_agent_engine_with_env_vars_dict(
                 labels=None,
                 class_methods=None,
                 source_packages=None,
+                developer_connect_source=None,
                 entrypoint_module=None,
                 entrypoint_object=None,
                 requirements_file=None,
@@ -1812,6 +1854,7 @@ def test_create_agent_engine_with_custom_service_account(
                 agent_server_mode=None,
                 class_methods=None,
                 source_packages=None,
+                developer_connect_source=None,
                 entrypoint_module=None,
                 entrypoint_object=None,
                 requirements_file=None,
@@ -1905,6 +1948,7 @@ def test_create_agent_engine_with_experimental_mode(
                 agent_server_mode=_genai_types.AgentServerMode.EXPERIMENTAL,
                 class_methods=None,
                 source_packages=None,
+                developer_connect_source=None,
                 entrypoint_module=None,
                 entrypoint_object=None,
                 requirements_file=None,
@@ -2061,6 +2105,7 @@ def test_create_agent_engine_with_class_methods(
                 agent_server_mode=None,
                 class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS,
                 source_packages=None,
+                developer_connect_source=None,
                 entrypoint_module=None,
                 entrypoint_object=None,
                 requirements_file=None,
@@ -2148,6 +2193,7 @@ def test_create_agent_engine_with_agent_framework(
                 agent_server_mode=None,
                 class_methods=None,
                 source_packages=None,
+                developer_connect_source=None,
                 entrypoint_module=None,
                 entrypoint_object=None,
                 requirements_file=None,

vertexai/_genai/agent_engines.py

@@ -903,6 +903,9 @@ def create(
         if context_spec is not None:
             # Conversion to a dict for _create_config
             context_spec = context_spec.model_dump()
+        developer_connect_source = config.developer_connect_source
+        if developer_connect_source is not None:
+            developer_connect_source = developer_connect_source.model_dump()
         if agent and agent_engine:
             raise ValueError("Please specify only one of `agent` or `agent_engine`.")
         elif agent_engine:
@@ -933,6 +936,7 @@ def create(
             labels=config.labels,
             class_methods=config.class_methods,
             source_packages=config.source_packages,
+            developer_connect_source=developer_connect_source,
             entrypoint_module=config.entrypoint_module,
             entrypoint_object=config.entrypoint_object,
             requirements_file=config.requirements_file,
@@ -1000,6 +1004,9 @@ def _create_config(
         agent_server_mode: Optional[types.AgentServerMode] = None,
         class_methods: Optional[Sequence[dict[str, Any]]] = None,
         source_packages: Optional[Sequence[str]] = None,
+        developer_connect_source: Optional[
+            types.ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceDict
+        ] = None,
         entrypoint_module: Optional[str] = None,
         entrypoint_object: Optional[str] = None,
         requirements_file: Optional[str] = None,
@@ -1049,6 +1056,12 @@ def _create_config(
                     "do not specify `agent` in `agent_engines.create()` or "
                     "`agent_engines.update()`."
                 )
+            if developer_connect_source is not None:
+                raise ValueError(
+                    "If you have provided `developer_connect_source` in `config`, please "
+                    "do not specify `agent` in `agent_engines.create()` or "
+                    "`agent_engines.update()`."
+                )
 
             project = self._api_client.project
             if project is None:
@@ -1149,6 +1162,11 @@ def _create_config(
                 }
             }
 
+        if developer_connect_source is not None:
+            update_masks.append("spec.source_code_spec.developer_connect_source")
+            source_code_spec = {"developer_connect_source": developer_connect_source}
+
+        if source_packages is not None or developer_connect_source is not None:
             update_masks.append("spec.source_code_spec.python_spec.version")
             python_spec = {
                 "version": sys_version,
@@ -1190,7 +1208,6 @@ def _create_config(
                     for class_method_spec in class_methods_spec
                 ],
             }
-
         if agent_engine_spec is not None:
             if (
                 env_vars is not None
@@ -1440,6 +1457,9 @@ def update(
         if context_spec is not None:
             # Conversion to a dict for _create_config
             context_spec = context_spec.model_dump()
+        developer_connect_source = config.developer_connect_source
+        if developer_connect_source is not None:
+            developer_connect_source = developer_connect_source.model_dump()
         if agent and agent_engine:
             raise ValueError("Please specify only one of `agent` or `agent_engine`.")
         elif agent_engine:
@@ -1468,6 +1488,7 @@ def update(
             labels=config.labels,
             class_methods=config.class_methods,
             source_packages=config.source_packages,
+            developer_connect_source=developer_connect_source,
             entrypoint_module=config.entrypoint_module,
             entrypoint_object=config.entrypoint_object,
             requirements_file=config.requirements_file,

vertexai/_genai/types/__init__.py

@@ -717,6 +717,12 @@
 from .common import ReasoningEngineSpecPackageSpecDict
 from .common import ReasoningEngineSpecPackageSpecOrDict
 from .common import ReasoningEngineSpecSourceCodeSpec
+from .common import ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfig
+from .common import ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict
+from .common import ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigOrDict
+from .common import ReasoningEngineSpecSourceCodeSpecDeveloperConnectSource
+from .common import ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceDict
+from .common import ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceOrDict
 from .common import ReasoningEngineSpecSourceCodeSpecDict
 from .common import ReasoningEngineSpecSourceCodeSpecInlineSource
 from .common import ReasoningEngineSpecSourceCodeSpecInlineSourceDict
@@ -1349,6 +1355,12 @@
     "ReasoningEngineSpecPackageSpec",
     "ReasoningEngineSpecPackageSpecDict",
     "ReasoningEngineSpecPackageSpecOrDict",
+    "ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfig",
+    "ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict",
+    "ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigOrDict",
+    "ReasoningEngineSpecSourceCodeSpecDeveloperConnectSource",
+    "ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceDict",
+    "ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceOrDict",
     "ReasoningEngineSpecSourceCodeSpecInlineSource",
     "ReasoningEngineSpecSourceCodeSpecInlineSourceDict",
     "ReasoningEngineSpecSourceCodeSpecInlineSourceOrDict",

vertexai/_genai/types/common.py

@@ -4703,6 +4703,68 @@ class ReasoningEngineSpecPackageSpecDict(TypedDict, total=False):
 ]
 
 
+class ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfig(_common.BaseModel):
+    """Specifies the configuration for fetching source code from a Git repository that is managed by Developer Connect. This includes the repository, revision, and directory to use."""
+
+    dir: Optional[str] = Field(
+        default=None,
+        description="""Required. Directory, relative to the source root, in which to run the build.""",
+    )
+    git_repository_link: Optional[str] = Field(
+        default=None,
+        description="""Required. The Developer Connect Git repository link, formatted as `projects/*/locations/*/connections/*/gitRepositoryLink/*`.""",
+    )
+    revision: Optional[str] = Field(
+        default=None,
+        description="""Required. The revision to fetch from the Git repository such as a branch, a tag, a commit SHA, or any Git ref.""",
+    )
+
+
+class ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict(
+    TypedDict, total=False
+):
+    """Specifies the configuration for fetching source code from a Git repository that is managed by Developer Connect. This includes the repository, revision, and directory to use."""
+
+    dir: Optional[str]
+    """Required. Directory, relative to the source root, in which to run the build."""
+
+    git_repository_link: Optional[str]
+    """Required. The Developer Connect Git repository link, formatted as `projects/*/locations/*/connections/*/gitRepositoryLink/*`."""
+
+    revision: Optional[str]
+    """Required. The revision to fetch from the Git repository such as a branch, a tag, a commit SHA, or any Git ref."""
+
+
+ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigOrDict = Union[
+    ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfig,
+    ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict,
+]
+
+
+class ReasoningEngineSpecSourceCodeSpecDeveloperConnectSource(_common.BaseModel):
+    """Specifies source code to be fetched from a Git repository managed through the Developer Connect service."""
+
+    config: Optional[ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfig] = Field(
+        default=None,
+        description="""Required. The Developer Connect configuration thats defines the specific repository, revision, and directory to use as the source code root.""",
+    )
+
+
+class ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceDict(
+    TypedDict, total=False
+):
+    """Specifies source code to be fetched from a Git repository managed through the Developer Connect service."""
+
+    config: Optional[ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict]
+    """Required. The Developer Connect configuration thats defines the specific repository, revision, and directory to use as the source code root."""
+
+
+ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceOrDict = Union[
+    ReasoningEngineSpecSourceCodeSpecDeveloperConnectSource,
+    ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceDict,
+]
+
+
 class ReasoningEngineSpecSourceCodeSpecInlineSource(_common.BaseModel):
     """Specifies source code provided as a byte stream."""
 
@@ -4771,6 +4833,12 @@ class ReasoningEngineSpecSourceCodeSpecPythonSpecDict(TypedDict, total=False):
 class ReasoningEngineSpecSourceCodeSpec(_common.BaseModel):
     """Specification for deploying from source code."""
 
+    developer_connect_source: Optional[
+        ReasoningEngineSpecSourceCodeSpecDeveloperConnectSource
+    ] = Field(
+        default=None,
+        description="""Source code is in a Git repository managed by Developer Connect.""",
+    )
     inline_source: Optional[ReasoningEngineSpecSourceCodeSpecInlineSource] = Field(
         default=None, description="""Source code is provided directly in the request."""
     )
@@ -4782,6 +4850,11 @@ class ReasoningEngineSpecSourceCodeSpec(_common.BaseModel):
 class ReasoningEngineSpecSourceCodeSpecDict(TypedDict, total=False):
     """Specification for deploying from source code."""
 
+    developer_connect_source: Optional[
+        ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceDict
+    ]
+    """Source code is in a Git repository managed by Developer Connect."""
+
     inline_source: Optional[ReasoningEngineSpecSourceCodeSpecInlineSourceDict]
     """Source code is provided directly in the request."""
 
@@ -5440,6 +5513,12 @@ class CreateAgentEngineConfig(_common.BaseModel):
         - class_methods     (required)
       """,
     )
+    developer_connect_source: Optional[
+        ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfig
+    ] = Field(
+        default=None,
+        description="""Specifies the configuration for fetching source code from a Git repository that is managed by Developer Connect. This includes the repository, revision, and directory to use.""",
+    )
     entrypoint_module: Optional[str] = Field(
         default=None,
         description="""The entrypoint module to be used for the Agent Engine
@@ -5572,6 +5651,11 @@ class CreateAgentEngineConfigDict(TypedDict, total=False):
         - class_methods     (required)
       """
 
+    developer_connect_source: Optional[
+        ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict
+    ]
+    """Specifies the configuration for fetching source code from a Git repository that is managed by Developer Connect. This includes the repository, revision, and directory to use."""
+
     entrypoint_module: Optional[str]
     """The entrypoint module to be used for the Agent Engine
       This field only used when source_packages is specified."""
@@ -6200,6 +6284,12 @@ class UpdateAgentEngineConfig(_common.BaseModel):
         - class_methods     (required)
       """,
     )
+    developer_connect_source: Optional[
+        ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfig
+    ] = Field(
+        default=None,
+        description="""Specifies the configuration for fetching source code from a Git repository that is managed by Developer Connect. This includes the repository, revision, and directory to use.""",
+    )
     entrypoint_module: Optional[str] = Field(
         default=None,
         description="""The entrypoint module to be used for the Agent Engine
@@ -6337,6 +6427,11 @@ class UpdateAgentEngineConfigDict(TypedDict, total=False):
         - class_methods     (required)
       """
 
+    developer_connect_source: Optional[
+        ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict
+    ]
+    """Specifies the configuration for fetching source code from a Git repository that is managed by Developer Connect. This includes the repository, revision, and directory to use."""
+
     entrypoint_module: Optional[str]
     """The entrypoint module to be used for the Agent Engine
       This field only used when source_packages is specified."""
@@ -13446,6 +13541,12 @@ class AgentEngineConfig(_common.BaseModel):
         - class_methods     (required)
       """,
     )
+    developer_connect_source: Optional[
+        ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfig
+    ] = Field(
+        default=None,
+        description="""Specifies the configuration for fetching source code from a Git repository that is managed by Developer Connect. This includes the repository, revision, and directory to use.""",
+    )
     entrypoint_module: Optional[str] = Field(
         default=None,
         description="""The entrypoint module to be used for the Agent Engine
@@ -13610,6 +13711,11 @@ class AgentEngineConfigDict(TypedDict, total=False):
         - class_methods     (required)
       """
 
+    developer_connect_source: Optional[
+        ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict
+    ]
+    """Specifies the configuration for fetching source code from a Git repository that is managed by Developer Connect. This includes the repository, revision, and directory to use."""
+
     entrypoint_module: Optional[str]
     """The entrypoint module to be used for the Agent Engine
       This field only used when source_packages is specified."""