feat: Refactor gepa sample code and clean-up user demo colab

PiperOrigin-RevId: 828293079

contributing/samples/gepa/adk_agent.py

@@ -31,8 +31,10 @@ from google.adk.agents import llm_agent
 from google.adk.agents import loop_agent
 from google.adk.events import event as event_lib
 from google.adk.models import google_llm
+from google.adk.planners import built_in_planner
 from google.adk.tools import base_tool
 from google.genai import types
+from retry import api as retry
 
 
 class EnvResponse(Protocol):
@@ -134,6 +136,11 @@ def _adk_agent(
           model=model or 'gemini-2.5-flash',
           retry_options=_default_retry_options(),
       ),
+      planner=built_in_planner.BuiltInPlanner(
+          thinking_config=types.ThinkingConfig(
+              thinking_budget=-1, include_thoughts=False
+          )
+      ),
       instruction=instruction,
       tools=tools,
       generate_content_config=types.GenerateContentConfig(
@@ -173,11 +180,17 @@ class _UserAgent(base_agent.BaseAgent):
       return
 
     if last_event.content and last_event.content.parts:
-      next_message = last_event.content.parts[-1].text
+      next_message = '\n\n'.join([p.text for p in last_event.content.parts])
     else:
       logging.warn('Empty content with event=%s', last_event)
       next_message = ''
-    env_response = self.env.step(types.Part(text=next_message))
+    env_response = retry.retry_call(
+        self.env.step,
+        fargs=(types.Part(text=next_message),),
+        tries=3,
+        delay=2,
+        backoff=2,
+    )
 
     output_event = event_lib.Event(
         content=types.Content(

contributing/samples/gepa/rater_lib.py

@@ -0,0 +1,178 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Library for rating agent trajectories."""
+
+from __future__ import annotations
+
+import re
+from typing import Any
+
+from google.genai import types
+import jinja2
+from retry import retry
+
+from google import genai
+
+
+def parse_rubric_validation_response(
+    rubric_val_response: str,
+) -> dict[str, str]:
+  """Parses rubric validation response text into a dictionary.
+
+  Args:
+    rubric_val_response: The text response from rubric validation.
+
+  Returns:
+    A dictionary containing parsed property, evidence, rationale, and verdict.
+  """
+  PROPERTY_PATTERN = (
+      r'Property:\s*([\s\S]*?)(?=(?:Evidence:|Rationale:|Verdict:|$))'
+  )
+  EVIDENCE_PATTERN = r'Evidence:\s*([\s\S]*?)(?=(?:Rationale:|Verdict:|$))'
+  RATIONALE_PATTERN = r'Rationale:\s*([\s\S]*?)(?=(?:Evidence:|Verdict:|$))'
+  VERDICT_PATTERN = r'Verdict:\s*([\s\S]*?)(?=(?:Evidence:|Rationale:|$))'
+
+  property_list = []
+  evidence_list = []
+  rationale_list = []
+  fulfillment_list = []
+  property_blocks = rubric_val_response.split('Property: ')[1:]
+  for property_block in property_blocks:
+    property_name = re.search(PROPERTY_PATTERN, 'Property: ' + property_block)
+    if property_name is None:
+      continue
+    property_name = property_name.group(1).strip()
+    property_list.append(property_name)
+
+    evidence_match = re.search(EVIDENCE_PATTERN, property_block, re.DOTALL)
+    evidence = evidence_match.group(1).strip() if evidence_match else ''
+    evidence_list.append(evidence)
+
+    rationale_match = re.search(RATIONALE_PATTERN, property_block, re.DOTALL)
+    rationale = rationale_match.group(1).strip() if rationale_match else ''
+    rationale_list.append(rationale)
+
+    verdict = re.search(VERDICT_PATTERN, property_block)
+    if verdict is None:
+      verdict_str = 'not_found'
+    else:
+      verdict_str = verdict.group(1).strip().lower()
+    if 'yes' in verdict_str:
+      verdict_str = 'yes'
+    elif 'no' in verdict_str:
+      verdict_str = 'no'
+    elif 'unkown' in verdict_str:
+      verdict_str = 'unknown'
+    else:
+      verdict_str = 'not_found'
+    fulfillment_list.append(verdict_str)
+  return dict(
+      property=property_list[0],
+      evidence=evidence_list[0],
+      rationale=rationale_list[0],
+      verdict=fulfillment_list[0],
+  )
+
+
+def format_user_agent_conversation(conv: list[dict[str, Any]]) -> str:
+  """Formats a conversation between user and agent into a string.
+
+  Args:
+    conv: A list of conversation turns.
+
+  Returns:
+    A formatted string representing the conversation.
+  """
+  # conv is a list in this eval data
+  # if not, manually convert to list to re-use these logics
+  # if not isinstance(conv, list):
+  #   conv = [conv]
+  res = ''
+  turn_idx = 1
+  for turn in conv:
+    # if 'request' in conv[turn]:\
+    role = turn['role']
+    for part in turn['parts']:
+      if role == 'user' and (txt := part.get('text')):
+        res = res + f'USER TURN {turn_idx}:\n' + txt + '\n'
+        turn_idx += 1
+      elif role == 'model' and (txt := part.get('text')):
+        res = res + f'The agent response is: {txt}' + '\n'
+      elif fc := part.get('function_call'):
+        res = (
+            res
+            + f'The agent called the function {fc["name"]} with the following'
+            f' function arguments: {fc["args"]}.\n'
+        )
+      elif fc := part.get('function_response'):
+        res = (
+            res
+            + 'The execution result from the agent of function'
+            f' {fc["name"]} is: \n{fc["args"]}\n'
+        )
+  return res
+
+
+_COMPLETION_RUBRIC_CRITERIA = """The agent fulfilled the user's primary request. Description: It measures if the agent successfully completed the action the user initiated the contact for (e.g., processed a return, provided a tracking number, answered a policy question). A "yes" requires confirmed completion within the transcript."""
+
+
+class Rater:
+  """Rates agent trajectories using an LLM based on rubrics."""
+
+  def __init__(self, tool_declarations: str):
+    """Initializes the Rater.
+
+    Args:
+      tool_declarations: JSON string of tool declarations for the agent.
+    """
+    self._client = genai.Client()
+    self._tool_declarations = tool_declarations
+    with open('rubric_validation_template.txt') as f:
+      self._rubric_validation_template = f.read().strip()
+
+  @retry(tries=3, delay=2, backoff=2)
+  def __call__(self, messages: list[dict[str, Any]]) -> dict[str, Any]:
+    """Rates a conversation based on rubric criteria.
+
+    Args:
+      messages: A list of conversation messages between user and agent.
+
+    Returns:
+      A dictionary containing rating information including score.
+    """
+    env = jinja2.Environment()
+    env.globals['user_input'] = (
+        messages[0].get('parts', [{}])[0].get('text', '') if messages else ''
+    )
+    env.globals['developer_instructions'] = ''
+    env.globals['tool_declarations'] = self._tool_declarations
+    env.globals['model_response'] = format_user_agent_conversation(messages)
+    env.globals['decomposed_rubric'] = '* ' + _COMPLETION_RUBRIC_CRITERIA
+    contents = env.from_string(self._rubric_validation_template).render()
+    resp = self._client.models.generate_content(
+        model='gemini-2.5-pro',
+        contents=contents,
+        config=types.GenerateContentConfig(
+            candidate_count=1,
+            thinking_config=types.ThinkingConfig(
+                include_thoughts=True, thinking_budget=-1
+            ),
+        ),
+    )
+    got = parse_rubric_validation_response(resp.text)
+    got = dict(got)
+    got['score'] = float(got['verdict'] == 'yes')
+    got['rating_criteria'] = got.pop('property')
+    return got

contributing/samples/gepa/rubric_validation_template.txt

@@ -0,0 +1,170 @@
+# Mission
+Your mission is to act as an impartial quality assurance analyst. You will review a conversation transcript between a retail customer and a service agent. Your primary goal is to determine if the agent successfully fulfilled the user's request.
+
+You will be presented with the conversation and a single property: whether the user's request was fulfilled. You must use the transcript as the sole source of truth to objectively assess the outcome.
+
+# Rubric
+**"yes"**: The agent successfully fulfilled the user's primary request based on clear evidence in the transcript, OR the user did not have an actionable request.
+**"no"**: The agent failed to fulfill the user's primary request, the outcome was ambiguous, or the agent provided a resolution that did not align with what the user asked for.
+
+# Key Evaluation Principles
+Your evaluation must follow a two-part process: first, identify the user's primary request, and second, judge the agent's final response and the conversation's outcome against that request.
+
+1.  **Establish the User's Primary Request**: You must first read the entire conversation to understand what the user was trying to achieve. The primary request is the main reason the user initiated the contact.
+    *   Your ONLY source of truth is the full conversation found in `<main_prompt>` and `<responses>`.
+    *   Examples of primary requests include:
+        *   Returning an item.
+        *   Checking an order status.
+        *   Asking for product information.
+        *   Filing a complaint about a product or service.
+        *   Updating account information.
+    *   If the user has multiple requests, focus on the main, initial one. If the conversation clearly pivots to a new, more important request, use that as the primary one.
+
+2.  **Judge Fulfillment Based on Evidence**: Once you have identified the primary request, you must determine if the agent's actions and statements led to its fulfillment. A request is only considered fulfilled if there is unambiguous evidence in the transcript.
+    *   **Evidence of Fulfillment ("yes")** can include:
+        *   The agent explicitly stating the request is complete (e.g., "I've now processed your refund," "Your tracking number is XYZ.").
+        *   The user explicitly confirming their issue is resolved (e.g., "Great, that's all I needed," "Thank you, that answers my question.").
+        *   The agent providing a complete and direct answer to a question (e.g., User asks for store hours, agent provides them).
+    *   **Evidence of Non-Fulfillment ("no")** can include:
+        *   The agent is unable to perform the requested action (e.g., "Our system is down, I can't process returns right now.").
+        *   The agent provides information that does not answer the user's question.
+        *   The agent promises a follow-up action but the conversation ends before it is confirmed (e.g., "Someone will call you back within 24 hours.").
+        *   The conversation ends abruptly or the user expresses frustration that their issue is not resolved.
+    *   **Crucial Clarification**: Do not make assumptions. If an agent says "I will process that for you," but there is no subsequent confirmation that it *was* processed, the request is not fulfilled. The action must be confirmed as completed within the conversation.
+
+For the property, follow these internal steps:
+1.  Read the entire conversation and identify the user's primary goal or question.
+2.  Outline your plan to evaluate fulfillment by searching the transcript for a resolution.
+3.  Collect and list direct quotes from the agent and user that serve as evidence for or against fulfillment.
+4.  Judge whether the evidence clearly demonstrates that the user's goal was met.
+5.  Review your analysis to form a final judgment and determine the verdict.
+6.  Output the final verdict in the required output format.
+
+# Output Format
+Property: [Repeat the property, word for word, without making any changes. Keep everything including punctuation and capitalization as-is.]
+Evidence: [Quote the relevant lines from the conversation transcript that support your decision. Reference the speaker (User or Agent).]
+Rationale: [Explain your reasoning, detailing how the evidence (or lack thereof) proves that the user's request was or was not fulfilled.]
+Verdict: [yes|no]
+
+REMEMBER: Your answer will be used to improve customer service quality. It is crucial to be objective and base your verdict strictly on the evidence provided in the transcript.
+
+# Example 1 (Request Fulfilled)
+## Input
+<user_prompt>
+  <available_tools>
+  {
+    "name": "get_order_status",
+    "description": "Retrieves the status and tracking information for a given order ID.",
+    "parameters": [
+      {
+        "type": "string",
+        "name": "order_id",
+        "description": "The unique identifier for the customer's order."
+      }
+    ]
+  },
+  {
+    "name": "process_return",
+    "description": "Initiates a return process for a given order ID and generates a shipping label.",
+    "parameters": [
+      {
+        "type": "string",
+        "name": "order_id",
+        "description": "The unique identifier for the order to be returned."
+      }
+    ]
+  }
+  </available_tools>
+
+  <main_prompt>
+  Hi, I need to check the status of my order, #98765.
+  </main_prompt>
+</user_prompt>
+
+<responses>
+Agent: Of course, I can help with that. One moment while I look it up.
+Agent: Okay, I see order #98765. It looks like it was shipped this morning. The tracking number is 1Z987ABC.
+User: Great, that's all I needed. Thank you!
+</responses>
+
+<properties>
+* The agent fulfilled the user's primary request.
+</properties>
+
+## Output
+Property: The agent fulfilled the user's primary request.
+Evidence: User: "Hi, I need to check the status of my order, #98765." Agent: "The tracking number is 1Z987ABC." User: "Great, that's all I needed. Thank you!"
+Rationale: The user's primary request was to check their order status. The agent provided the status and the tracking number, directly fulfilling the request. The user confirmed that their need was met.
+Verdict: yes
+
+# Example 2 (Request Not Fulfilled)
+## Input
+<user_prompt>
+  <available_tools>
+  {
+    "name": "get_order_status",
+    "description": "Retrieves the status and tracking information for a given order ID.",
+    "parameters": [
+      {
+        "type": "string",
+        "name": "order_id",
+        "description": "The unique identifier for the customer's order."
+      }
+    ]
+  },
+  {
+    "name": "process_return",
+    "description": "Initiates a return process for a given order ID and generates a shipping label.",
+    "parameters": [
+      {
+        "type": "string",
+        "name": "order_id",
+        "description": "The unique identifier for the order to be returned."
+      }
+    ]
+  }
+  </available_tools>
+
+  <main_prompt>
+  I'd like to return the shoes I bought last week. The order number is #54321.
+  </main_prompt>
+</user_prompt>
+
+<responses>
+Agent: I can help you with that. Can you confirm your shipping address?
+User: Yes, it's 123 Main St, Anytown.
+Agent: Thank you. Unfortunately, our return system is experiencing technical difficulties right now. I can't generate a return label. I can try again in a few hours.
+User: Oh. Okay, I guess just let me know.
+</responses>
+
+<properties>
+* The agent fulfilled the user's primary request.
+</properties>
+
+## Output
+Property: The agent fulfilled the user's primary request.
+Evidence: User: "I'd like to return the shoes I bought last week." Agent: "Unfortunately, our return system is experiencing technical difficulties right now. I can't generate a return label."
+Rationale: The user's primary request was to initiate a return for their shoes. The agent was unable to complete this action due to a system issue. The conversation ended without the user's request being fulfilled.
+Verdict: no
+
+# Your Turn
+## Input
+<user_prompt>
+  <available_tools>
+  {{tool_declarations}}
+  </available_tools>
+
+  <main_prompt>
+  {{user_input}}
+  </main_prompt>
+</user_prompt>
+
+<responses>
+{{model_response}}
+</responses>
+
+<properties>
+{{decomposed_rubric}}
+</properties>
+
+## Output

contributing/samples/gepa/run_experiment.py

@@ -0,0 +1,116 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Runs a GEPA experiment on Tau-Bench."""
+
+from collections.abc import Sequence
+import dataclasses
+from datetime import datetime
+import json
+import logging
+import os
+
+from absl import app
+from absl import flags
+import experiment
+from google.genai import types
+
+_OUTPUT_DIR = flags.DEFINE_string('output_dir', None, '')
+_EVAL_SET_SIZE = flags.DEFINE_integer('eval_set_size', None, '')
+_MAX_METRIC_CALLS = flags.DEFINE_integer('max_metric_calls', 500, '')
+_NUM_TEST_RECORDS = flags.DEFINE_integer('num_test_records', None, '')
+_NUM_EVAL_TRIALS = flags.DEFINE_integer('num_eval_trials', 4, '')
+_MAX_CONCURRENCY = flags.DEFINE_integer('max_concurrency', 8, '')
+_EVAL_MODE = flags.DEFINE_bool('eval_mode', False, '')
+_USE_RATER = flags.DEFINE_bool('use_rater', False, '')
+_TRAIN_BATCH_SIZE = flags.DEFINE_integer('train_batch_size', 3, '')
+
+
+def main(argv: Sequence[str]) -> None:
+  if len(argv) > 1:
+    raise app.UsageError('Too many command-line arguments.')
+
+  # Get a list of all existing loggers
+  # logging.root.manager.loggerDict contains all named loggers
+  # logging.getLogger(name) retrieves the logger object
+  loggers = [
+      logging.getLogger(name) for name in logging.root.manager.loggerDict
+  ]
+
+  # Iterate through the loggers and set their level to WARNING
+  for logger in loggers:
+    logger.setLevel(logging.WARNING)
+
+  types.logger.addFilter(experiment.FilterInferenceWarnings())
+  if not _OUTPUT_DIR.value:
+    raise ValueError('outptut dir must be specified')
+  output_dir = os.path.join(
+      _OUTPUT_DIR.value, datetime.now().strftime('%Y%m%d%H%M%S%f')
+  )
+  os.makedirs(output_dir)
+  logging.info('Writing to output_dir=%s', output_dir)
+  config = experiment.ExperimentConfig(
+      tau_bench_env='retail',
+      agent_model='gemini-2.5-flash',
+      agent_model_provider='vertex_ai',
+      user_model='gemini-2.5-flash',
+      user_model_provider='vertex_ai',
+      max_concurrency=_MAX_CONCURRENCY.value,
+      num_eval_trials=_NUM_EVAL_TRIALS.value,
+      rnd_seed=42,
+      max_metric_calls=_MAX_METRIC_CALLS.value,
+      reflection_model='gemini-2.5-pro',
+      reflection_minibatch_size=_TRAIN_BATCH_SIZE.value,
+      use_rater=_USE_RATER.value,
+      feedback_dataset=experiment.Dataset(split='train'),
+      pareto_dataset=experiment.Dataset(
+          split='dev', max_size=_EVAL_SET_SIZE.value
+      ),
+      eval_dataset=experiment.Dataset(
+          split='test', max_size=_NUM_TEST_RECORDS.value
+      ),
+  )
+  json.dump(
+      dataclasses.asdict(config),
+      open(os.path.join(output_dir, 'config.json'), 'w'),
+  )
+  logging.info('Using config=%s', config)
+
+  if _EVAL_MODE.value:
+    return experiment.run_eval(
+        output_dir=output_dir,
+        instructions=experiment.SEED_SYSTEM_INSTRUCTION,
+        config=config,
+    )
+
+  results = experiment.run_gepa(
+      config=config,
+      seed_instructions=experiment.SEED_SYSTEM_INSTRUCTION,
+      output_dir=output_dir,
+  )
+  print(list(enumerate(results.val_aggregate_scores)))
+
+  eval_dir = os.path.join(
+      output_dir, 'evals', datetime.now().strftime('%Y%m%d%H%M%S%f')
+  )
+  os.makedirs(eval_dir)
+  experiment.run_eval(
+      output_dir=eval_dir,
+      instructions=results.best_candidate['system_instruction'],
+      config=config,
+  )
+
+
+if __name__ == '__main__':
+  app.run(main)