fix: GenAI Client(evals) - skip agent config fetch when agent and client locations differ

PiperOrigin-RevId: 956029497

Author: vertex-sdk-bot

agentplatform/_genai/_evals_common.py

@@ -904,31 +904,56 @@ def _fetch_agent_config_dict(
         An AgentConfig with ``agent_id`` and, when available,
         ``instruction``, ``description``, ``agent_type``, and ``tools``.
     """
-    agent_short_id = agent_resource_name.split("/")[-1] or "agent"
+    parts = agent_resource_name.split("/")
+    agent_location = None
+    agent_short_id = "agent"
+
+    # Expected format: projects/{project}/locations/{location}/agents/{agent_id}
+    if (
+        len(parts) >= 6
+        and parts[0] == "projects"
+        and parts[2] == "locations"
+        and parts[4] == "agents"
+    ):
+        agent_location = parts[3]
+        agent_short_id = parts[5]
+    else:
+        agent_short_id = parts[-1] or "agent"
 
     instruction: Optional[str] = None
     description: Optional[str] = None
     agent_type: Optional[str] = None
     tools: Optional[list[genai_types.Tool]] = None
 
-    # TODO(b/539762376): This drops the location from `agent_resource_name` and
-    # sends a relative path, so the client re-qualifies it with its own
-    # location. Agents in `locations/global` are queried in the client's region
-    # and fail with 400, silently leaving the returned AgentConfig empty.
-    try:
-        agent_resp = api_client.request("get", f"agents/{agent_short_id}", {}, None)
-        if agent_resp.body:
-            agent_dict = json.loads(agent_resp.body)
-            instruction = agent_dict.get("system_instruction") or None
-            description = agent_dict.get("description") or None
-            agent_type = agent_dict.get("base_agent") or None
-            tools = _agent_tools_to_config_tools(agent_dict.get("tools"))
-    except Exception as e:  # pylint: disable=broad-exception-caught
+    client_location = _get_resolved_location(api_client)
+
+    if agent_location and client_location and agent_location != client_location:
         logger.warning(
-            "Failed to fetch agent config for '%s' (continuing without it): %s",
+            "Skipping agent config fetch for '%s' due to location mismatch. "
+            "Agent location is '%s', but client location is '%s'. "
+            "To fetch the agent config, configure a client with matching location.",
             agent_resource_name,
-            e,
+            agent_location,
+            client_location,
         )
+    else:
+        try:
+            request_path = (
+                agent_resource_name if agent_location else f"agents/{agent_short_id}"
+            )
+            agent_resp = api_client.request("get", request_path, {}, None)
+            if agent_resp.body:
+                agent_dict = json.loads(agent_resp.body)
+                instruction = agent_dict.get("system_instruction") or None
+                description = agent_dict.get("description") or None
+                agent_type = agent_dict.get("base_agent") or None
+                tools = _agent_tools_to_config_tools(agent_dict.get("tools"))
+        except Exception as e:  # pylint: disable=broad-exception-caught
+            logger.warning(
+                "Failed to fetch agent config for '%s' (continuing without it): %s",
+                agent_resource_name,
+                e,
+            )
 
     return types.evals.AgentConfig(  # pytype: disable=missing-parameter
         agent_id=agent_short_id,
@@ -941,7 +966,10 @@ def _fetch_agent_config_dict(
 
 def _get_resolved_location(api_client: Any) -> Optional[str]:
     """Returns the location configured on the API client."""
-    return getattr(api_client, "location", None)
+    loc = getattr(api_client, "location", None)
+    if isinstance(loc, str):
+        return loc
+    return None
 
 
 class _InteractionsRestClient:

tests/unit/agentplatform/genai/test_evals.py

@@ -10301,9 +10301,9 @@ def test_resolve_interactions_to_eval_cases_happy_path(self):
 
         def mock_request(method, path, *args, **kwargs):
             resp = mock.MagicMock()
-            if path.startswith("interactions/"):
+            if path.endswith("interactions/i1"):
                 resp.body = interaction_body
-            elif path.startswith("agents/"):
+            elif path.endswith("agents/my-agent"):
                 resp.body = agent_body
             else:
                 resp.body = None
@@ -10314,9 +10314,9 @@ def mock_request(method, path, *args, **kwargs):
         original_case = agentplatform_genai_types.EvalCase(
             interactions_data_source=agentplatform_genai_types.InteractionsDataSource(
                 gemini_agent_config=agentplatform_genai_types.GeminiAgentConfig(
-                    gemini_agent="projects/p/locations/l/agents/my-agent",
+                    gemini_agent="projects/p/locations/us-central1/agents/my-agent",
                 ),
-                interaction="projects/p/locations/l/interactions/i1",
+                interaction="projects/p/locations/us-central1/interactions/i1",
             ),
         )
 
@@ -11119,7 +11119,12 @@ def mock_request(method, path, *args, **kwargs):
         # Verify the interaction and agent were fetched.
         assert mock_api_client.request.call_count == 2
         mock_api_client.request.assert_any_call("get", "interactions/test-id", {}, None)
-        mock_api_client.request.assert_any_call("get", "agents/my-agent", {}, None)
+        mock_api_client.request.assert_any_call(
+            "get",
+            "projects/p/locations/l/agents/my-agent",
+            {},
+            None,
+        )
 
     def test_agents_map_populated_with_full_config(self):
         """The agents dict includes instruction and tools from the Agent API."""
@@ -12072,6 +12077,40 @@ def test_mcp_server_kept_as_named_declaration(self):
         assert len(mcp_tool) == 1
         assert "my-mcp" in mcp_tool[0].function_declarations[0].description
 
+    def test_skips_fetch_when_location_mismatches(self):
+        """Skips agent config fetch and logs warning if locations mismatch."""
+        mock_api_client = mock.MagicMock()
+        mock_api_client.location = "us-central1"
+        mock_api_client.request.return_value = self._make_api_response({})
+
+        result = _evals_common._fetch_agent_config_dict(
+            mock_api_client,
+            "projects/p/locations/global/agents/my-agent",
+        )
+        assert result.agent_id == "my-agent"
+        assert result.instruction is None
+        mock_api_client.request.assert_not_called()
+
+    def test_uses_full_resource_name_when_location_matches(self):
+        """Uses full resource name in GET request when client location matches."""
+        agent_json = {"system_instruction": "You are helpful."}
+        mock_api_client = mock.MagicMock()
+        mock_api_client.location = "us-central1"
+        mock_api_client.request.return_value = self._make_api_response(agent_json)
+
+        result = _evals_common._fetch_agent_config_dict(
+            mock_api_client,
+            "projects/p/locations/us-central1/agents/my-agent",
+        )
+        assert result.agent_id == "my-agent"
+        assert result.instruction == "You are helpful."
+        mock_api_client.request.assert_called_once_with(
+            "get",
+            "projects/p/locations/us-central1/agents/my-agent",
+            {},
+            None,
+        )
+
     def test_catalog_in_sync_with_server(self):
         """SDK catalog keys and function names match the server-side catalog.