chore: import upstream snapshot with attribution

contributing/samples/integrations/gepa/run_experiment.py

@@ -0,0 +1,167 @@
+# Copyright 2026 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
+import gepa_utils
+from google.genai import types
+
+_OUTPUT_DIR = flags.DEFINE_string(
+    'output_dir',
+    None,
+    'Directory to save experiment results and artifacts.',
+    required=True,
+)
+_EVAL_SET_SIZE = flags.DEFINE_integer(
+    'eval_set_size',
+    None,
+    'Size of the dev set to use for Pareto frontier evaluation in GEPA. If'
+    ' None, uses all available dev tasks. A few tens of examples might'
+    ' suffice more simpler tasks and up to a few hundreds for '
+    ' more complex and variable tasks. Increase the size to mitigate effect of'
+    ' variability at greater cost.',
+)
+_MAX_METRIC_CALLS = flags.DEFINE_integer(
+    'max_metric_calls',
+    500,
+    'Total budget for GEPA prompt evaluations. This is the main control for'
+    ' runtime/cost. One could start with 100 and increase to 500+ for further'
+    ' optimization.',
+)
+_NUM_TEST_RECORDS = flags.DEFINE_integer(
+    'num_test_records',
+    None,
+    'Size of the test set for final evaluation of the optimized prompt. If'
+    ' None, uses all available test tasks.',
+)
+_NUM_EVAL_TRIALS = flags.DEFINE_integer(
+    'num_eval_trials',
+    4,
+    'Number of times each task is run during evaluation. Higher values give'
+    ' more stable evaluation metrics but increase runtime. Recommended: 4-8.',
+)
+_MAX_CONCURRENCY = flags.DEFINE_integer(
+    'max_concurrency',
+    8,
+    'Maximum number of parallel agent-environment interactions. Increase if'
+    ' you have sufficient API quota.',
+)
+_EVAL_MODE = flags.DEFINE_bool(
+    'eval_mode',
+    False,
+    'If set, run evaluation only using the seed prompt, skipping GEPA'
+    ' optimization.',
+)
+_USE_RATER = flags.DEFINE_bool(
+    'use_rater',
+    False,
+    'If set, use an LLM rater to score trajectories.',
+)
+_TRAIN_BATCH_SIZE = flags.DEFINE_integer(
+    'train_batch_size',
+    3,
+    'Number of trajectories sampled from rollouts to be used by the'
+    ' reflection model in each GEPA step to generate prompt improvements.'
+    ' Increasing the batch size may help provide a more stable signal and'
+    ' estimate of a prompt quality but entails higher cost. One can start with'
+    ' a low value and increase the size if significant variations are'
+    ' observed.',
+)
+
+
+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(gepa_utils.FilterInferenceWarnings())
+  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)