Merge pull request #14 from gemini-cli-extensions/few-shot-po

feat: Adding Few-Shot Prompt Optimization as a tool

Author: ssusie

GEMINI.md

@@ -1,6 +1,6 @@
 # Vertex AI Gemini CLI Extension
 
-This extension provides tools to manage prompts and use the data-driven prompt
+This extension provides tools to manage prompts and use prompt
 optimization in Vertex AI directly from the Gemini CLI.
 
 ## Available Tools
@@ -12,8 +12,10 @@ optimization in Vertex AI directly from the Gemini CLI.
 - `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
+### Prompt Optimization Tools
+- `run_few_shot_optimization`: optimize user prompt based on examples provided 
+  in the CSV file stored in the GCS bucket.
+-`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.
@@ -275,6 +277,28 @@ If any tool call fails with an error indicating a project permission issue (e.g.
 
 ---
 
+## Few-Shot Prompt Optimization Guide
+Few-shot prompt optimization is a tool used to refine a prompt (or system instructions) based on user-provided examples. These examples are stored in a CSV file and user must provide a GCS path to this data. The examples must follow one of two formats:
+`("prompt", "model_response", "target_response")` or 
+`("prompt", "model_response", "rubrics", "rubrics_evaluations")`.
+Depending on the format of the examples provided in the user's CSV file, the optimization uses one of the following methods:
+`TARGET_RESPONSE`: The prompt is tuned to generate outputs that closely align with a provided  ground-truth answer.
+`RUBRICS`: The model analyzes rubric rules and evaluation scores to tune the prompt and address specific failures.
+For more details, please refer to this guide: @./src/vertex/prompt_optimizer/docs/few_shot_prompt_optimization_guide.md
+The tool is accessed through `run_few_shot_prompt_optimization(prompt_to_optimize: str,
+  example_path: str, method: str)`. This tool should be called directly on the raw prompt without invoking create_prompt or any other tool first.
+
+## Few-shot Optimization Workflows
+
+- When asked to run few-shot prompt optimization, identify if all required parameters are passed.
+  1. **Identify Essential Parameters**: Proactively ask the user for the
+     following required fields:
+     - `prompt_to_optimize`        
+     - `example_path`
+     - `method`
+
+  2. Once the parameters are set, call the `run_few_shot_prompt_optimization` with those params.
+
 ## Data-Driven Prompt Optimizer Overall Guide
 
 For a general understanding of Data-Driven Prompt Optimizer and its
@@ -283,7 +307,7 @@ refer to the Data-Driven Optimize Overall Guide:
 
 @./src/vertex/prompt_optimizer/docs/data_driven_optimize_overall_guide.md
 
-## Optimization Tool Details
+## Data-Driven 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
@@ -322,14 +346,14 @@ file system for analysis.
     Supports specifying the `prompt_optimizer_method`. The `config_gcs_path`
     must point to a JSON file in GCS.
 
-## Optimization Method Considerations
+## Data-Driven 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
+## Data-Driven Optimization Tuning Considerations
 
 Only suggest modifications for the parameters explicitly listed as tunable in
 the Data-Driven Optimize Tuning Guide:
@@ -341,7 +365,7 @@ 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
+## Data-Driven 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.

README.md

@@ -17,6 +17,10 @@ This extension provides tools to manage prompts and work with data-driven prompt
 *   **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.
 
+### Few-Shot Prompt Optimizer
+*  **Prompt Optimization**: tunes the given prompt by utilizing small set of user-provided examples using few-shot prompt optimization. For more details about the algorithm, please refer to this [documentation](https://docs.cloud.google.com/vertex-ai/generative-ai/docs/learn/prompts/few-shot-optimizer) or this [guide](@./src/vertex/prompt_optimizer/docs/few_shot_prompt_optimization_guide.md).
+
+
 ## Prerequisites
 
 *   You have the [Gemini CLI](https://github.com/google-gemini/gemini-cli) installed.
@@ -63,7 +67,7 @@ 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
+### Data-Driven Prompt Optimizer Examples
 
 *   **Analyze optimization results**:
     ```bash
@@ -82,6 +86,17 @@ For more detailed information on all available commands and their parameters, pl
 
 For more detailed information on all available commands and workflows, please refer to the `GEMINI.md` file in this repository.
 
+### Few-Shot Prompt Optimization Examples
+*   **Few-shot optimization with target response**:
+When user examples consit of `(question, model_response, target_response)`, few-shot optimization should be run with "target response" method.
+    ```bash
+    gemini "Apply few-shot prompt optimization using target response with data from gs://your_bucket/your_examples.csv. Here is the prompt: '...'"
+
+*   **Few-shot optimization with rubrics and evaluations**:
+When user examples consit of `(question, model_response, rubrics, rubrics_evaluations)`, few-shot optimization should be run with "rubrics" method.
+    ```bash
+    gemini "Apply few-shot prompt optimization using rubrics with data from gs://your_bucket/your_examples.csv. Here is the prompt: '...'"
+
 ## Development
 
 ### Local Development Setup

commands/vertex/run_few_shot_prompt_optimization.toml

@@ -0,0 +1,18 @@
+name = "run_few_shot_prompt_optimization"
+description = "Optimize a prompt based on few examples using few-shot optimization from Vertex AI SDK."
+prompt = """
+You are tasked to call `run_few_shot_prompt_optimization` tool call.
+
+**Parameters:**
+
+1. **`prompt_to_optimize` (string, required):**
+    *   **Source:** Prompt (or system instructions) that will be optimized.
+2. **`example_path` (string, required):**
+    *   **Source:** The GCS URI to the CSV file containing examples
+3. **`method:` (string, required):** 
+    *   **Source:** The method for prompt optimization. Either "TARGET_RESPONSE" or "RUBRICS"
+
+
+**Clarification Policy:**
+If `prompt_to_optimize`, `example_path` or `method` is missing, ask the user for it.
+"""

src/vertex/prompt_optimizer/docs/few_shot_prompt_optimization_guide.md

@@ -0,0 +1,28 @@
+# Few-Shot Prompt Optimizer
+
+The **Few-Shot Prompt Optimizer** helps you improve your system instructions by learning from examples.
+
+You provide your current instructions you want to tune along with a CSV file containing questions(prompts), the model's current responses, and your feedback. The optimizer analyzes where the model fell short and automatically updates your instructions to fix those gaps.
+
+Your example data must be saved as a CSV file in Google Cloud Storage (GCS). You can format your feedback using one of two methods:
+
+## Method 1: Target Responses
+
+Use this method to show the optimizer exactly what the model *should* have said. The tool compares the actual output to your ideal output and adjusts the instructions to bridge the gap.
+
+Your CSV must include these three columns:
+
+* **`prompt`**: The original user instruction or question.
+* **`model_response`**: The current model output that needs improvement.
+* **`target_response`**: The ideal, correct response you want the model to generate.
+
+## Method 2: Rubrics (Grading Criteria)
+
+Use this method to score responses based on specific rules. The tool learns from which rules the model passed or failed. Rubrics can be consistent across all rows or customized for each individual example.
+
+Your CSV must include these four columns:
+
+* **`prompt`**: The original user instruction or question.
+* **`model_response`**: The actual output generated by the model.
+* **`rubrics`**: A list of specific grading criteria (formatted as a string representation of a Python list of strings, e.g., `"['Tone is polite', 'Mentions pricing']"`).
+* **`rubrics_evaluations`**: True/False scores indicating if the model met each rule (formatted as a string representation of a Python list of booleans, e.g., `"['True', 'False']"`).

src/vertex/prompt_optimizer/prompt_optimizer.py

@@ -7,6 +7,7 @@
 from google.cloud import storage
 from vertexai._genai import Client
 
+from . import storage as storage_utils
 from . import utils
 
 
@@ -232,3 +233,39 @@ def run_data_driven_optimize(
             ) from re
         except Exception as e:
             raise RuntimeError(f"An unexpected error occurred: {e}") from e
+
+    def run_few_shot_optimization(
+        self,
+        prompt_to_optimize: str,
+        example_path: str,
+        method: str,
+    ) -> str:
+        """Applies few shot prompt optimization to a prompt using user provided dataset and method.
+
+        Args:
+            prompt_to_optimize: The zero-based index of the prompt to improve.
+            example_path: GCS path to the csv file containg few-shot examples
+            method: The optimization method to use for few shot prompt improvement. The
+                method should be one of the following:
+                - TARGET_RESPONSE: Optimize the prompt to match the target response.
+                - RUBRICS: Optimize the prompt to improve the rubrics scores.
+
+        Returns:
+            Optimized prompt.
+        """
+        if not self.client:
+            raise ValueError(
+                "Error: Client was not initialized properly. Try authenticating then rerunning"
+            )
+
+        example_df = storage_utils._read_csv_from_gcs(example_path)
+        optimization_target = utils._get_optimization_target(method)
+        config = vertexai.types.OptimizeConfig(
+            optimization_target=optimization_target,
+            examples_dataframe=example_df,
+        )
+        response = self.client.prompts.optimize(
+            prompt=prompt_to_optimize,
+            config=config,
+        )
+        return response.parsed_response.suggested_prompt

src/vertex/prompt_optimizer/storage.py

@@ -1,9 +1,11 @@
 """Utility functions for reading from and writing to GCS or local paths."""
 
+import io
 import json
 from pathlib import Path
 from typing import Any
 
+import pandas as pd
 from google.cloud import storage
 
 EXPECTED_GCS_PARTS = 2
@@ -33,6 +35,23 @@ def _read_json_from_gcs(gcs_uri: str) -> Any:
         raise RuntimeError(f"Failed to read from GCS URI {gcs_uri}: {e}") from e
 
 
+def _read_csv_from_gcs(path: str) -> pd.DataFrame:
+    if not _is_gcs_path(path) or not path.endswith(".csv"):
+        raise ValueError(
+            "Few-shot examples must be stored as CSV file in a GCS bucket, ",
+            f"but `{path}` was provided",
+        )
+    try:
+        storage_client = storage.Client()
+        bucket_name, blob_name = _parse_gcs_uri(path)
+        bucket = storage_client.bucket(bucket_name)
+        blob = bucket.blob(blob_name)
+        csv_bytes = blob.download_as_bytes()
+        return pd.read_csv(io.BytesIO(csv_bytes))
+    except Exception as e:
+        raise RuntimeError(f"Failed to read from GCS URI {path}: {e}") from e
+
+
 def _write_json_to_gcs(data: Any, gcs_uri: str):
     """Writes a JSON object to a GCS file."""
     storage_client = storage.Client()

src/vertex/prompt_optimizer/utils.py

@@ -2,6 +2,8 @@
 
 from typing import Any
 
+import vertexai
+
 
 def merge_configs(
     base_config: dict[str, Any], modifications: dict[str, Any]
@@ -19,3 +21,15 @@ def merge_configs(
         else:
             merged_config[key] = value
     return merged_config
+
+
+def _get_optimization_target(method: str) -> vertexai.types.OptimizeTarget:
+    """Returns the optimization target for the given method."""
+    if method == "TARGET_RESPONSE":
+        return vertexai.types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE
+    elif method == "RUBRICS":
+        return (
+            vertexai.types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_RUBRICS
+        )
+    else:
+        raise ValueError(f"Unsupported optimization method: {method}")

src/vertex/server.py

@@ -38,6 +38,10 @@ def main(argv: Sequence[str]) -> None:
         optimizer_tools.run_data_driven_optimize,
         name="run_data_driven_optimize",
     )
+    mcp.add_tool(
+        optimizer_tools.run_few_shot_optimization,
+        name="run_few_shot_optimization",
+    )
     mcp.add_tool(
         analyzer.analyze_results, name="analyze_data_driven_optimize_results"
     )