feat(dea): parameterize url_agent_name and agent_type_uri options

Author: James Nguyen

datasets/dea-tools/README.md

@@ -30,6 +30,26 @@ This file is the main driver that dictates how evaluations are executed. Key par
 *   `set_up_script` & `tear_down_script`: Scripts that prep and clean up the Dataform environment. You can comment out `tear_down_script` if you want to keep the workspace alive after the run for debugging.
 *   `dataform_workspace_gcs_archive`: Saves a ZIP archive of the workspace artifacts to a GCS bucket.
 
+### Parameterized Agent & Environment Configuration:
+The model configuration file is defined in `datasets/model_configs/gcp_data_engineering_agent_model.yaml`. In your model configuration YAML file, you can configure target environments and agent personas:
+
+```yaml
+generator: "data_engineering_agent"
+gcp_project_id: !ENV ${EVAL_GCP_PROJECT_ID}
+gcp_region: !ENV ${EVAL_GCP_PROJECT_REGION}
+
+# Parameterized Agent Options:
+model_env: local                # Target environment: "local", "staging", or "prod" (defaults to "prod")
+port: 9876                      # Required port for local Boq servers when env="local"
+agent_type: sparkagent          # Desired agent persona (e.g., "sparkagent", "dataengineeringagent")
+
+```
+
+* **`model_env`**: Configures the host environment (`local` uses `http://localhost:{port}`, `staging` uses `https://staging-geminidataanalytics.sandbox.googleapis.com`, `prod` uses `https://geminidataanalytics.googleapis.com`).
+
+* **`port`**: Specifies the HTTP port for local Boq servers (required when `env="local"`).
+* **`agent_type`**: Sets the desired agent persona (e.g., "sparkagent", "dataengineeringagent").
+
 ## 3. Run EvalBench
 
 You can run EvalBench for DEA in two modes:

datasets/model_configs/gcp_data_engineering_agent_model.yaml

@@ -3,4 +3,7 @@ gcp_project_id: !ENV ${EVAL_GCP_PROJECT_ID}
 gcp_region: !ENV ${EVAL_GCP_PROJECT_REGION:us-west4}
 dataform_repository: !ENV ${EVAL_DEA_REPOSITORY_ID}
 dataform_workspace: !ENV ${EVAL_DEA_WORKSPACE_ID}
+
+model_env: prod
+agent_type: dataengineeringagent
 execs_per_minute: 10

evalbench/generators/models/gcp_data_engineering_agent.py

@@ -50,11 +50,21 @@
     "https://geminidataanalytics.googleapis.com/a2a/extensions/"
     "finishreason/v1"
 )
+AGENT_TYPE_URI = (
+    "https://geminidataanalytics.googleapis.com/a2a/extensions/"
+    "agenttype/v1"
+)
+
+# Mapping from agent_type to AGENT_TYPE_URI value
+AGENT_TYPE_MAP = {
+    "sparkagent": "SPARK_AGENT",
+}
+
 
 # All required A2A Extension Headers combined
 ALL_EXTENSIONS = (
     f"{MESSAGE_LEVEL_URI},{INSTRUCTION_URI},{GCP_RESOURCE_URI},"
-    f"{CONVERSATION_TOKEN_URI},{FINISH_REASON_URI}"
+    f"{CONVERSATION_TOKEN_URI},{FINISH_REASON_URI},{AGENT_TYPE_URI}"
 )
 
 logger = logging.getLogger(__name__)
@@ -211,6 +221,28 @@ def __init__(self, querygenerator_config: dict[str, Any]):
         self.name = "data_engineering_agent"
         gcp_project_id = querygenerator_config.get("gcp_project_id", "")
         gcp_region = querygenerator_config.get("gcp_region", "")
+        env_val = querygenerator_config.get("model_env")
+
+        if not isinstance(env_val, str):
+            raise TypeError(
+                "Configuration key 'model_env' must be a string, "
+                f"got {type(env_val).__name__}."
+            )
+        env = env_val.lower()
+
+        self.agent_type = querygenerator_config.get(
+            "agent_type", "dataengineeringagent"
+        ).lower()
+        match self.agent_type:
+            case "dataengineeringagent":
+                self.agent_type_uri = None
+            case "sparkagent":
+                self.agent_type_uri = AGENT_TYPE_MAP.get(self.agent_type)
+            case _:
+                raise ValueError(
+                    f"Unsupported agent_type '{self.agent_type}'. "
+                    "Must be 'dataengineeringagent' or 'sparkagent'."
+                )
 
         if not gcp_project_id:
             raise ValueError(
@@ -223,10 +255,26 @@ def __init__(self, querygenerator_config: dict[str, Any]):
                 "DataEngineeringAgentGenerator."
             )
 
+        if env == "local":
+            port = querygenerator_config.get("port", 9876)
+
+            host = f"http://localhost:{port}"
+        elif env == "staging":
+            host = (
+                "https://staging-geminidataanalytics."
+                "sandbox.googleapis.com"
+            )
+        elif env == "prod":
+            host = "https://geminidataanalytics.googleapis.com"
+        else:
+            raise ValueError(
+                f"Unsupported env: '{env}'. "
+                "Expected 'local', 'staging', or 'prod'."
+            )
+
         self.endpoint = (
-            f"https://geminidataanalytics.googleapis.com/v1/a2a/projects/"
-            f"{gcp_project_id}/locations/{gcp_region}/"
-            f"agents/dataengineeringagent"
+            f"{host}/v1/a2a/projects/{gcp_project_id}/"
+            f"locations/{gcp_region}/agents/{self.agent_type}"
         )
 
         self.auth_interceptor = AuthInterceptor(GcpAdcCredentialService())
@@ -324,6 +372,9 @@ async def _run_client(
         message_req.metadata[GCP_RESOURCE_URI] = {
             "gcpResourceId": target_workspace
         }
+        # Configure Agent Type extension
+        if self.agent_type_uri:
+            message_req.metadata[AGENT_TYPE_URI] = self.agent_type_uri
 
         # Handle ConversationToken state memory thread-safely
         token = ""

evalbench/test/gcp_data_engineering_agent_test.py

@@ -55,6 +55,7 @@ def valid_config():
         "gcp_region": "us-east1",
         "dataform_repository": "test-repo",
         "dataform_workspace": "test-workspace",
+        "model_env": "prod",
     }
 
 
@@ -165,6 +166,97 @@ def test_data_engineering_agent_generator_setup(valid_config):
         assert generator.auth_interceptor is not None
 
 
+def test_generator_parameterized_staging_env(valid_config):
+    config = valid_config.copy()
+    config["model_env"] = "staging"
+    with patch("google.auth.default") as mock_auth_default:
+        mock_creds = MagicMock()
+        mock_creds.valid = True
+        mock_auth_default.return_value = (mock_creds, "test-project")
+
+        generator = DataEngineeringAgentGenerator(config)
+
+        expected_endpoint = (
+            "https://staging-geminidataanalytics.sandbox.googleapis.com/"
+            "v1/a2a/projects/test-project-123/locations/us-east1/agents/"
+            "dataengineeringagent"
+        )
+        assert generator.endpoint == expected_endpoint
+
+
+def test_generator_parameterized_sparkagent(valid_config):
+    config = valid_config.copy()
+    config["model_env"] = "staging"
+    config["agent_type"] = "sparkagent"
+    with patch("google.auth.default") as mock_auth_default:
+        mock_creds = MagicMock()
+        mock_creds.valid = True
+        mock_auth_default.return_value = (mock_creds, "test-project")
+
+        generator = DataEngineeringAgentGenerator(config)
+
+        expected_endpoint = (
+            "https://staging-geminidataanalytics.sandbox.googleapis.com/"
+            "v1/a2a/projects/test-project-123/locations/us-east1/agents/"
+            "sparkagent"
+        )
+        assert generator.endpoint == expected_endpoint
+        assert generator.agent_type == "sparkagent"
+        assert generator.agent_type_uri == "SPARK_AGENT"
+
+
+def test_generator_parameterized_local_env(valid_config):
+    config = valid_config.copy()
+    config["model_env"] = "local"
+    config["port"] = 9876
+    config["agent_type"] = "sparkagent"
+    with patch("google.auth.default") as mock_auth_default:
+        mock_creds = MagicMock()
+        mock_creds.valid = True
+        mock_auth_default.return_value = (mock_creds, "test-project")
+
+        generator = DataEngineeringAgentGenerator(config)
+
+        expected_endpoint = (
+            "http://localhost:9876/v1/a2a/projects/test-project-123/"
+            "locations/us-east1/agents/sparkagent"
+        )
+        assert generator.endpoint == expected_endpoint
+
+
+def test_generator_parameterized_local_port_override(valid_config):
+    config = valid_config.copy()
+    config["model_env"] = "local"
+    config["port"] = 8080
+    config["agent_type"] = "dataengineeringagent"
+    with patch("google.auth.default") as mock_auth_default:
+        mock_creds = MagicMock()
+        mock_creds.valid = True
+        mock_auth_default.return_value = (mock_creds, "test-project")
+
+        generator = DataEngineeringAgentGenerator(config)
+
+        expected_endpoint = (
+            "http://localhost:8080/v1/a2a/projects/test-project-123/"
+            "locations/us-east1/agents/dataengineeringagent"
+        )
+        assert generator.endpoint == expected_endpoint
+
+
+def test_generator_setup_unsupported_env_raises_value_error(valid_config):
+    config = valid_config.copy()
+    config["model_env"] = "invalid_env"
+    with patch("google.auth.default") as mock_auth_default:
+        mock_creds = MagicMock()
+        mock_creds.valid = True
+        mock_auth_default.return_value = (mock_creds, "test-project")
+
+        with pytest.raises(ValueError) as excinfo:
+            DataEngineeringAgentGenerator(config)
+
+        assert "Unsupported env: 'invalid_env'" in str(excinfo.value)
+
+
 @pytest.mark.anyio
 @patch("google.auth.default")
 @patch("generators.models.gcp_data_engineering_agent.create_client")
@@ -305,6 +397,23 @@ async def mock_close():
     assert called_card.capabilities.extended_agent_card is True
 
 
+def test_generator_setup_unsupported_agent_type_raises_error(
+    valid_config,
+):
+    config = valid_config.copy()
+    config["agent_type"] = "unsupportedagent"
+
+    with patch("google.auth.default") as mock_auth_default:
+        mock_creds = MagicMock()
+        mock_creds.valid = True
+        mock_auth_default.return_value = (mock_creds, "test-project")
+
+        with pytest.raises(ValueError) as excinfo:
+            DataEngineeringAgentGenerator(config)
+
+        assert "Unsupported agent_type 'unsupportedagent'" in str(excinfo.value)
+
+
 @patch("evaluator.dataengineeringagentevaluator.AgentScoreWork")
 @patch("evaluator.dataengineeringagentevaluator.SimulatedUser")
 def test_evaluator_process_scenario(