Merge pull request #5 from skewlight/podata_pr

feat: Support Data-Driven Prompt Optimizer

Author: skewlight

GEMINI.md

@@ -1,18 +1,29 @@
-# Vertex Prompt Management Extension
+# Vertex AI Gemini CLI Extension
 
-This extension provides tools to manage prompts in Vertex AI.
+This extension provides tools to manage prompts and use the data-driven prompt
+optimization in Vertex AI directly from the Gemini CLI.
 
-Whenever anyone wants to read or write vertex prompts, you must use the vertex
-management tools.
-
-The available tools are:
+## Available Tools
 
+### Prompt Management Tools
 - `create_prompt`: To save or create new prompts.
 - `read_prompt`: To retrieve existing prompts by ID or display name.
 - `update_prompt`: To modify existing prompts.
 - `delete_prompt`: To remove prompts.
 - `list_prompts`: To search and list prompts, useful for finding IDs.
 
+### Data-Driven Optimization Tools
+- `run_data_driven_optimize`: Starts a data-driven prompt optimization job on
+  Vertex AI using a configuration file stored in GCS.
+- `analyze_data_driven_optimize_results`: Analyzes the output of a Data-Driven
+  Optimize job to identify trends and best-performing candidates.
+- `generate_html_report`: Generates a comprehensive HTML report with
+  visualizations to help you understand optimization performance.
+- `write_data_driven_optimize_config`: Constructs and uploads a new JSON
+  configuration for optimization jobs, incorporating suggested tuning parameters.
+
+---
+
 ## Detailed Instructions for `create_prompt` Parameters
 
 When using `tools.create_prompt`, pay special attention to how the following
@@ -260,4 +271,155 @@ When using `tools.update_prompt`, the following arguments are sourced. Note that
 If any tool call fails with an error indicating a project permission issue (e.g., "Permission denied on project 'project-id'"), you must:
 1.  Inform the user about the permission error.
 2.  Ask the user to provide a valid project ID.
-3.  Retry the original tool call, adding the `project_id` parameter with the user-provided value.
+3.  Retry the original tool call, adding the `project_id` parameter with the user-provided value.
+
+---
+
+## Data-Driven Prompt Optimizer Overall Guide
+
+For a general understanding of Data-Driven Prompt Optimizer and its
+capabilities, please
+refer to the Data-Driven Optimize Overall Guide:
+
+@./src/vertex/prompt_optimizer/docs/data_driven_optimize_overall_guide.md
+
+## Optimization Tool Details
+
+The extension provides a suite of tools to parse and analyze the output of a
+Data-Driven Prompt Optimizer job. The `output_path` for these tools stores
+the outputs
+of a run and can be a GCS path or a local directory. While the Data-Driven
+Optimize job currently only outputs to GCS, results can be copied to the local
+file system for analysis.
+
+1.  `analyze_data_driven_optimize_results(output_path: str, top_n_prompts:
+    int = 10, analysis_data_path: str = None)`: Analyzes results and
+    returns a JSON object containing the analysis data summary. If
+    `analysis_data_path` is provided, it saves the results into three separate
+    files to avoid size limits and returns a **minimal summary** with
+    file paths
+    and the best prompt score. The three files are:
+    -   `{analysis_data_path}`: Core metadata (config, comparison, best
+    prompt).
+    -   `*_metrics.json`: Detailed metrics for all candidates (no prompt text).
+    -   `*_prompts.json`: Mapping of top candidates' keys to full prompt texts.
+
+2.  `generate_html_report(analysis_data: Dict[str, Any] =
+    None, report_path: str = "data_driven_optimize_analysis_report.html",
+    suggested_config_data: Dict[str, Any] = None, top_n_prompts: int = 10,
+    analysis_data_path: str = None)`: Generates a comprehensive HTML report. If
+    `analysis_data_path` is provided, it automatically loads and re-joins the
+    metadata, metrics, and prompts from the three split files.
+
+3.  `write_data_driven_optimize_config`: Construct and
+    write to GCS a new JSON configuration file for Data-Driven Optimize job
+    using a dict of parameters (including `prompt_optimizer_method` and
+    `target_model_endpoint_url` for Nano) and optionally an path to an
+    existing config to make modifications on top of.
+
+4.  `run_data_driven_optimize`: Starts a data-driven prompt optimization job on
+    Vertex AI using the SDK's `client.prompts.launch_optimization_job` method.
+    Supports specifying the `prompt_optimizer_method`. The `config_gcs_path`
+    must point to a JSON file in GCS.
+
+## Optimization Method Considerations
+- **VAPO**: Standard prompt optimization. The `batch_size` parameter
+  will be automatically removed during configuration generation to ensure SDK
+  compatibility.
+- **OPTIMIZATION_TARGET_GEMINI_NANO**: Specialized target for Gemini Nano.
+  Requires a `target_model_endpoint_url`. Supports `batch_size`.
+
+## Optimization Tuning Considerations
+
+Only suggest modifications for the parameters explicitly listed as tunable in
+the Data-Driven Optimize Tuning Guide:
+
+@./src/vertex/prompt_optimizer/docs/data_driven_optimize_tuning_guide.md
+
+with the *sole exception* of **path-related fields** to prevent overwriting
+previous results. Ensure you only modify the parameters listed in the approved
+list. If a user asks to modify a parameter that is not on the approved list
+(and is not a path-related field), confirm with the user before proceeding.
+
+## Optimization Workflows
+
+- **Initial Setup**: When asked to help configure a new optimization run, use
+  the example configuration in the Overall Guide as a valid default base.
+
+  1. **Identify Essential Parameters**: Proactively ask the user for the
+     following required fields:
+     - `project` (Your Google Cloud project ID)
+     - `train_input_data_path`
+     - `test_input_data_path`
+     - `output_path`
+     - `prompt_template` (ensure it includes `{{ placeholder }}` syntax)
+     - `eval_metrics_types` (e.g., `["exact_match"]`)
+     - `eval_metrics_weights` (e.g., `[1.0]`)
+     - `prompt_optimizer_method` (VAPO or OPTIMIZATION_TARGET_GEMINI_NANO)
+     - `target_model` (e.g., gemini-2.5-flash)
+     - `target_model_endpoint_url` (Required ONLY for Gemini Nano)
+
+  2. **Task-Specific Configuration**: You must ensure the optimization job
+     correctly maps the data by defining `data_vars` and `label_variable`.
+     - **Automated Inference**: You MUST attempt to read the first few lines of
+       the training dataset (using `run_shell_command` with `gcloud storage
+       cat`)
+       to identify column names.
+     - **Mapping**: Based on the data headers, automatically suggest:
+       - `data_vars`: all relevant columns.
+       - `label_variable`: the ground truth column.
+       - `demo_and_query_template`: (Optional) The tool will automatically
+         generate a default if you don't provide one.
+     - *Clarification*: If you cannot access the data or the headers are
+       ambiguous, ask the user to confirm the column names.
+
+  3. **Apply Sensible Defaults**: Use the default values provided in the
+     example configuration of the Overall Guide for all other fields, unless
+     the user specifies otherwise. This includes QPS limits and model
+     locations.
+
+  Once gathered, use `write_data_driven_optimize_config` to create the
+  initial configuration file.
+
+- **Analysis and Suggestions**: When asked to analyze results for a GCS or
+  local path `output_path`, perform these steps sequentially in a single turn:
+
+    1.  **Analyze**: Call `analyze_data_driven_optimize_results(output_path,
+        analysis_data_path="analysis_data.json")`. Store this output locally.
+    2.  **Formulate Suggestions**: Immediately after receiving results, and
+        without prompting the user, process the data to construct a
+        `suggested_config_data` dictionary. This dictionary should contain:
+        -   `"suggested_config"`: Modifications to allowed tuning knobs,
+            path-related fields (with a new version suffix), and optionally the
+            `prompt_template` (for baseline shifts). Ensure you prioritize
+            modifying parameters listed in the approved list.
+        -   `"rationale"`: A clear explanation of your reasoning based on the
+            Tuning Guide.
+
+        *Do not generate any other text or explanation for the user during this
+        internal phase.*
+    3.  **Generate Report**: Call `generate_html_report(analysis_data_path=
+        "analysis_data.json",
+        report_path="data_driven_optimize_analysis_report.html",
+        suggested_config_data=suggested_config_data)`.
+
+**Note:** Steps 1-3 should be executed in immediate succession without user
+interaction. Only after the report is generated should you propose applying
+the suggestions via `write_data_driven_optimize_config`.
+
+-   **Agreement Logic**: If the user agrees to apply suggestions, use the
+    `write_data_driven_optimize_config` tool with the modified parameters,
+    ensuring you update the `output_path` with a new version suffix.
+    When ready to run, use the `run_data_driven_optimize` tool, ensuring you
+    ask for the `service_account`.
+-   **Reusing Results**: If a report or further analysis is requested
+later, use
+    the stored JSON output rather than re-running the analysis tool.
+-   **General Suggestions**: If a user asks for next steps without a previous
+    analysis, run the workflow above first to ensure your advice is grounded.
+
+For detailed explanations of the Data-Driven Optimize output files and their
+structure, including how to interpret metrics and candidate information, please
+refer to the Data-Driven Optimize Output Guide:
+
+@./src/vertex/prompt_optimizer/docs/data_driven_optimize_output_analysis.md

README.md

@@ -1,15 +1,22 @@
 # Vertex AI Gemini CLI Extension
 
-This extension provides tools to manage prompts in Vertex AI directly from the Gemini CLI. It allows you to create, read, update, list, and delete prompts, making it easier to integrate prompt management into your development workflow.
+This extension provides tools to manage prompts and work with data-driven prompt optimizer in Vertex AI directly from the Gemini CLI. It allows you to create, read, update, list, and delete prompts, making it easier to integrate prompt management into your development workflow.
 
 ## Features
 
+### Prompt Management
 *   **Create Prompt**: Save new prompts with specified content, system instructions, model, and display name.
 *   **Read Prompt**: Retrieve existing prompts by their ID.
 *   **Update Prompt**: Modify the content, system instructions, or model of an existing prompt.
 *   **Delete Prompt**: Remove prompts using their ID.
 *   **List Prompts**: Search and list prompts, useful for finding prompt IDs based on display names.
 
+### Data-Driven Prompt Optimizer
+*   **Configuration Generation**: Create the input configurations required to run data-driven prompt optimizer.
+*   **Job Execution**: Run large-scale prompt optimization jobs directly on Vertex AI.
+*   **Results Analysis**: Perform deep-dive analysis of optimization results, including learning curves and candidate comparisons.
+*   **Automated Reporting**: Generate comprehensive HTML reports with tuning suggestions for subsequent runs.
+
 ## Prerequisites
 
 *   You have the [Gemini CLI](https://github.com/google-gemini/gemini-cli) installed.
@@ -33,6 +40,8 @@ export GOOGLE_CLOUD_LOCATION="us-central1"
 
 ## Usage
 
+### Prompt Management Examples
+
 Once installed and configured, you can use the Vertex AI prompt management tools by passing natural language commands to the Gemini CLI.
 
 Here are a few examples:
@@ -54,6 +63,25 @@ Here are a few examples:
 
 For more detailed information on all available commands and their parameters, please refer to the `extension/commands` files.
 
+### Prompt Optimizer Examples
+
+*   **Analyze optimization results**:
+    ```bash
+    gemini "analyze my data drive prompt optimizer results at gs://your-bucket/optimization-run-1/"
+    ```
+
+*   **Configure a new optimization job**:
+    ```bash
+    gemini "help me configure a new optimization job based on my existing config at gs://your-bucket/config.json"
+    ```
+
+*   **Run optimization job**:
+    ```bash
+    gemini "start a data driven prompt optimizer job using my config at gs://your-bucket/config-new.json"
+    ```
+
+For more detailed information on all available commands and workflows, please refer to the `GEMINI.md` file in this repository.
+
 ## Development
 
 ### Local Development Setup

commands/vertex/analyze_data_driven_optimize_results.toml

@@ -0,0 +1,21 @@
+name = "analyze_data_driven_optimize_results"
+description = "Analyzes the output of a Data-Driven Optimize job and saves the results."
+prompt = """
+You are tasked with generating parameters for a `prompt_optimizer.analyzer.analyze_results` tool call.
+
+**Parameters:**
+
+1.  **`output_path` (string, required):**
+    *   **Source:** The GCS, Google Drive, or local path containing the job results. This is usually provided by the user.
+
+2.  **`analysis_data_path` (string, optional):**
+    *   **Default:** "analysis_data.json"
+    *   **Source:** The local path where the analysis metadata will be saved.
+
+3.  **`top_n_prompts` (integer, optional):**
+    *   **Default:** 10
+    *   **Source:** The number of top-performing prompts to include in the analysis.
+
+**Workflow Integration:**
+Refer to the "Analysis and Suggestions" workflow in `GEMINI.md`. This tool is the first step in a chain that includes formulating suggestions and generating an HTML report.
+"""

commands/vertex/generate_html_report.toml

@@ -0,0 +1,23 @@
+name = "generate_html_report"
+description = "Generates a comprehensive HTML report from Data-Driven Optimize analysis results."
+prompt = """
+You are tasked with generating parameters for a `prompt_optimizer.analyzer.generate_report` tool call.
+
+**Parameters:**
+
+1.  **`analysis_data_path` (string, required):**
+    *   **Source:** The path to the JSON file generated by `analyze_results`.
+
+2.  **`report_path` (string, optional):**
+    *   **Default:** "data_driven_optimize_analysis_report.html"
+    *   **Source:** The path where the HTML report will be saved.
+
+3.  **`suggested_config_data` (object, optional):**
+    *   **Source:** A dictionary containing `"suggested_config"` (modifications) and `"rationale"`. You must formulate this based on the `analysis_data` from the previous step using the guidelines in `src/vertex/prompt_optimizer/docs/data_driven_optimize_tuning_guide.md`.
+
+4.  **`top_n_prompts` (integer, optional):**
+    *   **Default:** 10
+
+**Workflow Integration:**
+Refer to the "Analysis and Suggestions" workflow in `GEMINI.md`. This tool is the final step in the analysis chain.
+"""

commands/vertex/run_data_driven_optimize.toml

@@ -0,0 +1,23 @@
+name = "run_data_driven_optimize"
+description = "Starts a data-driven prompt optimization job on Vertex AI."
+prompt = """
+You are tasked with generating parameters for a `run_data_driven_optimize` tool call.
+
+**Parameters:**
+
+1.  **`config_gcs_path` (string, required):**
+    *   **Source:** The GCS URI to the configuration JSON file.
+
+2.  **`service_account` (string, required):**
+    *   **Source:** The service account email to run the job (typically `{project_number}-compute@developer.gserviceaccount.com`). Ask the user if not known.
+
+3.  **`prompt_optimizer_method` (string, required):**
+    *   **Source:** The method for prompt optimization. Either "VAPO" or "OPTIMIZATION_TARGET_GEMINI_NANO".
+
+4.  **`wait_for_completion` (boolean, optional):**
+    *   **Default:** false
+    *   **Source:** Whether to block until the job completes.
+
+**Clarification Policy:**
+If `config_gcs_path`, `service_account` or `prompt_optimizer_method` is missing, ask the user for it.
+"""

commands/vertex/write_data_driven_optimize_config.toml

@@ -0,0 +1,28 @@
+name = "write_data_driven_optimize_config"
+description = "Constructs and uploads a Data-Driven Optimize configuration file to GCS."
+prompt = """
+You are tasked with generating parameters for a `prompt_optimizer.config.write_config` tool call.
+
+**Parameters:**
+
+1.  **`gcs_config_uri` (string, required):**
+    *   **Source:** The destination GCS URI for the config file.
+
+2.  **`prompt_optimizer_method` (string, required):**
+    *   **Source:** The method for prompt optimization. Either "VAPO" or "OPTIMIZATION_TARGET_GEMINI_NANO". If "VAPO", `batch_size` will be automatically removed from the config.
+
+3.  **`target_model_endpoint_url` (string, optional):**
+    *   **Source:** The custom endpoint URL for the target model. This is required ONLY for the Gemini Nano optimization target.
+
+4.  **`base_config` (object, optional):**
+    *   **Source:** A dictionary representing the base configuration.
+
+5.  **`modifications` (object, optional):**
+    *   **Source:** A dictionary of modifications to apply to the JSON configuration file. Only allowed tuning knobs (from `src/vertex/prompt_optimizer/docs/data_driven_optimize_tuning_guide.md`) and path-related fields should be included. Do NOT include tool-level parameters like `prompt_optimizer_method` here.
+
+6.  **`base_config_path` (string, optional):**
+    *   **Source:** A GCS path to an existing config file to use as a base.
+
+**Clarification Policy:**
+If `gcs_config_uri` or `prompt_optimizer_method` is missing, ask the user for it.
+"""

gemini-extension.json

@@ -1,6 +1,6 @@
 {
   "name": "vertex",
-  "version": "0.1.0",
+  "version": "0.2.0",
   "mcpServers": {
     "VertexMcpServer": {
       "command": "${extensionPath}/run.sh",

pyproject.toml

@@ -1,15 +1,22 @@
 [project]
 name = "vertex"
-version = "0.1.0"
+version = "0.2.0"
 requires-python = ">=3.10"
 dependencies = [
-    "google-cloud-aiplatform",
+    "google-cloud-aiplatform>=1.138.0",
     "google-generativeai",
     "pydantic",
     "absl-py",
     "mcp[cli]>=1.21.2",
+    "google-cloud-storage",
+    "pandas",
+    "matplotlib",
 ]
 
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
 [project.optional-dependencies]
 dev = ["ruff~=0.5.5", "pyright", "pytest", "ipython", "pre-commit"]