fix: GenAI SDK client - do not register API methods in get or update for agent_engines with lightweight creation

yeesian · copybara-github · commit d76411a70416 · 2025-08-25T13:18:01.000-07:00
PiperOrigin-RevId: 799239717
diff --git a/tests/unit/vertexai/genai/test_agent_engines.py b/tests/unit/vertexai/genai/test_agent_engines.py
@@ -1882,19 +1882,19 @@ def setup_method(self):
         ],
     )
     @pytest.mark.usefixtures("caplog")
-    @mock.patch.object(_genai_types.AgentEngine, "operation_schemas")
     @mock.patch.object(agent_engines.AgentEngines, "_get")
     def test_invalid_operation_schema(
         self,
         mock_get,
-        mock_operation_schemas,
         test_case_name,
         test_operation_schemas,
         want_log_output,
         caplog,
     ):
-        mock_get.return_value = _genai_types.AgentEngine()  # just to avoid an API call
-        mock_operation_schemas.return_value = test_operation_schemas
+        mock_get.return_value = _genai_types.ReasoningEngine(
+            name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
+            spec=_genai_types.ReasoningEngineSpec(class_methods=test_operation_schemas),
+        )
         self.client.agent_engines.get(name=_TEST_AGENT_ENGINE_RESOURCE_NAME)
         assert want_log_output in caplog.text
 
diff --git a/vertexai/_genai/agent_engines.py b/vertexai/_genai/agent_engines.py
@@ -970,12 +970,15 @@ def get(
               "projects/123/locations/us-central1/reasoningEngines/456" or a
               shortened name such as "reasoningEngines/456".
         """
+        api_resource = self._get(name=name, config=config)
         agent_engine = types.AgentEngine(
             api_client=self,
             api_async_client=AsyncAgentEngines(api_client_=self._api_client),
-            api_resource=self._get(name=name, config=config),
+            api_resource=api_resource,
         )
-        return self._register_api_methods(agent_engine=agent_engine)
+        if api_resource.spec:
+            self._register_api_methods(agent_engine=agent_engine)
+        return agent_engine
 
     def create(
         self,
@@ -1513,7 +1516,9 @@ def update(
             logger.info(
                 f"agent_engine=client.agent_engines.get('{agent_engine.api_resource.name}')"
             )
-        return self._register_api_methods(agent_engine=agent_engine)
+        if agent_engine.api_resource.spec:
+            self._register_api_methods(agent_engine=agent_engine)
+        return agent_engine
 
     def _stream_query(
         self,
