# ----------------------------------------------------------------------------- # This file re-implements algorithms from the DSPy project: # Repo: https://github.com/stanfordnlp/dspy # Paper: "DSPy: Compiling Declarative Language Model Calls into Self-Improving Pipelines" # Authors: Khattab et al. # # Re-implementation integrated into EvoAgentX with permission from the authors. # All mistakes or modifications are our own. # # Code of Conduct: This project follows the Microsoft Open Source Code of Conduct. # https://opensource.microsoft.com/codeofconduct/ # ----------------------------------------------------------------------------- import os import tqdm # import types import inspect import threading from copy import deepcopy from functools import wraps from collections import defaultdict import optuna from typing import Optional, Callable, Literal, List, Any, Dict, Union, Tuple, Set import dspy from dspy import MIPROv2 from dspy.clients import LM, Provider from dspy.utils.callback import BaseCallback from dspy.utils.parallelizer import ParallelExecutor from dspy.propose.grounded_proposer import GroundedProposer from dspy.teleprompt.utils import ( create_n_fewshot_demo_sets, get_signature, create_minibatch, print_full_program, save_candidate_program, get_program_with_highest_avg_score ) from ..core.logging import logger from ..core.callbacks import suppress_cost_logging, suppress_logger_info from ..models.base_model import BaseLLM from ..benchmark.benchmark import Benchmark from .engine.base import BaseOptimizer from .engine.registry import ParamRegistry from ..utils.mipro_utils.register_utils import MiproRegistry from ..agents.agent_manager import AgentManager from ..workflow.workflow_graph import WorkFlowGraph from ..workflow.workflow import WorkFlow from ..evaluators.evaluator import Evaluator from ..prompts.template import PromptTemplate, MiproPromptTemplate # from ..utils.mipro_utils.signature_utils import signature_from_registry from ..utils.mipro_utils.module_utils import PromptTuningModule # Constants BOOTSTRAPPED_FEWSHOT_EXAMPLES_IN_CONTEXT = 3 LABELED_FEWSHOT_EXAMPLES_IN_CONTEXT = 0 MIN_MINIBATCH_SIZE = 50 AUTO_RUN_SETTINGS = { "light": {"n": 6, "val_size": 100}, "medium": {"n": 12, "val_size": 300}, "heavy": {"n": 18, "val_size": 1000}, } # ANSI escape codes for colors YELLOW = "\033[93m" GREEN = "\033[92m" BLUE = "\033[94m" BOLD = "\033[1m" ENDC = "\033[0m" # Resets the color to default class MiproLMWrapper(LM): """ A wrapper class for the LLM model. It converts the BaseLLM model in EvoAgentX to a dspy.LM object. """ def __init__( self, model: BaseLLM, model_type: Literal["chat", "text"] = "chat", temperature: float = 0.0, max_tokens: int = 4000, cache: bool = True, cache_in_memory: bool = True, callbacks: Optional[List[BaseCallback]] = None, num_retries: int = 3, provider=None, finetuning_model: Optional[str] = None, launch_kwargs: Optional[dict[str, Any]] = None, train_kwargs: Optional[dict[str, Any]] = None, **kwargs, ): self.model = model self.model_type = model_type self.cache = cache self.cache_in_memory = cache_in_memory self.callbacks = callbacks or [] self.history = [] self.provider = provider or Provider() self.num_retries = num_retries self.finetuning_model = finetuning_model self.launch_kwargs = launch_kwargs or {} self.train_kwargs = train_kwargs or {} self.kwargs = dict(temperature=temperature, max_tokens=max_tokens, **kwargs) def forward(self, prompt=None, messages=None, **kwargs): response = self.model.generate(prompt=prompt, messages=messages, **kwargs) return [response.content] def __call__(self, prompt=None, messages=None, **kwargs): return self.forward(prompt=prompt, messages=messages, **kwargs) def copy(self, **kwargs): new_config = deepcopy(self.model.config) new_kwargs = {} for key, value in kwargs.items(): if hasattr(new_config, key): setattr(new_config, key, value) if (key in self.kwargs) or (not hasattr(self, key)): new_kwargs[key] = value new_model = self.model.__class__(config=new_config) return MiproLMWrapper(new_model, **new_kwargs) def generate(self, *args, **kwargs): # to be compatible with BaseLLM.generate() return self.model.generate(*args, **kwargs) async def async_generate(self, *args, **kwargs): # to be compatible with BaseLLM.async_generate() return await self.model.async_generate(*args, **kwargs) class MiproEvaluator: def __init__( self, benchmark: Benchmark, num_threads: Optional[int] = None, display_progress: Optional[bool] = None, max_errors: int = 5, return_all_scores: bool = False, return_outputs: bool = False, provide_traceback: bool = False, failure_score: float = 0.0, metric_name: Optional[str] = None, **kwargs ): self.benchmark = benchmark self.num_threads = num_threads self.display_progress = display_progress self.max_errors = max_errors self.return_all_scores = return_all_scores self.return_outputs = return_outputs self.provide_traceback = provide_traceback self.failure_score = failure_score self.metric_name = metric_name self.kwargs = kwargs # Add a thread-safe counter for logging self._log_counter = 0 self._log_lock = threading.Lock() def _extract_score_from_dict(self, score_dict: Dict[str, float]) -> float: """Extract a single score from a dictionary of scores. Args: score_dict (Dict[str, float]): Dictionary containing metric scores Returns: float: The extracted score based on the following rules: 1. If dict has only one score, return that score 2. If metric_name is specified, return that metric's score 3. Otherwise, return average of all scores """ if len(score_dict) == 1: return list(score_dict.values())[0] elif self.metric_name is not None: return score_dict[self.metric_name] else: avg_score = sum(score_dict.values()) / len(score_dict) # Use thread-safe counter to ensure message is only logged once with self._log_lock: if self._log_counter == 0: logger.info(f"`{type(self.benchmark)}.evaluate` returned a dictionary of scores, but no metric name was provided. Will return the average score across all metrics.") self._log_counter += 1 return avg_score def metric(self, example: dspy.Example, prediction: Any, *args, **kwargs): if isinstance(self.benchmark.get_train_data()[0], dspy.Example): # the data in original benchmark is a dspy.Example score = self.benchmark.evaluate( prediction=prediction, label=self.benchmark.get_label(example) ) elif isinstance(self.benchmark.get_train_data()[0], dict): # the data in original benchmark is a dict, convert the dspy.Example to a dict score = self.benchmark.evaluate( prediction=prediction, label=self.benchmark.get_label(example.toDict()) # convert the dspy.Example to a dict ) else: raise ValueError(f"Unsupported example type in `{type(self.benchmark)}`! Expected `dspy.Example` or `dict`, got {type(self.benchmark.get_train_data()[0])}") if isinstance(score, dict): score = self._extract_score_from_dict(score) return score def __call__( self, program: Callable, evalset: List[Any], **kwargs, ) -> float: return_all_scores = kwargs.get("return_all_scores", None) or self.return_all_scores return_outputs = kwargs.get("return_outputs", None) or self.return_outputs tqdm.tqdm._instances.clear() # Get the current suppress_cost_logs value from ..core.callbacks import suppress_cost_logs current_suppress_cost = suppress_cost_logs.get() if self.num_threads and self.num_threads > 1: executor = ParallelExecutor( num_threads=self.num_threads, disable_progress_bar=not self.display_progress, max_errors=self.max_errors, provide_traceback=self.provide_traceback, compare_results=True, ) else: executor = None def process_item(example): # Set the suppress_cost_logs context in the worker thread token = suppress_cost_logs.set(current_suppress_cost) try: if not isinstance(example, dspy.Example): raise ValueError(f"Example from benchmark must be a dspy.Example object, got {type(example)}") try: # prediction = program(**example) # score = metric(example, prediction) prediction = program(**example.inputs()) score = self.metric(example, prediction) # score = self.benchmark.evaluate(prediction=prediction, label=self.benchmark.get_label(example)) except Exception as e: logger.error(f"Error evaluating example {example}: {e}") return None, self.failure_score # Increment assert and suggest failures to program's attributes if hasattr(program, "_assert_failures"): program._assert_failures += dspy.settings.get("assert_failures") if hasattr(program, "_suggest_failures"): program._suggest_failures += dspy.settings.get("suggest_failures") return prediction, score finally: # Reset the context suppress_cost_logs.reset(token) if executor: results = executor.execute(process_item, evalset) else: # Use tqdm for single-threaded execution results = [] pbar = tqdm.tqdm( total=len(evalset), dynamic_ncols=True, disable=not self.display_progress, desc="Processing examples" ) for example in evalset: result = process_item(example) results.append(result) # Update progress bar with current score if available if result and result[1] is not None: current_scores = [r[1] for r in results if r and r[1] is not None] avg_score = sum(current_scores) / len(current_scores) if current_scores else 0 pbar.set_description(f"Average Metric: {avg_score:.2f}") pbar.update(1) pbar.close() assert len(evalset) == len(results) results = [(example, prediction, score) for example, (prediction, score) in zip(evalset, results)] ncorrect, ntotal = sum(score for *_, score in results), len(evalset) logger.info(f"Average Metric: {ncorrect} / {ntotal} ({round(100 * ncorrect / ntotal, 1)}%)") if return_all_scores and return_outputs: return round(100 * ncorrect / ntotal, 2), results, [score for *_, score in results] if return_all_scores: return round(100 * ncorrect / ntotal, 2), [score for *_, score in results] if return_outputs: return round(100 * ncorrect / ntotal, 2), results return round(100 * ncorrect / ntotal, 2) class MiproOptimizer(BaseOptimizer, MIPROv2): def __init__( self, registry: ParamRegistry, program: Callable, optimizer_llm: BaseLLM, evaluator: Optional[Callable] = None, eval_rounds: Optional[int] = 1, metric_threshold: Optional[float] = None, max_bootstrapped_demos: int = 4, max_labeled_demos: int = 4, auto: Optional[Literal["light", "medium", "heavy"]] = "medium", max_steps: int = None, num_candidates: Optional[int] = None, num_threads: Optional[int] = None, max_errors: int = 10, seed: int = 9, init_temperature: float = 0.5, track_stats: bool = True, save_path: Optional[str] = None, minibatch: bool = True, minibatch_size: int = 35, minibatch_full_eval_steps: int = 5, program_aware_proposer: bool = True, data_aware_proposer: bool = True, view_data_batch_size: int = 10, tip_aware_proposer: bool = True, fewshot_aware_proposer: bool = True, requires_permission_to_run: bool = False, provide_traceback: Optional[bool] = None, verbose: bool = False, **kwargs ): """ Base MiproOptimizer class that supports plug-and-play usage. Args: registry (ParamRegistry): a ParamRegistry object that contains the parameters to optimize. program (Callable): a program to optimize. Must be a callable object with save(path) and load(path) methods. optimizer_llm (BaseLLM): a language model to use for optimization. evaluator (Optional[Callable]): a function that evaluates the performance of the program. Required to have a `__call__(program, evalset, *kwargs) -> float` method that receives a program and a list of examples from a benchmark's train/dev/test set and return a float score. Must also have a `metric(example, prediction) -> float` method that evaluates a single example. If not provided, will construct a default evaluator using the benchmark's evaluate method. eval_rounds (Optional[int]): number of rounds to evaluate the program. Defaults to 1. metric_threshold (Optional[float]): threshold for the metric score. If provided, only examples with scores above this threshold will be used as demonstrations. If not provided, examples with scores above 0 will be used as demonstrations. max_bootstrapped_demos (int): maximum number of bootstrapped demonstrations to use. Defaults to 4. max_labeled_demos (int): maximum number of labeled demonstrations to use. Defaults to 4. auto (Optional[Literal["light", "medium", "heavy"]]): automatic configuration mode. If set, will override num_candidates and max_steps. "light": n=6, val_size=100; "medium": n=12, val_size=300; "heavy": n=18, val_size=1000. Defaults to "medium". max_steps (int): maximum number of optimization steps. Required if auto is None. num_candidates (Optional[int]): number of candidates to generate for each optimization step. Required if auto is None. num_threads (Optional[int]): number of threads to use for parallel evaluation. If None, will use single thread. Only used if evaluator is not provided. max_errors (int): maximum number of errors allowed during evaluation before stopping. Defaults to 10. seed (int): random seed for reproducibility. Defaults to 9. init_temperature (float): initial temperature for instruction generation. Defaults to 0.5. track_stats (bool): whether to track optimization statistics. Defaults to True. save_path (Optional[str]): path to save optimization results. If None, results will not be saved. minibatch (bool): whether to use minibatch evaluation during optimization. Defaults to True. minibatch_size (int): size of minibatch for evaluation. Defaults to 35. minibatch_full_eval_steps (int): number of minibatch steps between full evaluations. Defaults to 5. program_aware_proposer (bool): whether to use program-aware instruction proposer. Defaults to True. data_aware_proposer (bool): whether to use data-aware instruction proposer. Defaults to True. view_data_batch_size (int): batch size for viewing data during instruction proposal. Defaults to 10. tip_aware_proposer (bool): whether to use tip-aware instruction proposer. Defaults to True. fewshot_aware_proposer (bool): whether to use fewshot-aware instruction proposer. Defaults to True. requires_permission_to_run (bool): whether to require user permission before running optimization. Defaults to False. provide_traceback (Optional[bool]): whether to provide traceback for evaluation errors. If None, will use default setting. **kwargs: additional keyword arguments to pass to the evaluator. Raises: TypeError: If program is not callable or evaluator doesn't return float ValueError: If program doesn't have required methods (save and load) or if evaluator doesn't have required methods """ # initialize base optimizer BaseOptimizer.__init__(self, registry=registry, program=program, evaluator=evaluator) # convert the registry and program to dspy-compatible module self._validate_program(program=program) self.model = self._convert_to_dspy_module(registry, program) self.optimizer_llm = MiproLMWrapper(optimizer_llm) dspy.configure(lm=self.optimizer_llm) self.task_model = dspy.settings.lm self.prompt_model = dspy.settings.lm self.metric_threshold = metric_threshold self.metric_name = None self.teacher_settings = {"use_teacher": True} # Validate 'auto' parameter allowed_modes = {None, "light", "medium", "heavy"} if auto not in allowed_modes: raise ValueError(f"Invalid value for auto: {auto}. Must be one of {allowed_modes}.") self.auto = auto self.num_fewshot_candidates = num_candidates self.num_instruct_candidates = num_candidates self.num_candidates = num_candidates self.init_temperature = init_temperature self.max_bootstrapped_demos = max_bootstrapped_demos self.max_labeled_demos = max_labeled_demos self.max_steps = max_steps self.num_threads = num_threads self.max_errors = max_errors self.track_stats = track_stats self.eval_rounds = eval_rounds self.save_path = save_path self.prompt_model_total_calls = 0 self.total_calls = 0 self.seed = seed self.rng = None self.minibatch = minibatch self.minibatch_size = minibatch_size self.minibatch_full_eval_steps = minibatch_full_eval_steps self.program_aware_proposer = program_aware_proposer self.data_aware_proposer = data_aware_proposer self.view_data_batch_size = view_data_batch_size self.tip_aware_proposer = tip_aware_proposer self.fewshot_aware_proposer = fewshot_aware_proposer self.requires_permission_to_run = requires_permission_to_run self.provide_traceback = provide_traceback self.verbose = verbose self.kwargs = kwargs def _validate_program(self, program: Callable): """ Validate that the program meets the required interface. Args: program (Callable): The program to validate Raises: TypeError: If program is not callable ValueError: If program doesn't have required methods (save and load) """ if not callable(program): raise TypeError("program must be callable") # Check if program has save method if not hasattr(program, 'save'): # raise ValueError("program must have a `save` method") logger.warning("program does not have a `save(path=...)` method, will use the default save method in dspy.Module") else: # Check save method signature save_sig = inspect.signature(program.save) save_params = list(save_sig.parameters.keys()) if 'path' not in save_params: raise ValueError("program.save must accept a 'path' parameter") # Check if program has load method if not hasattr(program, 'load'): # raise ValueError("program must have a `load` method") logger.warning("program does not have a `load(path=...)` method, will use the default load method in dspy.Module") else: # Check load method signature load_sig = inspect.signature(program.load) load_params = list(load_sig.parameters.keys()) if 'path' not in load_params: raise ValueError("program.load must accept a 'path' parameter") def _validate_evaluator(self, evaluator: Callable = None, benchmark: Benchmark = None, metric_name: Optional[str] = None) -> Callable: """ Validate that the evaluator meets the required interface and wrap it with runtime checks. Args: evaluator (Callable): The evaluator to validate. If provided, it must have a `__call__(program, evalset, *kwargs) -> float` method that receives a program and a list of examples from a benchmark's train/dev/test set and return a float score. It must also have a `metric(example: dspy.Example, prediction: Any) -> float/int/bool` method that evaluates a single example. benchmark (Benchmark): The benchmark to use for evaluation. Only used if evaluator is not provided. In this case, the evaluator will be constructed using the `evaluate` method (return a dictionary of scores) in the benchmark. metric_name (Optional[str]): The name of the metric to use for evaluation. Only used if evaluator is not provided. It will be used to select the metric for optimization from the dictionary of scores returned by the benchmark's `evaluate` method. Raises: TypeError: If evaluator is not callable or doesn't return float ValueError: If evaluator doesn't have required parameters """ if evaluator is None: if not hasattr(benchmark, "evaluate"): raise ValueError("`evaluator` is not provided and the benchmark does not have a `evaluate` method.") logger.info("`evaluator` is not provided. Will construct a default evaluator using the `evaluate` method in the benchmark.") evaluator = MiproEvaluator( benchmark=benchmark, num_threads=self.num_threads, max_errors=self.max_errors, display_progress=True, provide_traceback=self.provide_traceback, metric_name=metric_name, **self.kwargs ) if not callable(evaluator): raise TypeError("evaluator must be callable, i.e., a function or a class with interface `__call__(program, evalset, *kwargs) -> float`") # Check if evaluator has __call__ method with correct signature sig = inspect.signature(evaluator.__call__ if hasattr(evaluator, '__call__') else evaluator) params = list(sig.parameters.keys()) if len(params) < 2: raise ValueError("evaluator must accept at least two parameters (program and evalset)") # Check return type annotation if available if sig.return_annotation != inspect.Signature.empty: if sig.return_annotation not in [float, int, bool]: raise TypeError("evaluator must return a float, int, or bool") # check if the evaluator has a `metric` method with correct signature if not hasattr(evaluator, 'metric'): raise ValueError("evaluator must have a `metric(example: dspy.Example, prediction: Any) -> float/int/bool` method") metric_sig = inspect.signature(evaluator.metric) metric_params = list(metric_sig.parameters.keys()) if len(metric_params) < 2: raise ValueError("evaluator.metric must accept at least two parameters (example and prediction)") if metric_params[0] != 'example' or metric_params[1] != 'prediction': raise ValueError("evaluator.metric must have parameters in order: example, prediction") # if '*args' not in str(metric_sig): # raise ValueError("evaluator.metric must accept *args") # Wrap the metric method with runtime checks # original_metric = evaluator.metric # @wraps(original_metric) # def wrapped_metric(example, prediction, *args, **kwargs): # result = original_metric(example, prediction, *args, **kwargs) # # Runtime check for return value # if not isinstance(result, (float, int, bool)): # raise TypeError(f"evaluator.metric must return a float, int, or bool, got {type(result)}") # return result # evaluator.metric = types.MethodType(wrapped_metric, evaluator) # Wrap the evaluator with runtime checks original_evaluator = evaluator.__call__ if hasattr(evaluator, '__call__') else evaluator @wraps(original_evaluator) def wrapped_evaluator(*args, **kwargs): result = original_evaluator(*args, **kwargs) # Runtime check for return value if not isinstance(result, (float, int, bool)): raise TypeError(f"evaluator must return a float, int, or bool, got {type(result)}") return result # Replace the evaluator with our wrapped version if hasattr(evaluator, '__call__'): evaluator.__call__ = wrapped_evaluator else: # If it's a function, we need to create a new callable object class WrappedEvaluator: def __init__(self, func): self._func = func def __call__(self, *args, **kwargs): return wrapped_evaluator(*args, **kwargs) return WrappedEvaluator(evaluator) return evaluator def _convert_to_dspy_module(self, registry: ParamRegistry, program: Callable): if isinstance(program, dspy.Module): return program program = PromptTuningModule.from_registry( program=program, registry=registry, ) return program def optimize(self, dataset: Benchmark, metric_name: Optional[str] = None, **kwargs): """ Optimize the program using the Mipro algorithm. Args: dataset (Benchmark): a Benchmark object that contains the training and validation data. metric_name (Optional[str]): the name of the metric to use for optimization. Only used when `self.evaluator` is not provided. In this case, the evaluator will be constructed using the `evaluate` method (return a dictionary of scores) in the benchmark, and the metric specified by `metric_name` will be used for optimization. If not provided, the average of all scores returned by the evaluator will be used. If `self.evaluator` is provided, this argument will be ignored. **kwargs: additional keyword arguments to pass to the evaluator. """ zeroshot_opt = (self.max_bootstrapped_demos == 0) and (self.max_labeled_demos == 0) student = self.model num_trials = self.max_steps minibatch = self.minibatch self.metric_name = metric_name # If auto is None, and num_trials is not provided (but num_candidates is), raise an error that suggests a good num_trials value if self.auto is None and (self.num_candidates is not None and num_trials is None): raise ValueError(f"If auto is None, max_steps must also be provided. Given num_candidates={self.num_candidates}, we'd recommend setting max_steps to ~{self._set_num_trials_from_num_candidates(self.model, zeroshot_opt, self.num_candidates)}.") # If auto is None, and num_candidates or num_trials is None, raise an error if self.auto is None and (self.num_candidates is None or num_trials is None): raise ValueError("If auto is None, num_candidates must also be provided.") # If auto is provided, and either num_candidates or num_trials is not None, raise an error if self.auto is not None and (self.num_candidates is not None or num_trials is not None): raise ValueError("If auto is not None, num_candidates and max_steps cannot be set, since they would be overrided by the auto settings. Please either set auto to None, or do not specify num_candidates and max_steps.") # Set random seeds seed = self.seed self._set_random_seeds(seed) # Set training & validation sets trainset, valset = self._set_and_validate_datasets(dataset=dataset) # Set hyperparameters based on run mode (if set) num_trials, valset, minibatch = self._set_hyperparams_from_run_mode( student, num_trials, minibatch, zeroshot_opt, valset ) if self.auto: self._print_auto_run_settings(num_trials, minibatch, valset) if minibatch and self.minibatch_size > len(valset): raise ValueError(f"Minibatch size cannot exceed the size of the valset. Valset size: {len(valset)}.") # # Estimate LM calls and get user confirmation if self.requires_permission_to_run: if not self._get_user_confirmation( student, num_trials, minibatch, self.minibatch_size, self.minibatch_full_eval_steps, valset, self.program_aware_proposer, ): logger.info("Compilation aborted by the user.") return student # Return the original student program program = student.deepcopy() # check the evaluator (If None, will construct a default evaluator using the `evaluate` method in the benchmark) and wrap it with runtime checks evaluator = self._validate_evaluator(evaluator=self.evaluator, benchmark=dataset, metric_name=metric_name) self.metric = evaluator.metric # Step 1: Bootstrap few-shot examples demo_candidates = self._bootstrap_fewshot_examples(program, trainset, seed, teacher=None) # Step 2: Propose instruction candidates with suppress_cost_logging(): instruction_candidates = self._propose_instructions( program, trainset, demo_candidates, self.view_data_batch_size, self.program_aware_proposer, self.data_aware_proposer, self.tip_aware_proposer, self.fewshot_aware_proposer, ) # Step 3: Find optimal prompt parameters with suppress_cost_logging(): best_program = self._optimize_prompt_parameters( program, instruction_candidates, demo_candidates, evaluator, valset, num_trials, minibatch, self.minibatch_size, self.minibatch_full_eval_steps, seed, ) if self.save_path: os.makedirs(self.save_path, exist_ok=True) self.best_program_path = os.path.join(self.save_path, "best_program.json") best_program.save(self.best_program_path) # reset the self.model. After optimization, the model will be reset to the original state. # This is necessary to avoid the model being modified by the optimization process. # Use self.restore_best_program() to restore the best program. self.model.reset() def restore_best_program(self): # todo: implement this pass def _get_input_keys(self, dataset: Benchmark) -> Optional[List[str]]: input_keys = None if hasattr(dataset, "get_input_keys"): candidate_input_keys = dataset.get_input_keys() if isinstance(candidate_input_keys, (list, tuple)) and all(isinstance(key, str) for key in candidate_input_keys): input_keys = candidate_input_keys return input_keys def _set_and_validate_datasets(self, dataset: Benchmark): trainset = dataset.get_train_data() if not trainset: raise ValueError("No training data found in the dataset. Please set `_train_data` in the benchmark.") if trainset and not isinstance(trainset[0], (dict, dspy.Example)): raise ValueError("Training set in the benchmark must be a list of dictionaries or dspy.Example objects.") valset = dataset.get_dev_data() if not valset: if len(trainset) < 2: raise ValueError("Training set in the benchmark must have at least 2 examples if no validation set is provided.") valset_size = min(1000, max(1, int(len(trainset) * 0.80))) cutoff = len(trainset) - valset_size valset = trainset[cutoff:] trainset = trainset[:cutoff] else: if len(valset) < 1: raise ValueError("Validation set in the benchmark must have at least 1 example.") # convert the trainset and valset to a list of dspy Example input_keys = self._get_input_keys(dataset) if input_keys is None: logger.warning("`get_input_keys` is not implemented in the benchmark. Will use all keys as input keys. This may cause unexpected behavior if the program does not use all the keys.") input_keys = trainset[0].keys() dspy_trainset = self._convert_benchmark_data_to_dspy_examples(trainset, input_keys) dspy_valset = self._convert_benchmark_data_to_dspy_examples(valset, input_keys) return dspy_trainset, dspy_valset def _convert_benchmark_data_to_dspy_examples(self, data: List[dict], input_keys: List[str]) -> List[dspy.Example]: """ Convert the benchmark data to a list of dspy Example. This is required since the evaluator accepts a list of dspy Example. """ dspy_examples = [ example.with_inputs(*input_keys) if isinstance(example, dspy.Example) else dspy.Example(**example).with_inputs(*input_keys) for example in data ] return dspy_examples def _bootstrap_fewshot_examples(self, program: Any, trainset: List, seed: int, teacher: Any) -> Optional[List]: logger.info("==> STEP 1: BOOTSTRAP FEWSHOT EXAMPLES <==") if self.max_bootstrapped_demos > 0: logger.info( "These will be used as few-shot example candidates for our program and for creating instructions.\n" ) else: logger.info("These will be used for informing instruction proposal.\n") logger.info(f"Bootstrapping N={self.num_fewshot_candidates} sets of demonstrations...") zeroshot = self.max_bootstrapped_demos == 0 and self.max_labeled_demos == 0 try: with suppress_logger_info(): demo_candidates = create_n_fewshot_demo_sets( student=program, num_candidate_sets=self.num_fewshot_candidates, trainset=trainset, max_labeled_demos=(LABELED_FEWSHOT_EXAMPLES_IN_CONTEXT if zeroshot else self.max_labeled_demos), max_bootstrapped_demos=( BOOTSTRAPPED_FEWSHOT_EXAMPLES_IN_CONTEXT if zeroshot else self.max_bootstrapped_demos ), metric=self.metric, max_errors=self.max_errors, teacher=teacher, teacher_settings=self.teacher_settings, seed=seed, metric_threshold=self.metric_threshold, rng=self.rng, ) except Exception as e: logger.info(f"Error generating few-shot examples: {e}") logger.info("Running without few-shot examples.") demo_candidates = None return demo_candidates def _propose_instructions( self, program: Any, trainset: List, demo_candidates: Optional[List], view_data_batch_size: int, program_aware_proposer: bool, data_aware_proposer: bool, tip_aware_proposer: bool, fewshot_aware_proposer: bool, ) -> Dict[int, List[str]]: logger.info("==> STEP 2: PROPOSE INSTRUCTION CANDIDATES <==") logger.info( "We will use the few-shot examples from the previous step, a generated dataset summary, a summary of the program code, and a randomly selected prompting tip to propose instructions." ) proposer = GroundedProposer( program=program, trainset=trainset, prompt_model=self.prompt_model, view_data_batch_size=view_data_batch_size, program_aware=program_aware_proposer, use_dataset_summary=data_aware_proposer, use_task_demos=fewshot_aware_proposer, num_demos_in_context=BOOTSTRAPPED_FEWSHOT_EXAMPLES_IN_CONTEXT, use_tip=tip_aware_proposer, set_tip_randomly=tip_aware_proposer, use_instruct_history=False, set_history_randomly=False, verbose=self.verbose, rng=self.rng, ) logger.info(f"Proposing N={self.num_instruct_candidates} instructions...") instruction_candidates = proposer.propose_instructions_for_program( trainset=trainset, program=program, demo_candidates=demo_candidates, N=self.num_instruct_candidates, T=self.init_temperature, trial_logs={}, ) for i, pred in enumerate(program.predicts): logger.info(f"Proposed Instructions for Predictor {i}:\n") instruction_candidates[i][0] = get_signature(pred).instructions for j, instruction in enumerate(instruction_candidates[i]): logger.info(f"{j}: {instruction}\n") logger.info("\n") return instruction_candidates def _optimize_prompt_parameters( self, program: Any, instruction_candidates: Dict[int, List[str]], demo_candidates: Optional[List], evaluator: Callable, valset: List, num_trials: int, minibatch: bool, minibatch_size: int, minibatch_full_eval_steps: int, seed: int, ) -> Optional[Any]: # Run optimization optuna.logging.set_verbosity(optuna.logging.WARNING) logger.info("==> STEP 3: FINDING OPTIMAL PROMPT PARAMETERS <==") logger.info( "We will evaluate the program over a series of trials with different combinations of instructions and few-shot examples to find the optimal combination using Bayesian Optimization.\n" ) # Compute the adjusted total trials that we will run (including full evals) run_additional_full_eval_at_end = 1 if num_trials % minibatch_full_eval_steps != 0 else 0 adjusted_num_trials = int((num_trials + num_trials // minibatch_full_eval_steps + 1 + run_additional_full_eval_at_end) if minibatch else num_trials) logger.info(f"== Trial {1} / {adjusted_num_trials} - Full Evaluation of Default Program ==") # default_score = eval_candidate_program( # len(valset), valset, program, evaluator, self.rng, # ) default_score = self.evaluate( evalset=valset, program=program, evaluator=evaluator, batch_size=len(valset) ) logger.info(f"Default program score: {default_score}\n") trial_logs = {} trial_logs[1] = {} trial_logs[1]["full_eval_program_path"] = save_candidate_program(program, self.save_path, -1) trial_logs[1]["full_eval_score"] = default_score trial_logs[1]["total_eval_calls_so_far"] = len(valset) trial_logs[1]["full_eval_program"] = program.deepcopy() # Initialize optimization variables best_score = default_score best_program = program.deepcopy() total_eval_calls = len(valset) score_data = [{"score": best_score, "program": program.deepcopy(), "full_eval": True}] param_score_dict = defaultdict(list) fully_evaled_param_combos = {} # Define the objective function def objective(trial): nonlocal program, best_program, best_score, trial_logs, total_eval_calls, score_data trial_num = trial.number + 1 if minibatch: logger.info(f"== Trial {trial_num} / {adjusted_num_trials} - Minibatch ==") else: logger.info(f"===== Trial {trial_num} / {num_trials} =====") trial_logs[trial_num] = {} # Create a new candidate program candidate_program = program.deepcopy() # Choose instructions and demos, insert them into the program chosen_params, raw_chosen_params = self._select_and_insert_instructions_and_demos( candidate_program, instruction_candidates, demo_candidates, trial, trial_logs, trial_num, ) # Log assembled program if self.verbose: logger.info("Evaluating the following candidate program...\n") print_full_program(candidate_program) # Evaluate the candidate program (on minibatch if minibatch=True) batch_size = minibatch_size if minibatch else len(valset) # score = eval_candidate_program(batch_size, valset, candidate_program, evaluator, self.rng) score = self.evaluate( evalset=valset, program=candidate_program, evaluator=evaluator, batch_size=batch_size ) total_eval_calls += batch_size # Update best score and program if not minibatch and score > best_score: best_score = score best_program = candidate_program.deepcopy() logger.info(f"{GREEN}Best full score so far!{ENDC} Score: {score}") # Log evaluation results score_data.append( {"score": score, "program": candidate_program, "full_eval": batch_size >= len(valset)} ) # score, prog, full_eval if minibatch: self._log_minibatch_eval( score, best_score, batch_size, chosen_params, score_data, trial, adjusted_num_trials, trial_logs, trial_num, candidate_program, total_eval_calls, ) else: self._log_normal_eval( score, best_score, chosen_params, score_data, trial, num_trials, trial_logs, trial_num, valset, batch_size, candidate_program, total_eval_calls, ) categorical_key = ",".join(map(str, chosen_params)) param_score_dict[categorical_key].append( (score, candidate_program, raw_chosen_params), ) # If minibatch, perform full evaluation at intervals (and at the very end) if minibatch and ((trial_num % (minibatch_full_eval_steps+1) == 0) or (trial_num == (adjusted_num_trials-1))): best_score, best_program, total_eval_calls = self._perform_full_evaluation( trial_num, adjusted_num_trials, param_score_dict, fully_evaled_param_combos, evaluator, valset, trial_logs, total_eval_calls, score_data, best_score, best_program, study, instruction_candidates, demo_candidates, ) return score sampler = optuna.samplers.TPESampler(seed=seed, multivariate=True) study = optuna.create_study(direction="maximize", sampler=sampler) default_params = {f"{i}_predictor_instruction": 0 for i in range(len(program.predicts))} if demo_candidates: default_params.update({f"{i}_predictor_demos": 0 for i in range(len(program.predicts))}) # Add default run as a baseline in optuna (TODO: figure out how to weight this by # of samples evaluated on) trial = optuna.trial.create_trial( params=default_params, distributions=self._get_param_distributions(program, instruction_candidates, demo_candidates), value=default_score, ) study.add_trial(trial) study.optimize(objective, n_trials=num_trials) # Attach logs to best program if best_program is not None and self.track_stats: best_program.trial_logs = trial_logs best_program.score = best_score best_program.prompt_model_total_calls = self.prompt_model_total_calls best_program.total_calls = self.total_calls sorted_candidate_programs = sorted(score_data, key=lambda x: x["score"], reverse=True) # Attach all minibatch programs best_program.mb_candidate_programs = [ score_data for score_data in sorted_candidate_programs if not score_data["full_eval"] ] # Attach all programs that were evaluated on the full trainset, in descending order of score best_program.candidate_programs = [ score_data for score_data in sorted_candidate_programs if score_data["full_eval"] ] logger.info(f"Returning best identified program with score {best_score}!") return best_program def _select_and_insert_instructions_and_demos( self, candidate_program: Any, instruction_candidates: Dict[int, List[str]], demo_candidates: Optional[List], trial: optuna.trial.Trial, trial_logs: Dict, trial_num: int, ) -> List[str]: chosen_params = [] raw_chosen_params = {} for i, predictor in enumerate(candidate_program.predictors()): # Select instruction instruction_idx = trial.suggest_categorical( f"{i}_predictor_instruction", range(len(instruction_candidates[i])) ) selected_instruction = instruction_candidates[i][instruction_idx] # updated_signature = get_signature(predictor).with_instructions(selected_instruction) # set_signature(predictor, updated_signature) predictor.signature.instructions = selected_instruction trial_logs[trial_num][f"{i}_predictor_instruction"] = instruction_idx chosen_params.append(f"Predictor {i}: Instruction {instruction_idx}") raw_chosen_params[f"{i}_predictor_instruction"] = instruction_idx # Select demos if available if demo_candidates: demos_idx = trial.suggest_categorical(f"{i}_predictor_demos", range(len(demo_candidates[i]))) predictor.demos = demo_candidates[i][demos_idx] trial_logs[trial_num][f"{i}_predictor_demos"] = demos_idx chosen_params.append(f"Predictor {i}: Few-Shot Set {demos_idx}") raw_chosen_params[f"{i}_predictor_demos"] = instruction_idx return chosen_params, raw_chosen_params def _log_minibatch_eval( self, score, best_score, batch_size, chosen_params, score_data, trial, adjusted_num_trials, trial_logs, trial_num, candidate_program, total_eval_calls, ): trial_logs[trial_num]["mb_program_path"] = save_candidate_program(candidate_program, self.save_path, trial_num=trial_num, note="mb") trial_logs[trial_num]["mb_score"] = score trial_logs[trial_num]["total_eval_calls_so_far"] = total_eval_calls trial_logs[trial_num]["mb_program"] = candidate_program.deepcopy() logger.info(f"Score: {score} on minibatch of size {batch_size} with parameters {chosen_params}.") minibatch_scores = ", ".join([f"{s['score']}" for s in score_data if not s["full_eval"]]) logger.info(f"Minibatch scores so far: {'[' + minibatch_scores + ']'}") full_eval_scores = ", ".join([f"{s['score']}" for s in score_data if s["full_eval"]]) trajectory = "[" + full_eval_scores + "]" logger.info(f"Full eval scores so far: {trajectory}") logger.info(f"Best full score so far: {best_score}") logger.info( f"{'=' * len(f'== Trial {trial.number + 1} / {adjusted_num_trials} - Minibatch Evaluation ==')}\n\n" ) def _log_normal_eval( self, score, best_score, chosen_params, score_data, trial, num_trials, trial_logs, trial_num, valset, batch_size, candidate_program, total_eval_calls, ): trial_logs[trial_num]["full_eval_program_path"] = save_candidate_program( candidate_program, self.save_path, trial_num ) trial_logs[trial_num]["full_eval_score"] = score trial_logs[trial_num]["total_eval_calls_so_far"] = total_eval_calls trial_logs[trial_num]["full_eval_program"] = candidate_program.deepcopy() logger.info(f"Score: {score} with parameters {chosen_params}.") full_eval_scores = ", ".join([f"{s['score']}" for s in score_data if s["full_eval"]]) logger.info(f"Scores so far: {'[' + full_eval_scores + ']'}") logger.info(f"Best score so far: {best_score}") logger.info(f"{'=' * len(f'===== Trial {trial.number + 1} / {num_trials} =====')}\n\n") def _perform_full_evaluation( self, trial_num: int, adjusted_num_trials: int, param_score_dict: Dict, fully_evaled_param_combos: Dict, evaluator: Callable, valset: List, trial_logs: Dict, total_eval_calls: int, score_data, best_score: float, best_program: Any, study: optuna.Study, instruction_candidates: List, demo_candidates: List, ): logger.info(f"===== Trial {trial_num + 1} / {adjusted_num_trials} - Full Evaluation =====") # Identify best program to evaluate fully highest_mean_program, mean_score, combo_key, params = get_program_with_highest_avg_score( param_score_dict, fully_evaled_param_combos ) logger.info(f"Doing full eval on next top averaging program (Avg Score: {mean_score}) from minibatch trials...") # full_eval_score = eval_candidate_program(len(valset), valset, highest_mean_program, evaluator, self.rng) full_eval_score = self.evaluate( evalset=valset, program=highest_mean_program, evaluator=evaluator, batch_size=len(valset) ) score_data.append({"score": full_eval_score, "program": highest_mean_program, "full_eval": True}) # Log full eval as a trial so that optuna can learn from the new results trial = optuna.trial.create_trial( params=params, distributions=self._get_param_distributions(best_program, instruction_candidates, demo_candidates), value=full_eval_score, ) study.add_trial(trial) # Log full evaluation results fully_evaled_param_combos[combo_key] = { "program": highest_mean_program, "score": full_eval_score, } total_eval_calls += len(valset) trial_logs[trial_num + 1] = {} trial_logs[trial_num + 1]["total_eval_calls_so_far"] = total_eval_calls trial_logs[trial_num + 1]["full_eval_program_path"] = save_candidate_program( program=highest_mean_program, log_dir=self.save_path, trial_num=trial_num + 1, note="full_eval", ) trial_logs[trial_num + 1]["full_eval_program"] = highest_mean_program trial_logs[trial_num + 1]["full_eval_score"] = full_eval_score # Update best score and program if necessary if full_eval_score > best_score: logger.info(f"{GREEN}New best full eval score!{ENDC} Score: {full_eval_score}") best_score = full_eval_score best_program = highest_mean_program.deepcopy() full_eval_scores = ", ".join([f"{s['score']}" for s in score_data if s["full_eval"]]) trajectory = "[" + full_eval_scores + "]" logger.info(f"Full eval scores so far: {trajectory}") logger.info(f"Best full score so far: {best_score}") logger.info(len(f"===== Full Eval {len(fully_evaled_param_combos) + 1} =====") * "=") logger.info("\n") return best_score, best_program, total_eval_calls def evaluate( self, evalset: Optional[List[dspy.Example]] = None, dataset: Optional[Benchmark] = None, eval_mode: Optional[str] = "dev", program: Optional[PromptTuningModule] = None, evaluator: Optional[Callable] = None, indices: Optional[List[int]] = None, sample_k: Optional[int] = None, batch_size: Optional[int] = None, # if provided, sample `batch_size` examples from the evalset **kwargs ): # if program is not provided, use the model as the program if program is None: program = self.model if evaluator is None: evaluator = self._validate_evaluator(evaluator=self.evaluator, benchmark=dataset, metric_name=self.metric_name) # if evalset is not provided, use the dataset to get the evalset if evalset is None: assert dataset is not None, "Either `evalset` or `dataset` must be provided." data_map = {"train": dataset.get_train_data, "dev": dataset.get_dev_data, "test": dataset.get_test_data} evaldata = data_map[eval_mode](indices=indices, sample_k=sample_k) if not evaldata: logger.warning(f"No data found for {eval_mode} set. Return 0.0.") return 0.0 input_keys = self._get_input_keys(dataset=dataset) if not input_keys: input_keys = evaldata[0].keys() evalset = self._convert_benchmark_data_to_dspy_examples(evaldata, input_keys) batch_size = batch_size or len(evalset) score_list = [] for _ in range(self.eval_rounds): score = eval_candidate_program( batch_size=batch_size, evalset=evalset, candidate_program=program, evaluator=evaluator, rng=self.rng ) score_list.append(score) return sum(score_list) / len(score_list) def eval_candidate_program( batch_size: int, evalset: list, candidate_program: Any, evaluator: Callable, rng = None, return_all_scores: bool = False, ) -> Union[float, Tuple[float, List[float]]]: try: if batch_size >= len(evalset): return evaluator( program=candidate_program, evalset=evalset, return_all_scores=return_all_scores, ) else: return evaluator( program=candidate_program, evalset=create_minibatch(evalset, batch_size, rng), return_all_scores=return_all_scores, ) except Exception as e: logger.error(f"An exception occurred during evaluation: {str(e)}", exc_info=True) if return_all_scores: return 0.0, [0.0] * len(evalset) return 0.0 class WorkFlowGraphProgram: def __init__( self, graph: WorkFlowGraph, agent_manager: AgentManager, executor_llm: BaseLLM, collate_func: Optional[Callable] = None, output_postprocess_func: Optional[Callable] = None, ): self.graph = graph self.agent_manager = agent_manager self.executor_llm = executor_llm self.collate_func = collate_func or (lambda x: x) self.output_postprocess_func = output_postprocess_func or (lambda x: x) def __call__(self, **input_data): new_config = deepcopy(self.graph.get_config()) new_graph: WorkFlowGraph = WorkFlowGraph.from_dict(new_config) new_graph.reset_graph() # execute the graph with WorkFlow use_teacher = dspy.settings.get("use_teacher", False) if use_teacher: # use teacher model to execute the graph, used for optimization new_graph, new_agent_manager = self.inject_teacher_settings(new_graph, self.agent_manager) workflow = WorkFlow(llm=self.executor_llm, graph=new_graph, agent_manager=new_agent_manager) else: # use the original executor llm to execute the graph workflow = WorkFlow(llm=self.executor_llm, graph=new_graph, agent_manager=self.agent_manager) output: str = workflow.execute(inputs=self.collate_func(input_data)) output = self.output_postprocess_func(output) # extract all the input and output data from the workflow execution all_execution_data = workflow.environment.execution_data all_input_output_keys = self._extract_input_output_keys(new_graph) execution_data = {k: v for k, v in all_execution_data.items() if k in all_input_output_keys} return output, execution_data def inject_teacher_settings(self, graph: WorkFlowGraph, agent_manager: AgentManager): """ Inject the teacher settings into the graph and agent manager. """ # dspy.settings.lm is configured in MiproOptimizer, which is a MiproLMWrapper instance optimizer_llm_config = dspy.settings.lm.model.config.to_dict() for node in graph.nodes: for agent in node.agents: agent["llm_config"] = optimizer_llm_config # create a new agent manager with the teacher settings new_agent_manager = agent_manager.copy() new_agent_manager.clear_agents() new_agent_manager.add_agents_from_workflow(graph, llm_config=optimizer_llm_config) return graph, new_agent_manager def _extract_input_output_keys(self, graph: WorkFlowGraph) -> Set[str]: """ Extract all the input and output keys from the graph. """ all_input_output_keys = set() for node in graph.nodes: for inp in node.inputs: all_input_output_keys.add(inp.name) for out in node.outputs: all_input_output_keys.add(out.name) for agent in node.agents: for agent_inp in agent.get("inputs", []): agent_inp_name = agent_inp.get("name", None) if agent_inp_name: all_input_output_keys.add(agent_inp_name) for agent_out in agent.get("outputs", []): agent_out_name = agent_out.get("name", None) if agent_out_name: all_input_output_keys.add(agent_out_name) return all_input_output_keys def save(self, path: str): self.graph.save_module(path=path) def load(self, path: str): return WorkFlowGraph.from_file(path=path) class MiproEvaluatorWrapper(MiproEvaluator): def __init__( self, evaluator: Evaluator, benchmark: Benchmark, metric_name: str = None, return_all_scores: bool = False, return_outputs: bool = False, ): self.evaluator = evaluator self.benchmark = benchmark self.metric_name = metric_name self.return_all_scores = return_all_scores self.return_outputs = return_outputs def metric(self, example: dspy.Example, prediction: Any, *args, **kwargs): return super().metric(example, prediction, *args, **kwargs) def __call__(self, program: PromptTuningModule, evalset: List[dspy.Example], **kwargs) -> float: # sync the candidate prompts and instructions to the workflow graph program.sync_predict_inputs_to_program() return_all_scores = kwargs.get("return_all_scores", None) or self.return_all_scores return_outputs = kwargs.get("return_outputs", None) or self.return_outputs if isinstance(program, PromptTuningModule): graph = program.program.graph elif isinstance(program, WorkFlowGraphProgram): graph = program.graph else: raise ValueError(f"Invalid program type: {type(program)}. Must be PromptTuningModule or WorkFlowGraphProgram.") self.evaluator._evaluation_records.clear() # update agents self.evaluator.agent_manager.update_agents_from_workflow(workflow_graph=graph, llm_config=self.evaluator.llm.config, **kwargs) if isinstance(self.benchmark.get_train_data()[0], dspy.Example): data = evalset else: data = [example.toDict() for example in evalset] with suppress_logger_info(): metrics = self.evaluator._evaluate_graph( graph=graph, data=data, benchmark=self.benchmark, verbose=True, **kwargs ) if isinstance(metrics, dict): score = self._extract_score_from_dict(metrics) else: score = metrics # extract all outputs and predictions all_scores, all_predictions = [], [] for example in data: example_id = self.benchmark.get_id(example=example) evaluation_record = self.evaluator._evaluation_records.get(example_id, None) if evaluation_record is None: all_scores.append(0.0) all_predictions.append(None) else: example_metrics = evaluation_record["metrics"] example_score = self._extract_score_from_dict(example_metrics) if isinstance(example_metrics, dict) else example_metrics all_scores.append(example_score) all_predictions.append(evaluation_record["prediction"]) if return_all_scores and return_outputs: return score, all_predictions, all_scores if return_all_scores: return score, all_scores if return_outputs: return score, all_predictions return score class WorkFlowMiproOptimizer(MiproOptimizer): def __init__( self, graph: WorkFlowGraph, evaluator: Evaluator, optimizer_llm: Optional[BaseLLM] = None, **kwargs, ): """ MiproOptimizer tailored for workflow graphs. Args: graph (WorkFlowGraph): the workflow graph to optimize. evaluator (Evaluator): the evaluator to use for the optimization. optimizer_llm (BaseLLM): the LLM to use for the optimization. If None, will use the LLM model in the evaluator. **kwargs: additional keyword arguments to pass to the MiproOptimizer. Available options: - metric_threshold (Optional[int]): threshold for the metric score. If provided, only examples with scores above this threshold will be used as demonstrations. - max_bootstrapped_demos (int): maximum number of bootstrapped demonstrations to use. Defaults to 4. - max_labeled_demos (int): maximum number of labeled demonstrations to use. Defaults to 4. - auto (Optional[Literal["light", "medium", "heavy"]]): automatic configuration mode. If set, will override num_candidates and max_steps. "light": n=6, val_size=100; "medium": n=12, val_size=300; "heavy": n=18, val_size=1000. Defaults to "medium". - max_steps (int): maximum number of optimization steps. Required if auto is None. - num_candidates (Optional[int]): number of candidates to generate for each optimization step. Required if auto is None. - num_threads (Optional[int]): number of threads to use for parallel evaluation. If None, will use single thread. - max_errors (int): maximum number of errors allowed during evaluation before stopping. Defaults to 10. - seed (int): random seed for reproducibility. Defaults to 9. - init_temperature (float): initial temperature for instruction generation. Defaults to 0.5. - track_stats (bool): whether to track optimization statistics. Defaults to True. - save_path (Optional[str]): path to save optimization results. If None, results will not be saved. - minibatch (bool): whether to use minibatch evaluation during optimization. Defaults to True. - minibatch_size (int): size of minibatch for evaluation. Defaults to 35. - minibatch_full_eval_steps (int): number of minibatch steps between full evaluations. Defaults to 5. - program_aware_proposer (bool): whether to use program-aware instruction proposer. Defaults to True. - data_aware_proposer (bool): whether to use data-aware instruction proposer. Defaults to True. - view_data_batch_size (int): batch size for viewing data during instruction proposal. Defaults to 10. - tip_aware_proposer (bool): whether to use tip-aware instruction proposer. Defaults to True. - fewshot_aware_proposer (bool): whether to use fewshot-aware instruction proposer. Defaults to True. - requires_permission_to_run (bool): whether to require user permission before running optimization. Defaults to False. - provide_traceback (Optional[bool]): whether to provide traceback for evaluation errors. If None, will use default setting. """ # check if the graph is compatible with the WorkFlowMipro optimizer. graph = self._validate_graph_compatibility(graph=graph) # convert the workflow graph to a callable program workflow_graph_program = WorkFlowGraphProgram( graph=graph, agent_manager=evaluator.agent_manager, executor_llm=evaluator.llm, collate_func=evaluator.collate_func, output_postprocess_func=evaluator.output_postprocess_func, ) # register optimizable parameters registry = self._register_optimizable_parameters(program=workflow_graph_program) super().__init__( registry=registry, program=workflow_graph_program, optimizer_llm=optimizer_llm or evaluator.llm, evaluator=evaluator, **kwargs ) def _validate_graph_compatibility(self, graph: WorkFlowGraph): """ Check if the graph is compatible with the WorkFlowMipro optimizer. Also, convert the MiproPromptTemplate data to MiproPromptTemplate instances. """ for node in graph.nodes: if len(node.agents) > 1: raise ValueError("WorkFlowMiproOptimizer only supports workflows where every node only has a single agent.") else: agent = node.agents[0] if not isinstance(agent, dict): raise ValueError(f"Unsupported agent type {type(agent)}. Expected 'dict'.") else: if "actions" in agent: # Agent has actions in its dict # All agents have a `ContextExtraction` action, filter it out non_ContextExtraction_actions = [ action for action in agent["actions"] if action["class_name"] != "ContextExtraction" ] if len(non_ContextExtraction_actions) > 1: raise ValueError(f"WorkFlowMiproOptimizer only supports workflows where every agent only has a single action. {agent['name']} has {len(non_ContextExtraction_actions)} actions.") # if "prompt_template" not in non_ContextExtraction_actions[0]: if non_ContextExtraction_actions[0].get("prompt_template", None) is None: # raise ValueError(f"Please provide a PromptTemplate for {agent['name']}.") logger.warning(f"{agent['name']} does not have a MiproPromptTemplate, its prompt will not be optimized.") else: prompt_template = non_ContextExtraction_actions[0]["prompt_template"] if isinstance(prompt_template, dict): prompt_template = PromptTemplate.from_dict(prompt_template) if isinstance(prompt_template, MiproPromptTemplate): # in some cases, the raw `prompt_template` can be a dict, convert it to a MiproPromptTemplate instance non_ContextExtraction_actions[0]["prompt_template"] = prompt_template else: logger.warning(f"{agent['name']} has a non-MiproPromptTemplate, its prompt will not be optimized. You should use `MiproPromptTemplate` to define the optimizable prompt.") else: # CustomizeAgent does not have actions in its dict # if "prompt_template" not in agent: # raise ValueError(f"Please provide a PromptTemplate for {agent['name']}.") if agent.get("prompt_template", None) is None: logger.warning(f"{agent['name']} does not have a MiproPromptTemplate, its prompt will not be optimized.") else: prompt_template = agent["prompt_template"] if isinstance(prompt_template, dict): prompt_template = PromptTemplate.from_dict(prompt_template) if isinstance(prompt_template, MiproPromptTemplate): # in some cases, the raw `prompt_template` can be a dict, convert it to a MiproPromptTemplate instance agent["prompt_template"] = prompt_template else: logger.warning(f"{agent['name']} has a non-MiproPromptTemplate, its prompt will not be optimized. You should use `MiproPromptTemplate` to define the optimizable prompt.") return graph def _validate_evaluator(self, evaluator: Callable = None, benchmark: Benchmark = None, metric_name: str = None) -> Callable: if evaluator and isinstance(evaluator, Evaluator): # if evaluator is an Evaluator, convert it to a MiproEvaluatorWrapper evaluator = MiproEvaluatorWrapper(evaluator=evaluator, benchmark=benchmark, metric_name=metric_name) return super()._validate_evaluator(evaluator, benchmark, metric_name) def _register_optimizable_parameters(self, program: WorkFlowGraphProgram): registry = MiproRegistry() workflow_graph = program.graph for i, node in enumerate(workflow_graph.nodes): agent = node.agents[0] # only one agent per node is allowed if "actions" in agent: # Agent Instance for j, action in enumerate(agent["actions"]): # only one action is allowed per agent. Use for loop because all the agent # will have the ContextExtraction action, which will be filtered out by default # since it does not have a prompt template. action_prompt_template = action.get("prompt_template", None) if action_prompt_template and isinstance(action_prompt_template, MiproPromptTemplate): registry.track( root_or_obj=program, path_or_attr=f"graph.nodes[{i}].agents[0]['actions'][{j}]['prompt_template']", name=f"{agent['name']}_prompt_template", input_names=node.get_input_names(), output_names=node.get_output_names() ) else: # CustomizeAgent Instance prompt_template = agent.get("prompt_template", None) if prompt_template and isinstance(prompt_template, MiproPromptTemplate): registry.track( root_or_obj=program, path_or_attr=f"graph.nodes[{i}].agents[0]['prompt_template']", name=f"{agent['name']}_prompt_template", input_names=node.get_input_names(), output_names=node.get_output_names() ) if not registry.fields: raise ValueError( "No optimizable parameters found in the workflow graph. " "Please check if the workflow graph is compatible with the WorkFlowMiproOptimizer. " "You should use `MiproPromptTemplate` to define the optimizable prompt." ) return registry