From affc0885ccdb29bb37ea2365c14e12f7da698437 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Tue, 10 Nov 2026 08:22:25 -0800 Subject: [PATCH] feat: support `git_source` in `client.runtimes` for deploying from public Git repositories PiperOrigin-RevId: 1018241384 --- agentplatform/_genai/runtimes.py | 50 ++++-- agentplatform/_genai/types/common.py | 21 +++ .../agentplatform/genai/test_agent_engines.py | 160 ++++++++++++++++++ 3 files changed, 221 insertions(+), 10 deletions(-) diff --git a/agentplatform/_genai/runtimes.py b/agentplatform/_genai/runtimes.py index 479badf162..24f5333d60 100644 --- a/agentplatform/_genai/runtimes.py +++ b/agentplatform/_genai/runtimes.py @@ -1607,6 +1607,7 @@ def _is_lightweight_creation( agent or config.source_packages or config.developer_connect_source + or config.git_source or config.agent_config_source or config.container_spec ): @@ -1860,6 +1861,9 @@ def create( developer_connect_source = json.loads( developer_connect_source.model_dump_json() ) + git_source = config.git_source + if git_source is not None: + git_source = json.loads(git_source.model_dump_json(exclude_none=True)) agent_config_source = config.agent_config_source if agent_config_source is not None: agent_config_source = json.loads(agent_config_source.model_dump_json()) @@ -1900,6 +1904,7 @@ def create( class_methods=config.class_methods, source_packages=config.source_packages, developer_connect_source=developer_connect_source, + git_source=git_source, entrypoint_module=config.entrypoint_module, entrypoint_object=config.entrypoint_object, requirements_file=config.requirements_file, @@ -1969,6 +1974,9 @@ def _set_source_code_spec( developer_connect_source: Optional[ types.ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict ] = None, + git_source: Optional[ + types.ReasoningEngineSpecSourceCodeSpecGitSourceDict + ] = None, class_methods: Optional[Sequence[dict[str, Any]]] = None, entrypoint_module: Optional[str] = None, entrypoint_object: Optional[str] = None, @@ -2000,10 +2008,13 @@ def _set_source_code_spec( source_code_spec["developer_connect_source"] = { "config": developer_connect_source } + elif git_source: + update_masks.append("spec.source_code_spec.git_source") + source_code_spec["git_source"] = git_source elif not agent_config_source: raise ValueError( "Please specify one of `source_packages`, `developer_connect_source`, " - "or `agent_config_source`." + "`git_source`, or `agent_config_source`." ) if class_methods is not None: update_masks.append("spec.class_methods") @@ -2019,8 +2030,8 @@ def _set_source_code_spec( elif image_spec is None: raise ValueError( "`class_methods` must be specified if `source_packages`, " - "`developer_connect_source`, or `agent_config_source` is " - "specified without a Dockerfile or `image_spec`." + "`developer_connect_source`, `git_source`, or `agent_config_source` " + "is specified without a Dockerfile or `image_spec`." ) if image_spec is not None: if entrypoint_module or entrypoint_object or requirements_file: @@ -2077,13 +2088,13 @@ def _set_source_code_spec( if not entrypoint_module: raise ValueError( - "`entrypoint_module` must be specified if `source_packages` or `developer_connect_source` is specified." + "`entrypoint_module` must be specified if `source_packages`, `developer_connect_source`, or `git_source` is specified." ) update_masks.append("spec.source_code_spec.python_spec.entrypoint_module") python_spec["entrypoint_module"] = entrypoint_module if not entrypoint_object: raise ValueError( - "`entrypoint_object` must be specified if `source_packages` or `developer_connect_source` is specified." + "`entrypoint_object` must be specified if `source_packages`, `developer_connect_source`, or `git_source` is specified." ) update_masks.append("spec.source_code_spec.python_spec.entrypoint_object") python_spec["entrypoint_object"] = entrypoint_object @@ -2215,6 +2226,9 @@ def _create_config( developer_connect_source: Optional[ types.ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict ] = None, + git_source: Optional[ + types.ReasoningEngineSpecSourceCodeSpecGitSourceDict + ] = None, entrypoint_module: Optional[str] = None, entrypoint_object: Optional[str] = None, requirements_file: Optional[str] = None, @@ -2284,9 +2298,18 @@ def _create_config( "do not specify `agent` in `runtimes.create()` or " "`runtimes.update()`." ) - elif source_packages and developer_connect_source: + if git_source: + raise ValueError( + "If you have provided `git_source` in `config`, please " + "do not specify `agent` in `runtimes.create()` or " + "`runtimes.update()`." + ) + elif ( + sum(map(bool, (source_packages, developer_connect_source, git_source))) > 1 + ): raise ValueError( - "Please specify only one of `source_packages` or `developer_connect_source` in `config`." + "Please specify only one of `source_packages`, " + "`developer_connect_source`, or `git_source` in `config`." ) if container_spec: @@ -2296,10 +2319,11 @@ def _create_config( "do not specify `agent` in `runtimes.create()` or " "`runtimes.update()`." ) - if source_packages or developer_connect_source: + if source_packages or developer_connect_source or git_source: raise ValueError( "If you have provided `container_spec` in `config`, please " - "do not specify `source_packages` or `developer_connect_source` in `config`." + "do not specify `source_packages`, `developer_connect_source`, or " + "`git_source` in `config`." ) runtime_spec: Any = None @@ -2323,6 +2347,7 @@ def _create_config( elif ( source_packages or developer_connect_source + or git_source or image_spec or agent_config_source ): @@ -2332,6 +2357,7 @@ def _create_config( update_masks=update_masks, source_packages=source_packages, developer_connect_source=developer_connect_source, + git_source=git_source, class_methods=class_methods, entrypoint_module=entrypoint_module, entrypoint_object=entrypoint_object, @@ -2373,7 +2399,7 @@ def _create_config( "`max_instances`, `resource_limits`, `container_concurrency`, or " "`keep_alive_probe`, you must also provide the `agent` variable or " "the source code options (`source_packages`, " - "`developer_connect_source` or `agent_config_source`)." + "`developer_connect_source`, `git_source` or `agent_config_source`)." ) if runtime_spec is not None: @@ -2672,6 +2698,9 @@ def update( developer_connect_source = json.loads( developer_connect_source.model_dump_json() ) + git_source = config.git_source + if git_source is not None: + git_source = json.loads(git_source.model_dump_json(exclude_none=True)) agent_config_source = config.agent_config_source if agent_config_source is not None: agent_config_source = json.loads(agent_config_source.model_dump_json()) @@ -2721,6 +2750,7 @@ def update( class_methods=config.class_methods, source_packages=config.source_packages, developer_connect_source=developer_connect_source, + git_source=git_source, entrypoint_module=config.entrypoint_module, entrypoint_object=config.entrypoint_object, requirements_file=config.requirements_file, diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index 9f140a672c..2c7711a718 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -9759,6 +9759,10 @@ class CreateRuntimeConfig(_common.BaseModel): subdirectory and the path must be added to `extra_packages`. """, ) + git_source: Optional[ReasoningEngineSpecSourceCodeSpecGitSource] = Field( + default=None, + description="""Specifies the public Git repository to fetch the source code from. This includes the URL, revision, and directory to use.""", + ) class CreateRuntimeConfigDict(TypedDict, total=False): @@ -9900,6 +9904,9 @@ class CreateRuntimeConfigDict(TypedDict, total=False): subdirectory and the path must be added to `extra_packages`. """ + git_source: Optional[ReasoningEngineSpecSourceCodeSpecGitSourceDict] + """Specifies the public Git repository to fetch the source code from. This includes the URL, revision, and directory to use.""" + CreateRuntimeConfigOrDict = Union[CreateRuntimeConfig, CreateRuntimeConfigDict] @@ -10420,6 +10427,10 @@ class UpdateRuntimeConfig(_common.BaseModel): subdirectory and the path must be added to `extra_packages`. """, ) + git_source: Optional[ReasoningEngineSpecSourceCodeSpecGitSource] = Field( + default=None, + description="""Specifies the public Git repository to fetch the source code from. This includes the URL, revision, and directory to use.""", + ) update_mask: Optional[str] = Field( default=None, description="""The update mask to apply. For the `FieldMask` definition, see @@ -10570,6 +10581,9 @@ class UpdateRuntimeConfigDict(TypedDict, total=False): subdirectory and the path must be added to `extra_packages`. """ + git_source: Optional[ReasoningEngineSpecSourceCodeSpecGitSourceDict] + """Specifies the public Git repository to fetch the source code from. This includes the URL, revision, and directory to use.""" + update_mask: Optional[str] """The update mask to apply. For the `FieldMask` definition, see https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""" @@ -33290,6 +33304,10 @@ class AgentRuntimeConfig(_common.BaseModel): replaces the deployed card; pass `{}` to remove it. The service sets the card's `url` to the runtime's A2A endpoint.""", ) + git_source: Optional[ReasoningEngineSpecSourceCodeSpecGitSource] = Field( + default=None, + description="""Specifies the public Git repository to fetch the source code from. This includes the URL, revision, and directory to use.""", + ) class AgentRuntimeConfigDict(TypedDict, total=False): @@ -33489,6 +33507,9 @@ class AgentRuntimeConfigDict(TypedDict, total=False): replaces the deployed card; pass `{}` to remove it. The service sets the card's `url` to the runtime's A2A endpoint.""" + git_source: Optional[ReasoningEngineSpecSourceCodeSpecGitSourceDict] + """Specifies the public Git repository to fetch the source code from. This includes the URL, revision, and directory to use.""" + AgentRuntimeConfigOrDict = Union[AgentRuntimeConfig, AgentRuntimeConfigDict] diff --git a/tests/unit/agentplatform/genai/test_agent_engines.py b/tests/unit/agentplatform/genai/test_agent_engines.py index 5e52b86139..e4e9a504ee 100644 --- a/tests/unit/agentplatform/genai/test_agent_engines.py +++ b/tests/unit/agentplatform/genai/test_agent_engines.py @@ -469,6 +469,11 @@ def register_operations(self) -> Dict[str, List[str]]: _TEST_AGENT_ENGINE_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = _TEST_STANDARD_API_MODE _TEST_PYTHON_VERSION = f"{sys.version_info.major}.{sys.version_info.minor}" _TEST_PYTHON_VERSION_OVERRIDE = "3.11" +_TEST_GIT_SOURCE = { + "url": "https://github.com/test-owner/test-repo.git", + "dir": "agent", + "revision": "main", +} _TEST_AGENT_ENGINE_FRAMEWORK = _runtimes_utils._DEFAULT_AGENT_FRAMEWORK _TEST_AGENT_ENGINE_CLASS_METHOD_1 = { "description": "Runs the engine.", @@ -1203,6 +1208,97 @@ def test_create_runtime_config_with_developer_connect_source(self): ) assert "keep_alive_probe" not in config["spec"].get("deployment_spec", {}) + def test_create_runtime_config_with_git_source(self): + config = self.client.runtimes._create_config( + mode="create", + git_source=_TEST_GIT_SOURCE, + entrypoint_module="main", + entrypoint_object="app", + class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS, + python_version=_TEST_PYTHON_VERSION_OVERRIDE, + ) + # Unlike `developer_connect_source`, `git_source` is serialized flat: + # the proto has no wrapper message around it. + assert config["spec"]["source_code_spec"] == { + "git_source": _TEST_GIT_SOURCE, + "python_spec": { + "version": _TEST_PYTHON_VERSION_OVERRIDE, + "entrypoint_module": "main", + "entrypoint_object": "app", + }, + } + + def test_update_runtime_config_with_git_source(self): + config = self.client.runtimes._create_config( + mode="update", + git_source=_TEST_GIT_SOURCE, + entrypoint_module="main", + entrypoint_object="app", + class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS, + python_version=_TEST_PYTHON_VERSION_OVERRIDE, + ) + assert config["spec"]["source_code_spec"] == { + "git_source": _TEST_GIT_SOURCE, + "python_spec": { + "version": _TEST_PYTHON_VERSION_OVERRIDE, + "entrypoint_module": "main", + "entrypoint_object": "app", + }, + } + assert config["update_mask"] == ",".join( + [ + "spec.source_code_spec.git_source", + "spec.class_methods", + "spec.source_code_spec.python_spec.version", + "spec.source_code_spec.python_spec.entrypoint_module", + "spec.source_code_spec.python_spec.entrypoint_object", + "spec.agent_framework", + ] + ) + + def test_create_runtime_config_with_git_source_requires_entrypoint(self): + with pytest.raises(ValueError, match="`entrypoint_module` must be specified"): + self.client.runtimes._create_config( + mode="create", + git_source=_TEST_GIT_SOURCE, + entrypoint_object="app", + class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS, + ) + + def test_create_runtime_config_with_git_source_and_others_raises(self): + def create_config_with_git_source(**other_source_kwargs): + return self.client.runtimes._create_config( + mode="create", + git_source=_TEST_GIT_SOURCE, + entrypoint_module="main", + entrypoint_object="app", + class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS, + **other_source_kwargs, + ) + + with pytest.raises(ValueError) as excinfo: + create_config_with_git_source(agent=self.test_agent) + assert "`git_source`" in str(excinfo.value) + + with pytest.raises(ValueError) as excinfo: + create_config_with_git_source(source_packages=["."]) + assert "`git_source`" in str(excinfo.value) + + with pytest.raises(ValueError) as excinfo: + create_config_with_git_source( + developer_connect_source={ + "git_repository_link": "projects/p/locations/l/connections/c/gitRepositoryLinks/r", + "revision": "main", + }, + ) + assert "`git_source`" in str(excinfo.value) + + with pytest.raises(ValueError) as excinfo: + create_config_with_git_source( + container_spec={"image_uri": "gcr.io/test-project/test-image"}, + ) + assert "`git_source`" in str(excinfo.value) + @mock.patch.object( _runtimes_utils, "_create_base64_encoded_tarball", @@ -2422,6 +2518,7 @@ def test_create_runtime_with_env_vars_dict( class_methods=None, source_packages=None, developer_connect_source=None, + git_source=None, entrypoint_module=None, entrypoint_object=None, requirements_file=None, @@ -2530,6 +2627,7 @@ def test_create_runtime_with_custom_service_account( class_methods=None, source_packages=None, developer_connect_source=None, + git_source=None, entrypoint_module=None, entrypoint_object=None, requirements_file=None, @@ -2637,6 +2735,7 @@ def test_create_runtime_with_experimental_mode( class_methods=None, source_packages=None, developer_connect_source=None, + git_source=None, entrypoint_module=None, entrypoint_object=None, requirements_file=None, @@ -2813,6 +2912,7 @@ def test_create_runtime_with_class_methods( class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS, source_packages=None, developer_connect_source=None, + git_source=None, entrypoint_module=None, entrypoint_object=None, requirements_file=None, @@ -2914,6 +3014,7 @@ def test_create_runtime_with_agent_framework( class_methods=None, source_packages=None, developer_connect_source=None, + git_source=None, entrypoint_module=None, entrypoint_object=None, requirements_file=None, @@ -3005,6 +3106,65 @@ def test_update_runtime_requirements( f"name='{_TEST_AGENT_ENGINE_RESOURCE_NAME}')" in caplog.text ) + @mock.patch.object(_runtimes_utils, "_await_operation") + @mock.patch.object( + _runtimes_utils, + "_get_reasoning_engine_id", + return_value=_TEST_RESOURCE_ID, + ) + def test_create_runtime_git_source( + self, mock_get_reasoning_engine_id, mock_await_operation + ): + mock_await_operation.return_value = _genai_types.RuntimeOperation( + response=_genai_types.ReasoningEngine( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + spec=_TEST_AGENT_ENGINE_SPEC, + ) + ) + git_source = {"url": _TEST_GIT_SOURCE["url"], "revision": "main"} + with mock.patch.object( + self.client.runtimes._api_client, "request" + ) as request_mock: + request_mock.return_value = genai_types.HttpResponse(body="") + self.client.runtimes.create( + config=_genai_types.AgentRuntimeConfig( + git_source=git_source, + entrypoint_module="main", + entrypoint_object="app", + class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS, + ), + ) + body = request_mock.call_args.args[2] + # No "dir": null for the field the caller left out. + assert body["spec"]["source_code_spec"]["git_source"] == git_source + + @mock.patch.object(_runtimes_utils, "_await_operation") + def test_update_runtime_git_source(self, mock_await_operation): + mock_await_operation.return_value = _genai_types.RuntimeOperation( + response=_genai_types.ReasoningEngine( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + spec=_TEST_AGENT_ENGINE_SPEC, + ) + ) + with mock.patch.object( + self.client.runtimes._api_client, "request" + ) as request_mock: + request_mock.return_value = genai_types.HttpResponse(body="") + self.client.runtimes.update( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + config=_genai_types.AgentRuntimeConfig( + git_source=_TEST_GIT_SOURCE, + entrypoint_module="main", + entrypoint_object="app", + class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS, + ), + ) + body = request_mock.call_args.args[2] + assert body["spec"]["source_code_spec"]["git_source"] == _TEST_GIT_SOURCE + assert body["_query"]["updateMask"].startswith( + "spec.source_code_spec.git_source," + ) + @mock.patch.object(_runtimes_utils, "_prepare") @mock.patch.object(_runtimes_utils, "_await_operation") def test_update_runtime_extra_packages(