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..d29767b9 --- /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.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: !ENV ${EVAL_GCP_PROJECT_ID} + csv: {} diff --git a/datasets/model_configs/agent_runtime.yaml b/datasets/model_configs/agent_runtime.yaml new file mode 100644 index 00000000..1b699640 --- /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 AGENT_ENGINE_RESOURCE from env. +resource_name: !ENV ${AGENT_ENGINE_RESOURCE:""} + +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 +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..5dee2ef4 --- /dev/null +++ b/docs/examples/agent_runtime_deployment.ipynb @@ -0,0 +1,470 @@ +{ + "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 + }, + { + "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 +} 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..9ab6733b --- /dev/null +++ b/evalbench/generators/models/agent_runtime.py @@ -0,0 +1,141 @@ +"""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.api_core.exceptions import ResourceExhausted +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 +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: + """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.""" + # 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 sanitize_sql(str(data["sql"])) + except json.JSONDecodeError: + pass + + # 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()) + + # 3. Fallback: Treat as raw SQL query + return sanitize_sql(text) + + +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 = ( + get_gcp_project(querygenerator_config.get("gcp_project_id")) + ) + location = get_gcp_region(querygenerator_config.get("gcp_region")) + 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 ResourceExhausted as e: + raise ResourceExhaustedError(e) + 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..67744d9b --- /dev/null +++ b/evalbench/test/agent_runtime_test.py @@ -0,0 +1,165 @@ +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_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; " + 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/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/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" diff --git a/pyproject.toml b/pyproject.toml index 2d001fb8..a6fb86ad 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",