From 52e29fdc056f8c53f5e1a8824fc3b3737444b0a5 Mon Sep 17 00:00:00 2001 From: Ariel Jassan Date: Tue, 28 Jul 2026 15:06:13 +0000 Subject: [PATCH 1/5] feat(evalbench): add native Agent Runtime generator support and deployment guide --- .../bird/test_agent_runtime_run_config.yaml | 41 ++ datasets/model_configs/agent_runtime.yaml | 14 + docs/configs/model-config.md | 28 +- docs/examples/agent_runtime_deployment.ipynb | 464 ++++++++++++++++++ evalbench/generators/models/__init__.py | 2 + evalbench/generators/models/agent_runtime.py | 149 ++++++ evalbench/test/agent_runtime_test.py | 139 ++++++ evalbench/work/sqlexecwork.py | 2 +- 8 files changed, 836 insertions(+), 3 deletions(-) create mode 100644 datasets/bird/test_agent_runtime_run_config.yaml create mode 100644 datasets/model_configs/agent_runtime.yaml create mode 100644 docs/examples/agent_runtime_deployment.ipynb create mode 100644 evalbench/generators/models/agent_runtime.py create mode 100644 evalbench/test/agent_runtime_test.py diff --git a/datasets/bird/test_agent_runtime_run_config.yaml b/datasets/bird/test_agent_runtime_run_config.yaml new file mode 100644 index 00000000..f59ff76c --- /dev/null +++ b/datasets/bird/test_agent_runtime_run_config.yaml @@ -0,0 +1,41 @@ +############################################################ +### Dataset / Eval Items (Vertex AI Multi-Agent Test) +############################################################ +dataset_config: datasets/bird/prompts_california.json + +databases: + - california_schools +num_trials: 1 + +database_configs: + - datasets/bat/db_configs/bigquery.yaml + - datasets/bird/db_configs/sqlite.yaml +dialects: + - bigquery +dialect: bigquery +query_types: + - dql +dataset_format: bird-standard-format + +############################################################ +### Prompt and Generation Modules +############################################################ +model_config: datasets/model_configs/agent_runtime.yaml +prompt_generator: 'NOOPGenerator' + +############################################################ +### Scorer Related Configs +############################################################ +scorers: + python_scorer: + script_path: 'evalbench/scorers/judges/hybrid_xa_judge.py' + scorer_name: 'hybrid_cross_db' + +############################################################ +### Reporting Related Configs +############################################################ +reporting: + bigquery: + dataset_location: "US" + gcp_project_id: "arielj-argolis-1" + csv: {} diff --git a/datasets/model_configs/agent_runtime.yaml b/datasets/model_configs/agent_runtime.yaml new file mode 100644 index 00000000..2f944f0c --- /dev/null +++ b/datasets/model_configs/agent_runtime.yaml @@ -0,0 +1,14 @@ +# Agent Runtime (Gemini Enterprise Agent Platform) Generator Manifest +# Enables native NL2SQL evaluation benchmarking against deployed Agent Runtime endpoints. +generator: agent_runtime + +# Explicit resource URI string (e.g. "projects/12345/locations/us-central1/reasoningEngines/67890") +# Leave empty string "" to dynamically read MULTI_AGENT_ENGINE_RESOURCE or AGENT_ENGINE_RESOURCE from env. +resource_name: "projects/743100398377/locations/us-central1/reasoningEngines/7921856483556327424" + +gcp_project_id: "arielj-argolis-1" +gcp_region: "us-central1" + +# Rate limiting and retry mechanics +execs_per_minute: 10 +max_attempts: 3 diff --git a/docs/configs/model-config.md b/docs/configs/model-config.md index b03c7e82..3a70f0e1 100644 --- a/docs/configs/model-config.md +++ b/docs/configs/model-config.md @@ -28,6 +28,16 @@ These settings are **required only** for generators that utilize Google Cloud Ve > Required*, you can globally set your GCP project_id and gcp_region using the environment variables `EVAL_GCP_PROJECT_ID` and `EVAL_GCP_PROJECT_REGION`. +## Agent Runtime (Gemini Enterprise Agent Platform) Configuration + +These settings are **required only** for the `agent_runtime` generator, which connects to a live deployed Agent Runtime (Gemini Enterprise Agent Platform) instance. + +| **Key** | **Required** | **Default Value** | **Description** | +| ----------------- | ------------ | ----------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `resource_name` | Yes | N/A | The fully-qualified resource ID of the deployed Agent Engine resource, e.g. `projects//locations//reasoningEngines/`. | +| `gcp_project_id` | Optional | `""` | The Google Cloud Project ID that hosts your Vertex AI resources. Can also be set via `EVAL_GCP_PROJECT_ID` environment variable. | +| `gcp_region` | Optional | `""` | The Google Cloud region where the Vertex AI service is deployed. Can also be set via `EVAL_GCP_PROJECT_REGION` environment variable. | + ## Important Notes - **Customization:** This configuration is fully customizable to the needs of the selected generator. You can add or remove keys as necessary. @@ -37,10 +47,11 @@ These settings are **required only** for generators that utilize Google Cloud Ve - **Rate Limiting & Retries:** The `execs_per_minute` and `max_attempts` keys help control the query generation process, ensuring that you can stay below project quota limits. -## Example Configuration +## Example Configurations -Below is an example of the updated YAML configuration file: +Below are examples of the YAML configuration files: +### Gemini Model Example ```yaml # General Generator Configuration generator: gcp_vertex_gemini @@ -54,3 +65,16 @@ gcp_project_id: my_cool_gcp_project gcp_region: us-east5 vertex_model: gemini-2.0-pro-exp-02-05 ``` + +### Agent Runtime Example +```yaml +# General Generator Configuration +generator: agent_runtime +execs_per_minute: 10 +max_attempts: 3 + +# Agent Runtime Configuration (Required for agent_runtime) +resource_name: "projects/my-gcp-project/locations/us-central1/reasoningEngines/1234567890" +gcp_project_id: "my-gcp-project" +gcp_region: "us-central1" +``` diff --git a/docs/examples/agent_runtime_deployment.ipynb b/docs/examples/agent_runtime_deployment.ipynb new file mode 100644 index 00000000..8d722318 --- /dev/null +++ b/docs/examples/agent_runtime_deployment.ipynb @@ -0,0 +1,464 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "5b5c8393", + "metadata": { + "id": "5b5c8393" + }, + "source": [ + "# Evaluating an ADK Agent to Agent Runtime (Gemini Enterprise Agent Platform)\n", + "\n", + "This notebook guides you through the process of:\n", + "1. Writing a minimal SQL-generating agent using **Agent Development Kit (ADK)**.\n", + "2. Deploying it to **Agent Runtime (Gemini Enterprise Agent Platform)** using the `AdkApp` template.\n", + "3. Querying the live endpoint to verify it streams valid SQL.\n", + "4. Evaluating the agent in Agent Runtime." + ] + }, + { + "cell_type": "markdown", + "id": "96e96c58", + "metadata": { + "id": "96e96c58" + }, + "source": [ + "## Step 1: Repository Setup" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3a2541a2", + "metadata": { + "id": "3a2541a2" + }, + "outputs": [], + "source": [ + "# Install/Upgrade the Google Cloud AI Platform SDK, ADK, and authentication libraries\n", + "!pip install --quiet --upgrade google-cloud-aiplatform google-adk google-auth==2.49.0" + ] + }, + { + "cell_type": "markdown", + "id": "4a1023be", + "metadata": {}, + "source": [ + "If running in a hosted environment like Colab Enterprise, we clone the `evalbench` repository to load evaluation datasets, model configurations, and pipeline execution shell scripts." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ff6a777c", + "metadata": {}, + "outputs": [], + "source": [ + "!git clone https://github.com/GoogleCloudPlatform/evalbench.git" + ] + }, + { + "cell_type": "markdown", + "id": "8fa1abd1", + "metadata": {}, + "source": [ + "We install uv to manage python packages, sync the virtual environment, install the mcp library package, and compile the protobuf message files." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "94d69b26", + "metadata": { + "id": "94d69b26" + }, + "outputs": [], + "source": [ + "# Re-install dependencies and compile protobufs\n", + "!pip install --quiet uv\n", + "!cd /content/evalbench && uv sync\n", + "!cd /content/evalbench && .venv/bin/python3 -m grpc_tools.protoc \\\n", + " --proto_path=evalbench/evalproto \\\n", + " --python_out=evalbench/evalproto \\\n", + " --pyi_out=evalbench/evalproto \\\n", + " --grpc_python_out=evalbench/evalproto \\\n", + " --experimental_editions evalbench/evalproto/*.proto\n", + "\n", + "import os\n", + "os.chdir(\"/content/evalbench/docs/examples\")\n", + "print(f\"✅ Working directory set to: {os.getcwd()}\")" + ] + }, + { + "cell_type": "markdown", + "id": "476091be", + "metadata": { + "id": "476091be" + }, + "source": [ + "Download the BIRD Dataset and Database Connections:\n", + "This script downloads the natural language questions, database schemas, and SQLite database connection files required to evaluate the SQL generator." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7994b603", + "metadata": { + "id": "7994b603" + }, + "outputs": [], + "source": [ + "# Download the BIRD dataset and database connections\n", + "!cd ../.. && bash datasets/bird/download_dataset.sh" + ] + }, + { + "cell_type": "markdown", + "id": "25850460", + "metadata": { + "id": "25850460" + }, + "source": [ + "## Step 2: Define the ADK Agent\n", + "\n", + "We write the ADK agent definition code to a local file named `minimal_agent.py` using standard ADK syntax. We specify `gemini-2.5-flash` as the model name and provide custom instruction formatting rules." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8481f972", + "metadata": { + "id": "8481f972" + }, + "outputs": [], + "source": [ + "%%writefile minimal_agent.py\n", + "import google.adk as adk\n", + "\n", + "def create_agent() -> adk.Agent:\n", + " \"\"\"Creates a minimal SQL-generating ADK agent.\"\"\"\n", + " return adk.Agent(\n", + " name=\"minimal_adk_sql_agent\",\n", + " model=\"gemini-2.5-flash\",\n", + " instruction=(\n", + " \"You are a SQL generator. Generate a BigQuery SQL query based on \"\n", + " \"the user prompt. Wrap your final generated SQL query inside a \"\n", + " \"```sql code block.\"\n", + " )\n", + " )" + ] + }, + { + "cell_type": "markdown", + "id": "d153df7b", + "metadata": { + "id": "d153df7b" + }, + "source": [ + "## Step 3: Configure and initialize Agent Platform client\n", + "\n", + "Set your target Google Cloud Project and region, and initialize the SDK." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "57e7af39", + "metadata": { + "id": "57e7af39" + }, + "outputs": [], + "source": [ + "from google.colab import auth\n", + "auth.authenticate_user()" + ] + }, + { + "cell_type": "markdown", + "id": "fd22b631", + "metadata": { + "id": "fd22b631" + }, + "source": [ + "Set your target Google Cloud Project and region, and initialize the SDK." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8adb6260", + "metadata": { + "id": "8adb6260" + }, + "outputs": [], + "source": [ + "import os\n", + "import google.cloud.aiplatform as aiplatform\n", + "import agentplatform\n", + "import google.auth\n", + "\n", + "# Edit these variables with your Google Cloud details.\n", + "PROJECT_ID = \"YOUR-PROJECT-ID\"\n", + "LOCATION = \"us-central1\"\n", + "STAGING_BUCKET = f\"gs://{PROJECT_ID}-agents\"\n", + "\n", + "# Get the Colab authenticated credentials.\n", + "print(f\"Initializing global context and client...\")\n", + "credentials, _ = google.auth.default()\n", + "\n", + "# Initialize global context and client with explicit credentials.\n", + "aiplatform.init(project=PROJECT_ID, location=LOCATION, credentials=credentials)\n", + "client = (\n", + " agentplatform.Client(\n", + " project=PROJECT_ID, \n", + " location=LOCATION, \n", + " credentials=credentials\n", + " )\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "80748249", + "metadata": { + "id": "80748249" + }, + "source": [ + "## Step 4: Deploy to Agent Runtime using the AdkApp Template\n", + "\n", + "We load the local ADK agent and wrap it inside the `AdkApp` class. Then, we call `client.agent_engines.create` to deploy the agent.\n", + "\n", + "Since our agent is self-contained in `minimal_agent.py`, we only need to include `./minimal_agent.py` in the `extra_packages` deployment configuration." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fe59c109", + "metadata": { + "id": "fe59c109" + }, + "outputs": [], + "source": [ + "from agentplatform.agent_engines.templates.adk import AdkApp\n", + "\n", + "import minimal_agent\n", + "\n", + "\n", + "# Instantiate the local ADK agent and wrap it in the template.\n", + "local_agent = minimal_agent.create_agent()\n", + "adk_app = AdkApp(agent=local_agent, app_name=\"minimal-adk-sql-app\")\n", + "\n", + "print(\"Deploying minimal ADK agent to Agent Runtime...\")\n", + "remote_agent = client.agent_engines.create(\n", + " agent=adk_app,\n", + " config=dict(\n", + " display_name=\"minimal_adk_agent\",\n", + " staging_bucket=STAGING_BUCKET,\n", + " requirements=[\n", + " \"google-cloud-aiplatform[agent_engines,adk]\",\n", + " \"google-genai\",\n", + " \"google-adk[all]==2.1.0\"\n", + " ],\n", + " extra_packages=[\"./minimal_agent.py\"]\n", + " )\n", + ")\n", + "\n", + "print(f\"\\n✅ Agent deployed successfully!\")\n", + "print(f\"Resource Name: {remote_agent.api_resource.name}\")" + ] + }, + { + "cell_type": "markdown", + "id": "e277933c", + "metadata": { + "id": "e277933c" + }, + "source": [ + "## Step 5: Verify and test the live agent\n", + "\n", + "Query your live deployed agent over streaming to verify it returns SQL wrapped inside markdown code blocks.\n", + "\n", + "Once verified, copy the `Resource Name` printed in the previous step and paste it into the `resource_name` configuration field of your EvalBench `model_config` YAML file (e.g. `datasets/model_configs/agent_runtime.yaml`) to run evaluations!" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0164110c", + "metadata": { + "id": "0164110c" + }, + "outputs": [], + "source": [ + "# Test query the remote app client over streaming.\n", + "print(\"Querying remote agent endpoint...\")\n", + "response_stream = remote_agent.stream_query(\n", + " message=\"How many schools are in California?\",\n", + " user_id=\"evalbench_test\"\n", + ")\n", + "\n", + "print(\"\\nStreaming response chunks:\")\n", + "for chunk in response_stream:\n", + " print(chunk, end=\"\")" + ] + }, + { + "cell_type": "markdown", + "id": "66a03a42", + "metadata": { + "id": "66a03a42" + }, + "source": [ + "## Step 6: Generate temporary config files for evaluation\n", + "\n", + "To run evaluations without modifying the default files in the repository, we create temporary copies of both the model configuration and the run configuration.\n", + "\n", + "The code block below:\n", + "1. Loads the default `agent_runtime.yaml` model config, updates the `resource_name` with your deployed agent ID, and writes it to a new temporary file `agent_runtime_temp.yaml`.\n", + "2. Loads the default run configuration, updates its `model_config` path to point to the temporary model config, and writes it to a new temporary file `test_agent_runtime_run_config_temp.yaml`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "99f82cfd", + "metadata": { + "id": "99f82cfd" + }, + "outputs": [], + "source": [ + "import os\n", + "import json\n", + "import yaml\n", + "\n", + "# Using absolute paths to ensure files are located correctly\n", + "base_path = \"/content/evalbench\"\n", + "model_config_path = (\n", + " os.path.join(base_path, \"datasets/model_configs/agent_runtime.yaml\")\n", + ")\n", + "prompts_path = os.path.join(base_path, \"datasets/bird/prompts.json\")\n", + "\n", + "# Temporary file paths\n", + "temp_model_config_path = (\n", + " os.path.join(base_path, \"datasets/model_configs/agent_runtime_temp.yaml\")\n", + ")\n", + "temp_prompts_path = (\n", + " os.path.join(base_path, \"datasets/bird/prompts_california_temp.json\")\n", + ")\n", + "temp_run_config_path = (\n", + " os.path.join(\n", + " base_path, \n", + " \"datasets/bird/test_agent_runtime_run_config_temp.yaml\"\n", + " )\n", + ")\n", + "\n", + "# Load, update, and write temporary model config.\n", + "with open(model_config_path, \"r\") as f:\n", + " model_data = yaml.safe_load(f)\n", + "\n", + "# Update the resource name with your deployed agent\n", + "model_data[\"resource_name\"] = remote_agent.api_resource.name\n", + "\n", + "with open(temp_model_config_path, \"w\") as f:\n", + " yaml.dump(model_data, f)\n", + "print(f\"Created temporary model config: {temp_model_config_path}\")\n", + "\n", + "# Slice prompts.json to extract California Schools prompts.\n", + "with open(prompts_path, \"r\") as f:\n", + " all_prompts = json.load(f)\n", + "\n", + "california_prompts = [\n", + " p for p in all_prompts if p.get(\"db_id\") == \"california_schools\"\n", + "]\n", + "sliced_prompts = california_prompts[:3]\n", + "\n", + "with open(temp_prompts_path, \"w\") as f:\n", + " json.dump(sliced_prompts, f, indent=2)\n", + "print(f\"Created temporary prompts slice: {temp_prompts_path}\")\n", + "\n", + "# Generate the temporary run configuration with absolute paths.\n", + "run_config_content = f\"\"\"\n", + "dataset_config: {temp_prompts_path}\n", + "\n", + "database_configs:\n", + " - {os.path.join(base_path, 'datasets/bird/db_configs/sqlite.yaml')}\n", + "\n", + "dialects:\n", + " - sqlite\n", + "query_types:\n", + " - dql\n", + "dataset_format: bird-standard-format\n", + "\n", + "model_config: {temp_model_config_path}\n", + "prompt_generator: 'SQLGenBasePromptGenerator'\n", + "\n", + "scorers:\n", + " python_scorer:\n", + " script_path: '{os.path.join(base_path, 'evalbench/scorers/judges/sqlite_exact_match.py')}'\n", + " scorer_name: 'sqlite_exact_match'\n", + "\"\"\"\n", + "\n", + "with open(temp_run_config_path, \"w\") as f:\n", + " f.write(run_config_content.strip())\n", + "print(f\"Created temporary run config: {temp_run_config_path}\")" + ] + }, + { + "cell_type": "markdown", + "id": "f71131b9", + "metadata": { + "id": "f71131b9" + }, + "source": [ + "## Step 7: Run EvalBench evaluation pipeline\n", + "\n", + "We can now trigger the EvalBench evaluation run using the updated configuration. This executes query generation on your deployed agent and performs BigQuery execution and correctness scoring." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d88ccc37", + "metadata": { + "id": "d88ccc37" + }, + "outputs": [], + "source": [ + "import os\n", + "# Change to the evalbench root directory\n", + "%cd /content/evalbench\n", + "\n", + "# Run the simplified execution command without the Python protobuf override\n", + "! .venv/bin/python3 evalbench/evalbench.py \\\n", + " --experiment_config=\"datasets/bird/test_agent_runtime_run_config_temp.yaml\"" + ] + } + ], + "metadata": { + "colab": { + "provenance": [] + }, + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.14" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/evalbench/generators/models/__init__.py b/evalbench/generators/models/__init__.py index 93485533..15b9841d 100644 --- a/evalbench/generators/models/__init__.py +++ b/evalbench/generators/models/__init__.py @@ -14,6 +14,7 @@ from .agy_cli import AgyCliGenerator from .mcp_tools import McpToolsGenerator from .noop_agent import NoopAgentGenerator +from .agent_runtime import AgentRuntimeGenerator from util.config import load_yaml_config @@ -40,6 +41,7 @@ def get_generator(global_models, model_config_path: str, db: DB = None): "agy_cli": lambda: AgyCliGenerator(config), "mcp_tools": lambda: McpToolsGenerator(config), "noop_agent": lambda: NoopAgentGenerator(config), + "agent_runtime": lambda: AgentRuntimeGenerator(config), } generator = config["generator"] if generator not in generators: diff --git a/evalbench/generators/models/agent_runtime.py b/evalbench/generators/models/agent_runtime.py new file mode 100644 index 00000000..130adc9f --- /dev/null +++ b/evalbench/generators/models/agent_runtime.py @@ -0,0 +1,149 @@ +"""Agent Runtime (Gemini Enterprise Agent Platform) generator for EvalBench. + +The remote Agent Engine instance is expected to return the generated query +in one of the following formats: + +1. A raw SQL query string, optionally wrapped in markdown code blocks: + SELECT * FROM table; + +2. A JSON envelope wrapped in a ```json block containing a "sql" key: + ```json + { + "sql": "SELECT * FROM table;", + "explain": "Optional explanation of the query..." + } + ``` +""" + +import json +import logging +import os +import re + +from google.cloud.aiplatform_v1.types import ( + reasoning_engine_execution_service as aip_types, +) +import vertexai +from vertexai import agent_engines + +from generators.models.generator import QueryGenerator + + +def _parse_stream_response(response) -> str: + """Accumulates text chunks from the streaming response.""" + complete_text = "" + for chunk in response: + data = getattr(chunk, "data", b"") + if not data: + continue + try: + utf8_data = data.decode("utf-8") + for line in utf8_data.split("\n"): + if not line.strip(): + continue + parsed_json = json.loads(line) + parts = parsed_json.get("content", {}).get("parts", []) + for part in parts: + text = part.get("text") + if text: + complete_text += text + except (UnicodeDecodeError, json.JSONDecodeError): + # Fallback if chunk payload is not valid JSON. + continue + return complete_text + + +def _extract_sql(text: str) -> str: + """Extracts SQL query from JSON envelope or markdown blocks in text.""" + # Try to locate and extract a JSON envelope anywhere in the text first. + json_match = re.search(r"```json\s*(\{.*?\})\s*```", text, re.DOTALL) + if json_match: + try: + data = json.loads(json_match.group(1).strip()) + if isinstance(data, dict) and "sql" in data: + return str(data["sql"]) + except json.JSONDecodeError: + pass + + # Fallback: Split by markdown blocks and take the first segment. + if "```" in text: + candidate = text.split("```")[0].strip() + if candidate: + return candidate + + return text.strip() + + +class AgentRuntimeGenerator(QueryGenerator): + """Generates SQL queries using Agent Runtime.""" + + def __init__(self, querygenerator_config): + super().__init__(querygenerator_config) + self.name = "agent_runtime" + self.resource_name = ( + querygenerator_config.get("resource_name") + or os.environ.get("AGENT_ENGINE_RESOURCE") + ) + if not self.resource_name: + raise ValueError( + "AgentRuntimeGenerator requires `resource_name` in model " + "config YAML or AGENT_ENGINE_RESOURCE env variable." + ) + + project_id = ( + querygenerator_config.get("gcp_project_id") + or os.environ.get("GOOGLE_CLOUD_PROJECT") + ) + if not project_id: + raise ValueError( + "AgentRuntimeGenerator requires `gcp_project_id` in model " + "config YAML or GOOGLE_CLOUD_PROJECT env variable." + ) + + location = ( + querygenerator_config.get("gcp_region") + or os.environ.get("GOOGLE_CLOUD_REGION") + or os.environ.get("GOOGLE_CLOUD_LOCATION") + ) + if not location: + raise ValueError( + "AgentRuntimeGenerator requires `gcp_region` in model config " + "YAML or GOOGLE_CLOUD_REGION/LOCATION env vars." + ) + logging.info( + "Initializing Vertex AI (Project: %s, Location: %s)", + project_id, + location, + ) + vertexai.init(project=project_id, location=location) + + logging.info( + "Connecting to live Agent Runtime: " + f"{self.resource_name}" + ) + self.remote_app = agent_engines.AgentEngine(self.resource_name) + + def generate_internal(self, prompt: str) -> str: + """Queries remote agent endpoint and extracts SQL reply string.""" + try: + # Query the agent using the raw streaming endpoint. + client = self.remote_app.execution_api_client + response = client.stream_query_reasoning_engine( + request=aip_types.StreamQueryReasoningEngineRequest( + name=self.resource_name, + input={ + "message": prompt, + "user_id": "evalbench_user", + }, + class_method="stream_query", + ), + ) + + complete_text = _parse_stream_response(response) + return _extract_sql(complete_text) + + except Exception as e: + logging.exception( + f"Error querying remote Agent Runtime: {e}" + ) + return "" diff --git a/evalbench/test/agent_runtime_test.py b/evalbench/test/agent_runtime_test.py new file mode 100644 index 00000000..2f14d69f --- /dev/null +++ b/evalbench/test/agent_runtime_test.py @@ -0,0 +1,139 @@ +import os +import sys +from unittest.mock import MagicMock, patch + +# Add parent directory to path so we can import generators. +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from generators.models.agent_runtime import ( + _extract_sql, + _parse_stream_response, + AgentRuntimeGenerator, +) + + +# Unit tests for helper functions. +def test_parse_stream_response_accumulates_text(): + # Mock gRPC chunks returned by the stream client. + chunk1 = MagicMock() + chunk1.data = b'{"content": {"parts": [{"text": "SELECT "}]}}' + + chunk2 = MagicMock() + chunk2.data = ( + b'{"content": {"parts": [{"text": "1 FROM "}]}}\n' + b'{"content": {"parts": [{"text": "table;"}]}}' + ) + + response_stream = [chunk1, chunk2] + + result = _parse_stream_response(response_stream) + assert result == "SELECT 1 FROM table;" + + +def test_parse_stream_response_handles_invalid_chunks(): + # Test that decoding or JSON errors in one chunk don't crash the loop. + chunk_valid = MagicMock() + chunk_valid.data = b'{"content": {"parts": [{"text": "SELECT * "}]}}' + + chunk_invalid = MagicMock() + chunk_invalid.data = b'invalid-non-json-bytes' + + response_stream = [chunk_valid, chunk_invalid] + result = _parse_stream_response(response_stream) + assert result == "SELECT * " + + +def test_extract_sql_from_clean_json(): + text = """```json +{ + "explain": "Retrieves values.", + "sql": "SELECT * FROM users;" +} +```""" + assert _extract_sql(text) == "SELECT * FROM users;" + + +def test_extract_sql_from_mixed_response(): + # Mix of raw SQL followed by markdown block. + text = """SELECT * FROM users; +```json +{ + "explain": "Retrieves values.", + "sql": "SELECT * FROM users;" +} +```""" + assert _extract_sql(text) == "SELECT * FROM users;" + + +def test_extract_sql_from_plain_markdown_fallback(): + # Fallback when no JSON envelope can be parsed, splits by block wrapper. + text = """SELECT * FROM users; +``` +Explanation of query below... +```""" + assert _extract_sql(text) == "SELECT * FROM users;" + + +def test_extract_sql_raw_fallback(): + # Fallback to trim when no markdown blocks are present. + text = " SELECT * FROM users; " + assert _extract_sql(text) == "SELECT * FROM users;" + + +# Unit tests for Generator class. +@patch("generators.models.agent_runtime.vertexai.init") +@patch("generators.models.agent_runtime.agent_engines.AgentEngine") +def test_generator_initialization(mock_agent_engine, mock_vertexai_init): + config = { + "resource_name": "projects/p/locations/l/reasoningEngines/r", + "gcp_project_id": "test-project", + "gcp_region": "us-central1" + } + + generator = AgentRuntimeGenerator(config) + + # Assert Vertex AI is initialized correctly. + mock_vertexai_init.assert_called_once_with( + project="test-project", location="us-central1" + ) + + # Assert AgentEngine is instantiated with correct resource ID. + mock_agent_engine.assert_called_once_with( + "projects/p/locations/l/reasoningEngines/r" + ) + assert generator.remote_app == mock_agent_engine.return_value + assert generator.name == "agent_runtime" + + +@patch("generators.models.agent_runtime.vertexai.init") +@patch("generators.models.agent_runtime.agent_engines.AgentEngine") +def test_generate_internal_queries_stream_api(mock_agent_engine, mock_vertexai_init): + config = { + "resource_name": "projects/p/locations/l/reasoningEngines/r", + "gcp_project_id": "test-project", + "gcp_region": "us-central1" + } + + # Setup mock clients and streaming responses. + mock_client = MagicMock() + mock_agent_engine.return_value.execution_api_client = mock_client + + mock_chunk = MagicMock() + mock_chunk.data = b'{"content": {"parts": [{"text": "SELECT 5;"}]}}' + mock_client.stream_query_reasoning_engine.return_value = [mock_chunk] + + generator = AgentRuntimeGenerator(config) + sql = generator.generate_internal("run query 5") + + # Verify stream query is called with correct name, prompts, method. + mock_client.stream_query_reasoning_engine.assert_called_once() + call_args = ( + mock_client.stream_query_reasoning_engine.call_args[1]["request"] + ) + + assert call_args.name == "projects/p/locations/l/reasoningEngines/r" + assert call_args.input["message"] == "run query 5" + assert call_args.class_method == "stream_query" + + # Verify query output was parsed and returned. + assert sql == "SELECT 5;" diff --git a/evalbench/work/sqlexecwork.py b/evalbench/work/sqlexecwork.py index 183e3c67..05e4a5d8 100644 --- a/evalbench/work/sqlexecwork.py +++ b/evalbench/work/sqlexecwork.py @@ -162,7 +162,7 @@ def _evaluate_execution_results( def _sanitize_sql(self): if ( self.experiment_config["prompt_generator"] == "NOOPGenerator" - and self.experiment_config["dialect"] != "googlesql" + and self.experiment_config.get("dialect") != "googlesql" ): self.eval_result["sanitized_sql"] = self.eval_result[ "generated_sql" From 9a73b0d611cebb7ef47f2ec93ea4b4aee6cb72b7 Mon Sep 17 00:00:00 2001 From: Ariel Jassan Date: Wed, 29 Jul 2026 00:35:07 +0300 Subject: [PATCH 2/5] Remove unnecessary blank lines in tests --- evalbench/test/agent_runtime_test.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/evalbench/test/agent_runtime_test.py b/evalbench/test/agent_runtime_test.py index 2f14d69f..ca99eb2c 100644 --- a/evalbench/test/agent_runtime_test.py +++ b/evalbench/test/agent_runtime_test.py @@ -17,7 +17,7 @@ def test_parse_stream_response_accumulates_text(): # Mock gRPC chunks returned by the stream client. chunk1 = MagicMock() chunk1.data = b'{"content": {"parts": [{"text": "SELECT "}]}}' - + chunk2 = MagicMock() chunk2.data = ( b'{"content": {"parts": [{"text": "1 FROM "}]}}\n' @@ -34,10 +34,10 @@ def test_parse_stream_response_handles_invalid_chunks(): # Test that decoding or JSON errors in one chunk don't crash the loop. chunk_valid = MagicMock() chunk_valid.data = b'{"content": {"parts": [{"text": "SELECT * "}]}}' - + chunk_invalid = MagicMock() chunk_invalid.data = b'invalid-non-json-bytes' - + response_stream = [chunk_valid, chunk_invalid] result = _parse_stream_response(response_stream) assert result == "SELECT * " @@ -113,11 +113,11 @@ def test_generate_internal_queries_stream_api(mock_agent_engine, mock_vertexai_i "gcp_project_id": "test-project", "gcp_region": "us-central1" } - + # Setup mock clients and streaming responses. mock_client = MagicMock() mock_agent_engine.return_value.execution_api_client = mock_client - + mock_chunk = MagicMock() mock_chunk.data = b'{"content": {"parts": [{"text": "SELECT 5;"}]}}' mock_client.stream_query_reasoning_engine.return_value = [mock_chunk] @@ -130,7 +130,7 @@ def test_generate_internal_queries_stream_api(mock_agent_engine, mock_vertexai_i call_args = ( mock_client.stream_query_reasoning_engine.call_args[1]["request"] ) - + assert call_args.name == "projects/p/locations/l/reasoningEngines/r" assert call_args.input["message"] == "run query 5" assert call_args.class_method == "stream_query" From 41827b660d04c221ba5ed29655e39796e50d455a Mon Sep 17 00:00:00 2001 From: Ariel Jassan Date: Wed, 29 Jul 2026 00:39:35 +0300 Subject: [PATCH 3/5] Fix formatting issue in agent_runtime_test.py --- evalbench/test/agent_runtime_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/evalbench/test/agent_runtime_test.py b/evalbench/test/agent_runtime_test.py index ca99eb2c..d1c41053 100644 --- a/evalbench/test/agent_runtime_test.py +++ b/evalbench/test/agent_runtime_test.py @@ -25,7 +25,7 @@ def test_parse_stream_response_accumulates_text(): ) response_stream = [chunk1, chunk2] - + result = _parse_stream_response(response_stream) assert result == "SELECT 1 FROM table;" From 78a03fbc4c5656e95fa7c4e5f4b1e7b67f249b8f Mon Sep 17 00:00:00 2001 From: Ariel Jassan Date: Thu, 30 Jul 2026 13:36:09 +0000 Subject: [PATCH 4/5] refactor(evalbench): address PR feedback on Agent Runtime generator - Fix SQL extractor logic to correctly parse queries inside markdown fences - Add robust unit tests covering different fenced and raw outputs - Use !ENV interpolation for project/resource names in configs - Align environment variable helpers and fix rate limiting propagation - Declare google-cloud-aiplatform dependency and fix KeyError in interactsqlexecwork.py --- .../bird/test_agent_runtime_run_config.yaml | 4 +- datasets/model_configs/agent_runtime.yaml | 8 ++-- evalbench/generators/models/agent_runtime.py | 44 ++++++++----------- evalbench/test/agent_runtime_test.py | 36 ++++++++++++--- evalbench/work/interactsqlexecwork.py | 2 +- pyproject.toml | 1 + 6 files changed, 57 insertions(+), 38 deletions(-) diff --git a/datasets/bird/test_agent_runtime_run_config.yaml b/datasets/bird/test_agent_runtime_run_config.yaml index f59ff76c..d29767b9 100644 --- a/datasets/bird/test_agent_runtime_run_config.yaml +++ b/datasets/bird/test_agent_runtime_run_config.yaml @@ -1,7 +1,7 @@ ############################################################ ### Dataset / Eval Items (Vertex AI Multi-Agent Test) ############################################################ -dataset_config: datasets/bird/prompts_california.json +dataset_config: datasets/bird/prompts.json databases: - california_schools @@ -37,5 +37,5 @@ scorers: reporting: bigquery: dataset_location: "US" - gcp_project_id: "arielj-argolis-1" + gcp_project_id: !ENV ${EVAL_GCP_PROJECT_ID} csv: {} diff --git a/datasets/model_configs/agent_runtime.yaml b/datasets/model_configs/agent_runtime.yaml index 2f944f0c..1b699640 100644 --- a/datasets/model_configs/agent_runtime.yaml +++ b/datasets/model_configs/agent_runtime.yaml @@ -3,11 +3,11 @@ generator: agent_runtime # Explicit resource URI string (e.g. "projects/12345/locations/us-central1/reasoningEngines/67890") -# Leave empty string "" to dynamically read MULTI_AGENT_ENGINE_RESOURCE or AGENT_ENGINE_RESOURCE from env. -resource_name: "projects/743100398377/locations/us-central1/reasoningEngines/7921856483556327424" +# Leave empty string "" to dynamically read AGENT_ENGINE_RESOURCE from env. +resource_name: !ENV ${AGENT_ENGINE_RESOURCE:""} -gcp_project_id: "arielj-argolis-1" -gcp_region: "us-central1" +gcp_project_id: !ENV ${EVAL_GCP_PROJECT_ID} +gcp_region: !ENV ${EVAL_GCP_PROJECT_REGION:us-central1} # Rate limiting and retry mechanics execs_per_minute: 10 diff --git a/evalbench/generators/models/agent_runtime.py b/evalbench/generators/models/agent_runtime.py index 130adc9f..9ab6733b 100644 --- a/evalbench/generators/models/agent_runtime.py +++ b/evalbench/generators/models/agent_runtime.py @@ -20,6 +20,7 @@ import os import re +from google.api_core.exceptions import ResourceExhausted from google.cloud.aiplatform_v1.types import ( reasoning_engine_execution_service as aip_types, ) @@ -27,6 +28,9 @@ from vertexai import agent_engines from generators.models.generator import QueryGenerator +from util.gcp import get_gcp_project, get_gcp_region +from util.rate_limit import ResourceExhaustedError +from util.sanitizer import sanitize_sql def _parse_stream_response(response) -> str: @@ -55,23 +59,25 @@ def _parse_stream_response(response) -> str: def _extract_sql(text: str) -> str: """Extracts SQL query from JSON envelope or markdown blocks in text.""" - # Try to locate and extract a JSON envelope anywhere in the text first. + # 1. Try to locate and extract a JSON envelope anywhere in the text first. json_match = re.search(r"```json\s*(\{.*?\})\s*```", text, re.DOTALL) if json_match: try: data = json.loads(json_match.group(1).strip()) if isinstance(data, dict) and "sql" in data: - return str(data["sql"]) + return sanitize_sql(str(data["sql"])) except json.JSONDecodeError: pass - # Fallback: Split by markdown blocks and take the first segment. - if "```" in text: - candidate = text.split("```")[0].strip() - if candidate: - return candidate + # 2. Try to locate any markdown code block (e.g. ```sql or ```) + sql_match = ( + re.search(r"```(?:sql)?\s*(.*?)\s*```", text, re.DOTALL | re.IGNORECASE) + ) + if sql_match: + return sanitize_sql(sql_match.group(1).strip()) - return text.strip() + # 3. Fallback: Treat as raw SQL query + return sanitize_sql(text) class AgentRuntimeGenerator(QueryGenerator): @@ -91,25 +97,9 @@ def __init__(self, querygenerator_config): ) project_id = ( - querygenerator_config.get("gcp_project_id") - or os.environ.get("GOOGLE_CLOUD_PROJECT") + get_gcp_project(querygenerator_config.get("gcp_project_id")) ) - if not project_id: - raise ValueError( - "AgentRuntimeGenerator requires `gcp_project_id` in model " - "config YAML or GOOGLE_CLOUD_PROJECT env variable." - ) - - location = ( - querygenerator_config.get("gcp_region") - or os.environ.get("GOOGLE_CLOUD_REGION") - or os.environ.get("GOOGLE_CLOUD_LOCATION") - ) - if not location: - raise ValueError( - "AgentRuntimeGenerator requires `gcp_region` in model config " - "YAML or GOOGLE_CLOUD_REGION/LOCATION env vars." - ) + location = get_gcp_region(querygenerator_config.get("gcp_region")) logging.info( "Initializing Vertex AI (Project: %s, Location: %s)", project_id, @@ -142,6 +132,8 @@ def generate_internal(self, prompt: str) -> str: complete_text = _parse_stream_response(response) return _extract_sql(complete_text) + except ResourceExhausted as e: + raise ResourceExhaustedError(e) except Exception as e: logging.exception( f"Error querying remote Agent Runtime: {e}" diff --git a/evalbench/test/agent_runtime_test.py b/evalbench/test/agent_runtime_test.py index d1c41053..67744d9b 100644 --- a/evalbench/test/agent_runtime_test.py +++ b/evalbench/test/agent_runtime_test.py @@ -65,15 +65,41 @@ def test_extract_sql_from_mixed_response(): assert _extract_sql(text) == "SELECT * FROM users;" -def test_extract_sql_from_plain_markdown_fallback(): - # Fallback when no JSON envelope can be parsed, splits by block wrapper. - text = """SELECT * FROM users; -``` -Explanation of query below... +def test_extract_sql_from_plain_markdown_fence(): + # Content inside plain fence is extracted. + text = """``` +SELECT * FROM users; ```""" assert _extract_sql(text) == "SELECT * FROM users;" +def test_extract_sql_from_sql_markdown_fence(): + # Case 1: Simple fenced block without any prose + text = """```sql +SELECT COUNT(*) FROM schools; +```""" + assert _extract_sql(text) == "SELECT COUNT(*) FROM schools;" + + +def test_extract_sql_from_fenced_block_with_preamble(): + # Case 2: Fenced block preceded by prose preamble + text = """Here is the query: +```sql +SELECT COUNT(*) FROM schools; +```""" + assert _extract_sql(text) == "SELECT COUNT(*) FROM schools;" + + +def test_extract_sql_from_fenced_block_with_preamble_and_postamble(): + # Case 3: Fenced block with preamble and post-amble text + text = """Sure! +```sql +SELECT 1; +``` +Hope that helps.""" + assert _extract_sql(text) == "SELECT 1;" + + def test_extract_sql_raw_fallback(): # Fallback to trim when no markdown blocks are present. text = " SELECT * FROM users; " diff --git a/evalbench/work/interactsqlexecwork.py b/evalbench/work/interactsqlexecwork.py index 8d5e251c..427717bf 100644 --- a/evalbench/work/interactsqlexecwork.py +++ b/evalbench/work/interactsqlexecwork.py @@ -110,7 +110,7 @@ def _evaluate_execution_results( def _sanitize_sql(self): if ( self.experiment_config["prompt_generator"] == "NOOPGenerator" - and self.experiment_config["dialect"] != "googlesql" + and self.experiment_config.get("dialect") != "googlesql" ): self.eval_result["sanitized_sql"] = self.eval_result["generated_sql"] else: diff --git a/pyproject.toml b/pyproject.toml index ea389209..c0eb0033 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,6 +22,7 @@ dependencies = [ "cloud-sql-python-connector[pytds]", "sqlalchemy-pytds", "google-genai", + "google-cloud-aiplatform", "google-cloud-secret-manager", "google-cloud-bigtable", "google-cloud-spanner", From c5fc71fd1fc218aaf148897577a4f3afdbc0b02c Mon Sep 17 00:00:00 2001 From: Ariel Jassan Date: Thu, 6 Aug 2026 18:34:19 +0000 Subject: [PATCH 5/5] docs(examples): fix agent_runtime_deployment notebook imports, scorers, and environment setup - Replace private agentplatform imports with public vertexai native SDK APIs - Install pyaml_env globally in Colab and use it to resolve !ENV YAML tags during config parsing - Export EVAL_GCP_PROJECT_ID and EVAL_GCP_PROJECT_REGION to environment for subprocesses - Use built-in set_match and exact_match scorers instead of the missing sqlite_exact_match.py script - Enable csv reporting in the temporary run config so results are written to disk - Update agent instructions to output SQLite SQL instead of BigQuery SQL --- docs/examples/agent_runtime_deployment.ipynb | 928 ++++++++++--------- 1 file changed, 467 insertions(+), 461 deletions(-) diff --git a/docs/examples/agent_runtime_deployment.ipynb b/docs/examples/agent_runtime_deployment.ipynb index 8d722318..5dee2ef4 100644 --- a/docs/examples/agent_runtime_deployment.ipynb +++ b/docs/examples/agent_runtime_deployment.ipynb @@ -1,464 +1,470 @@ { - "cells": [ - { - "cell_type": "markdown", - "id": "5b5c8393", - "metadata": { - "id": "5b5c8393" - }, - "source": [ - "# Evaluating an ADK Agent to Agent Runtime (Gemini Enterprise Agent Platform)\n", - "\n", - "This notebook guides you through the process of:\n", - "1. Writing a minimal SQL-generating agent using **Agent Development Kit (ADK)**.\n", - "2. Deploying it to **Agent Runtime (Gemini Enterprise Agent Platform)** using the `AdkApp` template.\n", - "3. Querying the live endpoint to verify it streams valid SQL.\n", - "4. Evaluating the agent in Agent Runtime." - ] - }, - { - "cell_type": "markdown", - "id": "96e96c58", - "metadata": { - "id": "96e96c58" - }, - "source": [ - "## Step 1: Repository Setup" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "3a2541a2", - "metadata": { - "id": "3a2541a2" - }, - "outputs": [], - "source": [ - "# Install/Upgrade the Google Cloud AI Platform SDK, ADK, and authentication libraries\n", - "!pip install --quiet --upgrade google-cloud-aiplatform google-adk google-auth==2.49.0" - ] - }, - { - "cell_type": "markdown", - "id": "4a1023be", - "metadata": {}, - "source": [ - "If running in a hosted environment like Colab Enterprise, we clone the `evalbench` repository to load evaluation datasets, model configurations, and pipeline execution shell scripts." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "ff6a777c", - "metadata": {}, - "outputs": [], - "source": [ - "!git clone https://github.com/GoogleCloudPlatform/evalbench.git" - ] - }, - { - "cell_type": "markdown", - "id": "8fa1abd1", - "metadata": {}, - "source": [ - "We install uv to manage python packages, sync the virtual environment, install the mcp library package, and compile the protobuf message files." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "94d69b26", - "metadata": { - "id": "94d69b26" - }, - "outputs": [], - "source": [ - "# Re-install dependencies and compile protobufs\n", - "!pip install --quiet uv\n", - "!cd /content/evalbench && uv sync\n", - "!cd /content/evalbench && .venv/bin/python3 -m grpc_tools.protoc \\\n", - " --proto_path=evalbench/evalproto \\\n", - " --python_out=evalbench/evalproto \\\n", - " --pyi_out=evalbench/evalproto \\\n", - " --grpc_python_out=evalbench/evalproto \\\n", - " --experimental_editions evalbench/evalproto/*.proto\n", - "\n", - "import os\n", - "os.chdir(\"/content/evalbench/docs/examples\")\n", - "print(f\"✅ Working directory set to: {os.getcwd()}\")" - ] - }, - { - "cell_type": "markdown", - "id": "476091be", - "metadata": { - "id": "476091be" - }, - "source": [ - "Download the BIRD Dataset and Database Connections:\n", - "This script downloads the natural language questions, database schemas, and SQLite database connection files required to evaluate the SQL generator." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "7994b603", - "metadata": { - "id": "7994b603" - }, - "outputs": [], - "source": [ - "# Download the BIRD dataset and database connections\n", - "!cd ../.. && bash datasets/bird/download_dataset.sh" - ] - }, - { - "cell_type": "markdown", - "id": "25850460", - "metadata": { - "id": "25850460" - }, - "source": [ - "## Step 2: Define the ADK Agent\n", - "\n", - "We write the ADK agent definition code to a local file named `minimal_agent.py` using standard ADK syntax. We specify `gemini-2.5-flash` as the model name and provide custom instruction formatting rules." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "8481f972", - "metadata": { - "id": "8481f972" - }, - "outputs": [], - "source": [ - "%%writefile minimal_agent.py\n", - "import google.adk as adk\n", - "\n", - "def create_agent() -> adk.Agent:\n", - " \"\"\"Creates a minimal SQL-generating ADK agent.\"\"\"\n", - " return adk.Agent(\n", - " name=\"minimal_adk_sql_agent\",\n", - " model=\"gemini-2.5-flash\",\n", - " instruction=(\n", - " \"You are a SQL generator. Generate a BigQuery SQL query based on \"\n", - " \"the user prompt. Wrap your final generated SQL query inside a \"\n", - " \"```sql code block.\"\n", - " )\n", - " )" - ] - }, - { - "cell_type": "markdown", - "id": "d153df7b", - "metadata": { - "id": "d153df7b" - }, - "source": [ - "## Step 3: Configure and initialize Agent Platform client\n", - "\n", - "Set your target Google Cloud Project and region, and initialize the SDK." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "57e7af39", - "metadata": { - "id": "57e7af39" - }, - "outputs": [], - "source": [ - "from google.colab import auth\n", - "auth.authenticate_user()" - ] - }, - { - "cell_type": "markdown", - "id": "fd22b631", - "metadata": { - "id": "fd22b631" - }, - "source": [ - "Set your target Google Cloud Project and region, and initialize the SDK." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "8adb6260", - "metadata": { - "id": "8adb6260" - }, - "outputs": [], - "source": [ - "import os\n", - "import google.cloud.aiplatform as aiplatform\n", - "import agentplatform\n", - "import google.auth\n", - "\n", - "# Edit these variables with your Google Cloud details.\n", - "PROJECT_ID = \"YOUR-PROJECT-ID\"\n", - "LOCATION = \"us-central1\"\n", - "STAGING_BUCKET = f\"gs://{PROJECT_ID}-agents\"\n", - "\n", - "# Get the Colab authenticated credentials.\n", - "print(f\"Initializing global context and client...\")\n", - "credentials, _ = google.auth.default()\n", - "\n", - "# Initialize global context and client with explicit credentials.\n", - "aiplatform.init(project=PROJECT_ID, location=LOCATION, credentials=credentials)\n", - "client = (\n", - " agentplatform.Client(\n", - " project=PROJECT_ID, \n", - " location=LOCATION, \n", - " credentials=credentials\n", - " )\n", - ")" - ] - }, - { - "cell_type": "markdown", - "id": "80748249", - "metadata": { - "id": "80748249" - }, - "source": [ - "## Step 4: Deploy to Agent Runtime using the AdkApp Template\n", - "\n", - "We load the local ADK agent and wrap it inside the `AdkApp` class. Then, we call `client.agent_engines.create` to deploy the agent.\n", - "\n", - "Since our agent is self-contained in `minimal_agent.py`, we only need to include `./minimal_agent.py` in the `extra_packages` deployment configuration." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "fe59c109", - "metadata": { - "id": "fe59c109" - }, - "outputs": [], - "source": [ - "from agentplatform.agent_engines.templates.adk import AdkApp\n", - "\n", - "import minimal_agent\n", - "\n", - "\n", - "# Instantiate the local ADK agent and wrap it in the template.\n", - "local_agent = minimal_agent.create_agent()\n", - "adk_app = AdkApp(agent=local_agent, app_name=\"minimal-adk-sql-app\")\n", - "\n", - "print(\"Deploying minimal ADK agent to Agent Runtime...\")\n", - "remote_agent = client.agent_engines.create(\n", - " agent=adk_app,\n", - " config=dict(\n", - " display_name=\"minimal_adk_agent\",\n", - " staging_bucket=STAGING_BUCKET,\n", - " requirements=[\n", - " \"google-cloud-aiplatform[agent_engines,adk]\",\n", - " \"google-genai\",\n", - " \"google-adk[all]==2.1.0\"\n", - " ],\n", - " extra_packages=[\"./minimal_agent.py\"]\n", - " )\n", - ")\n", - "\n", - "print(f\"\\n✅ Agent deployed successfully!\")\n", - "print(f\"Resource Name: {remote_agent.api_resource.name}\")" - ] - }, - { - "cell_type": "markdown", - "id": "e277933c", - "metadata": { - "id": "e277933c" - }, - "source": [ - "## Step 5: Verify and test the live agent\n", - "\n", - "Query your live deployed agent over streaming to verify it returns SQL wrapped inside markdown code blocks.\n", - "\n", - "Once verified, copy the `Resource Name` printed in the previous step and paste it into the `resource_name` configuration field of your EvalBench `model_config` YAML file (e.g. `datasets/model_configs/agent_runtime.yaml`) to run evaluations!" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "0164110c", - "metadata": { - "id": "0164110c" - }, - "outputs": [], - "source": [ - "# Test query the remote app client over streaming.\n", - "print(\"Querying remote agent endpoint...\")\n", - "response_stream = remote_agent.stream_query(\n", - " message=\"How many schools are in California?\",\n", - " user_id=\"evalbench_test\"\n", - ")\n", - "\n", - "print(\"\\nStreaming response chunks:\")\n", - "for chunk in response_stream:\n", - " print(chunk, end=\"\")" - ] - }, - { - "cell_type": "markdown", - "id": "66a03a42", - "metadata": { - "id": "66a03a42" - }, - "source": [ - "## Step 6: Generate temporary config files for evaluation\n", - "\n", - "To run evaluations without modifying the default files in the repository, we create temporary copies of both the model configuration and the run configuration.\n", - "\n", - "The code block below:\n", - "1. Loads the default `agent_runtime.yaml` model config, updates the `resource_name` with your deployed agent ID, and writes it to a new temporary file `agent_runtime_temp.yaml`.\n", - "2. Loads the default run configuration, updates its `model_config` path to point to the temporary model config, and writes it to a new temporary file `test_agent_runtime_run_config_temp.yaml`." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "99f82cfd", - "metadata": { - "id": "99f82cfd" - }, - "outputs": [], - "source": [ - "import os\n", - "import json\n", - "import yaml\n", - "\n", - "# Using absolute paths to ensure files are located correctly\n", - "base_path = \"/content/evalbench\"\n", - "model_config_path = (\n", - " os.path.join(base_path, \"datasets/model_configs/agent_runtime.yaml\")\n", - ")\n", - "prompts_path = os.path.join(base_path, \"datasets/bird/prompts.json\")\n", - "\n", - "# Temporary file paths\n", - "temp_model_config_path = (\n", - " os.path.join(base_path, \"datasets/model_configs/agent_runtime_temp.yaml\")\n", - ")\n", - "temp_prompts_path = (\n", - " os.path.join(base_path, \"datasets/bird/prompts_california_temp.json\")\n", - ")\n", - "temp_run_config_path = (\n", - " os.path.join(\n", - " base_path, \n", - " \"datasets/bird/test_agent_runtime_run_config_temp.yaml\"\n", - " )\n", - ")\n", - "\n", - "# Load, update, and write temporary model config.\n", - "with open(model_config_path, \"r\") as f:\n", - " model_data = yaml.safe_load(f)\n", - "\n", - "# Update the resource name with your deployed agent\n", - "model_data[\"resource_name\"] = remote_agent.api_resource.name\n", - "\n", - "with open(temp_model_config_path, \"w\") as f:\n", - " yaml.dump(model_data, f)\n", - "print(f\"Created temporary model config: {temp_model_config_path}\")\n", - "\n", - "# Slice prompts.json to extract California Schools prompts.\n", - "with open(prompts_path, \"r\") as f:\n", - " all_prompts = json.load(f)\n", - "\n", - "california_prompts = [\n", - " p for p in all_prompts if p.get(\"db_id\") == \"california_schools\"\n", - "]\n", - "sliced_prompts = california_prompts[:3]\n", - "\n", - "with open(temp_prompts_path, \"w\") as f:\n", - " json.dump(sliced_prompts, f, indent=2)\n", - "print(f\"Created temporary prompts slice: {temp_prompts_path}\")\n", - "\n", - "# Generate the temporary run configuration with absolute paths.\n", - "run_config_content = f\"\"\"\n", - "dataset_config: {temp_prompts_path}\n", - "\n", - "database_configs:\n", - " - {os.path.join(base_path, 'datasets/bird/db_configs/sqlite.yaml')}\n", - "\n", - "dialects:\n", - " - sqlite\n", - "query_types:\n", - " - dql\n", - "dataset_format: bird-standard-format\n", - "\n", - "model_config: {temp_model_config_path}\n", - "prompt_generator: 'SQLGenBasePromptGenerator'\n", - "\n", - "scorers:\n", - " python_scorer:\n", - " script_path: '{os.path.join(base_path, 'evalbench/scorers/judges/sqlite_exact_match.py')}'\n", - " scorer_name: 'sqlite_exact_match'\n", - "\"\"\"\n", - "\n", - "with open(temp_run_config_path, \"w\") as f:\n", - " f.write(run_config_content.strip())\n", - "print(f\"Created temporary run config: {temp_run_config_path}\")" - ] - }, - { - "cell_type": "markdown", - "id": "f71131b9", - "metadata": { - "id": "f71131b9" - }, - "source": [ - "## Step 7: Run EvalBench evaluation pipeline\n", - "\n", - "We can now trigger the EvalBench evaluation run using the updated configuration. This executes query generation on your deployed agent and performs BigQuery execution and correctness scoring." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "d88ccc37", - "metadata": { - "id": "d88ccc37" - }, - "outputs": [], - "source": [ - "import os\n", - "# Change to the evalbench root directory\n", - "%cd /content/evalbench\n", - "\n", - "# Run the simplified execution command without the Python protobuf override\n", - "! .venv/bin/python3 evalbench/evalbench.py \\\n", - " --experiment_config=\"datasets/bird/test_agent_runtime_run_config_temp.yaml\"" - ] - } - ], - "metadata": { - "colab": { - "provenance": [] - }, - "kernelspec": { - "display_name": ".venv", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.13.14" - } + "cells": [ + { + "id": "5b5c8393", + "cell_type": "markdown", + "source": [ + "# Evaluating an ADK Agent to Agent Runtime (Gemini Enterprise Agent Platform)\n", + "\n", + "This notebook guides you through the process of:\n", + "1. Writing a minimal SQL-generating agent using **Agent Development Kit (ADK)**.\n", + "2. Deploying it to **Agent Runtime (Gemini Enterprise Agent Platform)** using the `AdkApp` template.\n", + "3. Querying the live endpoint to verify it streams valid SQL.\n", + "4. Evaluating the agent in Agent Runtime." + ], + "metadata": { + "id": "5b5c8393" + }, + "execution_count": null }, - "nbformat": 4, - "nbformat_minor": 5 + { + "id": "96e96c58", + "cell_type": "markdown", + "source": [ + "## Step 1: Repository Setup" + ], + "metadata": { + "id": "96e96c58" + }, + "execution_count": null + }, + { + "id": "3a2541a2", + "cell_type": "code", + "source": [ + "# Install/Upgrade the Google Cloud AI Platform SDK, ADK, and authentication libraries\n", + "!pip install --quiet --upgrade google-cloud-aiplatform google-adk google-auth==2.49.0 pyaml_env" + ], + "metadata": { + "id": "3a2541a2" + }, + "execution_count": null + }, + { + "id": "4a1023be", + "cell_type": "markdown", + "source": [ + "If running in a hosted environment like Colab Enterprise, we clone the `evalbench` repository to load evaluation datasets, model configurations, and pipeline execution shell scripts." + ], + "metadata": {}, + "execution_count": null + }, + { + "id": "ff6a777c", + "cell_type": "code", + "source": [ + "!git clone https://github.com/GoogleCloudPlatform/evalbench.git" + ], + "metadata": {}, + "execution_count": null + }, + { + "id": "8fa1abd1", + "cell_type": "markdown", + "source": [ + "We install uv to manage python packages, sync the virtual environment, install the mcp library package, and compile the protobuf message files." + ], + "metadata": {}, + "execution_count": null + }, + { + "id": "94d69b26", + "cell_type": "code", + "source": [ + "# Re-install dependencies and compile protobufs\n", + "!pip install --quiet uv\n", + "!cd /content/evalbench \u0026\u0026 uv sync\n", + "!cd /content/evalbench \u0026\u0026 .venv/bin/python3 -m grpc_tools.protoc \\\n", + " --proto_path=evalbench/evalproto \\\n", + " --python_out=evalbench/evalproto \\\n", + " --pyi_out=evalbench/evalproto \\\n", + " --grpc_python_out=evalbench/evalproto \\\n", + " --experimental_editions evalbench/evalproto/*.proto\n", + "\n", + "import os\n", + "os.chdir(\"/content/evalbench/docs/examples\")\n", + "print(f\"✅ Working directory set to: {os.getcwd()}\")" + ], + "metadata": { + "id": "94d69b26" + }, + "execution_count": null + }, + { + "id": "476091be", + "cell_type": "markdown", + "source": [ + "Download the BIRD Dataset and Database Connections:\n", + "This script downloads the natural language questions, database schemas, and SQLite database connection files required to evaluate the SQL generator." + ], + "metadata": { + "id": "476091be" + }, + "execution_count": null + }, + { + "id": "7994b603", + "cell_type": "code", + "source": [ + "# Download the BIRD dataset and database connections\n", + "!cd ../.. \u0026\u0026 bash datasets/bird/download_dataset.sh" + ], + "metadata": { + "id": "7994b603" + }, + "execution_count": null + }, + { + "id": "25850460", + "cell_type": "markdown", + "source": [ + "## Step 2: Define the ADK Agent\n", + "\n", + "We write the ADK agent definition code to a local file named `minimal_agent.py` using standard ADK syntax. We specify `gemini-2.5-flash` as the model name and provide custom instruction formatting rules." + ], + "metadata": { + "id": "25850460" + }, + "execution_count": null + }, + { + "id": "8481f972", + "cell_type": "code", + "source": [ + "%%writefile minimal_agent.py\n", + "import google.adk as adk\n", + "\n", + "def create_agent() -\u003e adk.Agent:\n", + " \"\"\"Creates a minimal SQL-generating ADK agent.\"\"\"\n", + " return adk.Agent(\n", + " name=\"minimal_adk_sql_agent\",\n", + " model=\"gemini-2.5-flash\",\n", + " instruction=(\n", + " \"You are a SQL generator. Generate a SQLite SQL query based on \"\n", + " \"the user prompt. Wrap your final generated SQL query inside a \"\n", + " \"```sql code block.\"\n", + " )\n", + " )" + ], + "metadata": { + "id": "8481f972" + }, + "execution_count": null + }, + { + "id": "d153df7b", + "cell_type": "markdown", + "source": [ + "## Step 3: Configure and initialize Agent Platform client\n", + "\n", + "Set your target Google Cloud Project and region, and initialize the SDK." + ], + "metadata": { + "id": "d153df7b" + }, + "execution_count": null + }, + { + "id": "57e7af39", + "cell_type": "code", + "source": [ + "from google.colab import auth\n", + "auth.authenticate_user()" + ], + "metadata": { + "id": "57e7af39" + }, + "execution_count": null + }, + { + "id": "fd22b631", + "cell_type": "markdown", + "source": [ + "Set your target Google Cloud Project and region, and initialize the SDK." + ], + "metadata": { + "id": "fd22b631" + }, + "execution_count": null + }, + { + "id": "8adb6260", + "cell_type": "code", + "source": [ + "import os\n", + "import vertexai\n", + "import google.auth\n", + "\n", + "# Edit these variables with your Google Cloud details.\n", + "PROJECT_ID = \"YOUR-PROJECT-ID\"\n", + "LOCATION = \"us-central1\"\n", + "STAGING_BUCKET = f\"gs://{PROJECT_ID}-agents\"\n", + "\n", + "# Get the Colab authenticated credentials.\n", + "print(f\"Initializing global context and client...\")\n", + "credentials, _ = google.auth.default()\n", + "\n", + "# Initialize global context and client with explicit credentials.\n", + "vertexai.init(project=PROJECT_ID, location=LOCATION, credentials=credentials)\n", + "client = (\n", + " vertexai.Client(\n", + " project=PROJECT_ID, \n", + " location=LOCATION, \n", + " credentials=credentials\n", + " )\n", + ")\n", + "\n", + "# Export to environment variables for subprocesses (like evalbench.py)\n", + "os.environ[\"EVAL_GCP_PROJECT_ID\"] = PROJECT_ID\n", + "os.environ[\"EVAL_GCP_PROJECT_REGION\"] = LOCATION" + ], + "metadata": { + "id": "8adb6260" + }, + "execution_count": null + }, + { + "id": "80748249", + "cell_type": "markdown", + "source": [ + "## Step 4: Deploy to Agent Runtime using the AdkApp Template\n", + "\n", + "We load the local ADK agent and wrap it inside the `AdkApp` class. Then, we call `client.agent_engines.create` to deploy the agent.\n", + "\n", + "Since our agent is self-contained in `minimal_agent.py`, we only need to include `./minimal_agent.py` in the `extra_packages` deployment configuration." + ], + "metadata": { + "id": "80748249" + }, + "execution_count": null + }, + { + "id": "fe59c109", + "cell_type": "code", + "source": [ + "from vertexai.agent_engines.templates.adk import AdkApp\n", + "\n", + "import minimal_agent\n", + "\n", + "\n", + "# Instantiate the local ADK agent and wrap it in the template.\n", + "local_agent = minimal_agent.create_agent()\n", + "adk_app = AdkApp(agent=local_agent, app_name=\"minimal-adk-sql-app\")\n", + "\n", + "print(\"Deploying minimal ADK agent to Agent Runtime...\")\n", + "remote_agent = client.agent_engines.create(\n", + " agent=adk_app,\n", + " config=dict(\n", + " display_name=\"minimal_adk_agent\",\n", + " staging_bucket=STAGING_BUCKET,\n", + " requirements=[\n", + " \"google-cloud-aiplatform[agent_engines,adk]\",\n", + " \"google-genai\",\n", + " \"google-adk[all]==2.1.0\"\n", + " ],\n", + " extra_packages=[\"./minimal_agent.py\"]\n", + " )\n", + ")\n", + "\n", + "print(f\"\\n✅ Agent deployed successfully!\")\n", + "print(f\"Resource Name: {remote_agent.api_resource.name}\")" + ], + "metadata": { + "id": "fe59c109" + }, + "execution_count": null + }, + { + "id": "e277933c", + "cell_type": "markdown", + "source": [ + "## Step 5: Verify and test the live agent\n", + "\n", + "Query your live deployed agent over streaming to verify it returns SQL wrapped inside markdown code blocks.\n", + "\n", + "Once verified, copy the `Resource Name` printed in the previous step and paste it into the `resource_name` configuration field of your EvalBench `model_config` YAML file (e.g. `datasets/model_configs/agent_runtime.yaml`) to run evaluations!" + ], + "metadata": { + "id": "e277933c" + }, + "execution_count": null + }, + { + "id": "0164110c", + "cell_type": "code", + "source": [ + "# Test query the remote app client over streaming.\n", + "print(\"Querying remote agent endpoint...\")\n", + "response_stream = remote_agent.stream_query(\n", + " message=\"How many schools are in California?\",\n", + " user_id=\"evalbench_test\"\n", + ")\n", + "\n", + "print(\"\\nStreaming response chunks:\")\n", + "for chunk in response_stream:\n", + " print(chunk, end=\"\")" + ], + "metadata": { + "id": "0164110c" + }, + "execution_count": null + }, + { + "id": "66a03a42", + "cell_type": "markdown", + "source": [ + "## Step 6: Generate temporary config files for evaluation\n", + "\n", + "To run evaluations without modifying the default files in the repository, we create temporary copies of both the model configuration and the run configuration.\n", + "\n", + "The code block below:\n", + "1. Loads the default `agent_runtime.yaml` model config, updates the `resource_name` with your deployed agent ID, and writes it to a new temporary file `agent_runtime_temp.yaml`.\n", + "2. Loads the default run configuration, updates its `model_config` path to point to the temporary model config, and writes it to a new temporary file `test_agent_runtime_run_config_temp.yaml`." + ], + "metadata": { + "id": "66a03a42" + }, + "execution_count": null + }, + { + "id": "99f82cfd", + "cell_type": "code", + "source": [ + "import os\n", + "import json\n", + "import yaml\n", + "import pyaml_env\n", + "\n", + "# Using absolute paths to ensure files are located correctly\n", + "base_path = \"/content/evalbench\"\n", + "model_config_path = (\n", + " os.path.join(base_path, \"datasets/model_configs/agent_runtime.yaml\")\n", + ")\n", + "prompts_path = os.path.join(base_path, \"datasets/bird/prompts.json\")\n", + "\n", + "# Temporary file paths\n", + "temp_model_config_path = (\n", + " os.path.join(base_path, \"datasets/model_configs/agent_runtime_temp.yaml\")\n", + ")\n", + "temp_prompts_path = (\n", + " os.path.join(base_path, \"datasets/bird/prompts_california_temp.json\")\n", + ")\n", + "temp_run_config_path = (\n", + " os.path.join(\n", + " base_path, \n", + " \"datasets/bird/test_agent_runtime_run_config_temp.yaml\"\n", + " )\n", + ")\n", + "\n", + "# Load, update, and write temporary model config using pyaml_env to resolve !ENV tags.\n", + "model_data = pyaml_env.parse_config(model_config_path)\n", + "\n", + "# Update the resource name with your deployed agent\n", + "model_data[\"resource_name\"] = remote_agent.api_resource.name\n", + "\n", + "with open(temp_model_config_path, \"w\") as f:\n", + " yaml.dump(model_data, f)\n", + "print(f\"Created temporary model config: {temp_model_config_path}\")\n", + "\n", + "# Slice prompts.json to extract California Schools prompts.\n", + "with open(prompts_path, \"r\") as f:\n", + " all_prompts = json.load(f)\n", + "\n", + "california_prompts = [\n", + " p for p in all_prompts if p.get(\"db_id\") == \"california_schools\"\n", + "]\n", + "sliced_prompts = california_prompts[:3]\n", + "\n", + "with open(temp_prompts_path, \"w\") as f:\n", + " json.dump(sliced_prompts, f, indent=2)\n", + "print(f\"Created temporary prompts slice: {temp_prompts_path}\")\n", + "\n", + "# Generate the temporary run configuration with absolute paths.\n", + "run_config_content = f\"\"\"\n", + "dataset_config: {temp_prompts_path}\n", + "\n", + "database_configs:\n", + " - {os.path.join(base_path, 'datasets/bird/db_configs/sqlite.yaml')}\n", + "\n", + "dialects:\n", + " - sqlite\n", + "query_types:\n", + " - dql\n", + "dataset_format: bird-standard-format\n", + "\n", + "model_config: {temp_model_config_path}\n", + "prompt_generator: 'NOOPGenerator'\n", + "\n", + "scorers:\n", + " set_match: null\n", + " exact_match: null\n", + "\n", + "reporting:\n", + " csv: {{}}\n", + "\"\"\"\n", + "\n", + "with open(temp_run_config_path, \"w\") as f:\n", + " f.write(run_config_content.strip())\n", + "print(f\"Created temporary run config: {temp_run_config_path}\")" + ], + "metadata": { + "id": "99f82cfd" + }, + "execution_count": null + }, + { + "id": "f71131b9", + "cell_type": "markdown", + "source": [ + "## Step 7: Run EvalBench evaluation pipeline\n", + "\n", + "We can now trigger the EvalBench evaluation run using the updated configuration. This executes query generation on your deployed agent and performs BigQuery execution and correctness scoring." + ], + "metadata": { + "id": "f71131b9" + }, + "execution_count": null + }, + { + "id": "d88ccc37", + "cell_type": "code", + "source": [ + "import os\n", + "# Change to the evalbench root directory\n", + "%cd /content/evalbench\n", + "\n", + "# Run the simplified execution command without the Python protobuf override\n", + "! .venv/bin/python3 evalbench/evalbench.py \\\n", + " --experiment_config=\"datasets/bird/test_agent_runtime_run_config_temp.yaml\"" + ], + "metadata": { + "id": "d88ccc37" + }, + "execution_count": null + } + ], + "metadata": { + "colab": { + "provenance": [] + }, + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.14" + } + }, + "nbformat_minor": 5, + "nbformat": 4 }