feat(dea): parameterize agent_name, env, domain and add AGENT_TYPE_URI extension

Author: James Nguyen

datasets/dea-tools/README.md

@@ -23,6 +23,24 @@ This file is the main driver that dictates how evaluations are executed. Key par
 *   `simulated_user_model_config`: Points to the model acting as the user (Gemini 2.5 Pro).
 *   `dataset_config`: Points to the Dataset JSON file (described in Step 1).
 *   `env` block: Sets local variables like `EVAL_DATAFORM_SETUP_ENV_FILES_DIR` to inject existing files (e.g., workspace settings, schema definitions) into the workspace before the agent runs. This allows the agent to build on top of an established environment instead of starting from scratch.
+
+### Parameterized Agent & Environment Configuration:
+The model configuration file is defined in `datasets/dea-tools/example_model_config.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:
+env: staging            # Target environment: "staging" or "prod" (defaults to "prod")
+agent_name: sparkagent  # Persona/Agent name: "sparkagent" or "dataengineeringagent"
+domain: ""              # Optional custom host override (e.g., "https://staging-geminidataanalytics.sandbox.googleapis.com")
+```
+
+* **`env`**: Configures the host environment (`staging` uses `https://staging-geminidataanalytics.sandbox.googleapis.com`, `prod` uses `https://geminidataanalytics.googleapis.com`).
+* **`agent_name`**: Sets the agent persona path parameter (`/v1/a2a/.../agents/{agent_name}`) and automatically attaches the `AGENT_TYPE_URI` metadata header extension (`"SPARK_AGENT"`).
+* **`domain`**: (Optional) Fully overrides the target base host.
 *   `scorers`: Defines the evaluation scorers:
     *   `binary_rubric_scorer`: Grader that evaluates the conversation against the `binary_rubric` (gives 0 or 100).
     *   `dataform_cloud_compile` & `dataform_cloud_run`: Cloud verifiers that compile and run the generated SQLX files in GCP to ensure they are physically valid.

datasets/dea-tools/example_model_config.yaml

@@ -0,0 +1,6 @@
+generator: "data_engineering_agent"
+gcp_project_id: !ENV ${EVAL_GCP_PROJECT_ID}
+gcp_region: !ENV ${EVAL_GCP_PROJECT_REGION}
+env: prod
+agent_name: dataengineeringagent
+execs_per_minute: 10

datasets/dea-tools/example_run_config.yaml

@@ -1,6 +1,6 @@
 orchestrator: dea
 dataset_format: dea-format
-model_config: datasets/model_configs/gcp_data_engineering_agent_model.yaml
+model_config: datasets/dea-tools/example_model_config.yaml
 simulated_user_model_config: datasets/model_configs/gemini_2.5_pro_model.yaml
 dataset_config: datasets/dea-tools/dea-live-conversational.evalset.json
 env:

evalbench/generators/models/gcp_data_engineering_agent.py

@@ -50,11 +50,15 @@
     "https://geminidataanalytics.googleapis.com/a2a/extensions/"
     "finishreason/v1"
 )
+AGENT_TYPE_URI = (
+    "https://geminidataanalytics.googleapis.com/a2a/extensions/"
+    "agenttype/v1"
+)
 
 # 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 +215,11 @@ 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("env", "prod")
+        env = env_val.lower() if isinstance(env_val, str) else "prod"
+        agent_name = querygenerator_config.get("agent_name", "dataengineeringagent")
+        domain = querygenerator_config.get("domain", "")
+        self.agent_name = agent_name
 
         if not gcp_project_id:
             raise ValueError(
@@ -223,10 +232,19 @@ def __init__(self, querygenerator_config: dict[str, Any]):
                 "DataEngineeringAgentGenerator."
             )
 
+        if domain:
+            if not (domain.startswith("http://") or domain.startswith("https://")):
+                host = f"https://{domain}"
+            else:
+                host = domain
+        elif env == "staging":
+            host = "https://staging-geminidataanalytics.sandbox.googleapis.com"
+        else:
+            host = "https://geminidataanalytics.googleapis.com"
+
         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}/locations/{gcp_region}/"
+            f"agents/{agent_name}"
         )
 
         self.auth_interceptor = AuthInterceptor(GcpAdcCredentialService())
@@ -324,6 +342,11 @@ async def _run_client(
         message_req.metadata[GCP_RESOURCE_URI] = {
             "gcpResourceId": target_workspace
         }
+        # Configure Agent Type extension
+        if self.agent_name.lower() in ("sparkagent", "spark_agent", "spark"):
+            message_req.metadata[AGENT_TYPE_URI] = "SPARK_AGENT"
+        elif self.agent_name.lower() not in ("dataengineeringagent", "data_engineering_agent", "dea"):
+            message_req.metadata[AGENT_TYPE_URI] = self.agent_name.upper()
 
         # Handle ConversationToken state memory thread-safely
         token = ""