import yaml import regex import random import inspect import numpy as np from pydantic import Field from copy import deepcopy import xml.etree.ElementTree as ET from typing import Literal, Union, Optional, List from evoagentx.models import OpenAILLMConfig, OpenAILLM from evoagentx.evaluators import Evaluator from .optimizer import Optimizer from ..core.logging import logger from ..models.base_model import BaseLLM from ..benchmark.benchmark import Benchmark from ..workflow.action_graph import ActionGraph from ..core.callbacks import suppress_logger_info from ..workflow.workflow_graph import SequentialWorkFlowGraph,WorkFlowGraph from ..prompts.workflow.sew_optimizer import mutation_prompts, thinking_styles VALID_SCHEMES = ["python", "yaml", "code", "core", "bpmn"] import difflib def find_closest_name(inputname, name_refernece): name_reference_correct = [step["task_name"] for step in name_refernece] print("inputname", inputname) print("correct_list", name_reference_correct) correct_name = difflib.get_close_matches(inputname, name_reference_correct, n=1, cutoff=0.1) print(correct_name) correct_step = name_refernece[name_reference_correct.index(correct_name[0])] return correct_step class STRUCTUREWorkFlowScheme: """ The scheme of the workflow for SEW optimizer. """ def __init__(self, graph: WorkFlowGraph, **kwargs): self.graph = graph # the workflow graph to be transformed self.kwargs = kwargs def convert_to_scheme(self, scheme: str) -> str: """ Transform the WorkflowGraph to the desired scheme. """ if scheme not in VALID_SCHEMES: raise ValueError(f"Invalid scheme: {scheme}. The scheme should be one of {VALID_SCHEMES}.") if scheme == "python": repr = self.get_workflow_python_repr() elif scheme == "yaml": repr = self.get_workflow_yaml_repr() elif scheme == "code": repr = self.get_workflow_code_repr() elif scheme == "core": repr = self.get_workflow_core_repr() elif scheme == "bpmn": repr = self.get_workflow_bpmn_repr() return repr def parse_from_scheme(self, scheme: str, repr: str) -> WorkFlowGraph: """ Parse the SequentialWorkFlowGraph from the given scheme and representation. """ if scheme not in VALID_SCHEMES: raise ValueError(f"Invalid scheme: {scheme}. The scheme should be one of {VALID_SCHEMES}.") if scheme == "python": graph = self.parse_workflow_python_repr(repr) elif scheme == "yaml": graph = self.parse_workflow_yaml_repr(repr) elif scheme == "code": graph = self.parse_workflow_code_repr(repr) elif scheme == "core": graph = self.parse_workflow_core_repr(repr) elif scheme == "bpmn": graph = self.parse_workflow_bpmn_repr(repr) return graph def _get_workflow_repr_info(self) -> List[dict]: """ Get the information for the workflow representation. """ info = [] for node in self.graph.nodes: task_name = node.name input_names = [param.name for param in node.inputs] output_names = [param.name for param in node.outputs] task_info = { "task_name": task_name, "input_names": input_names, "output_names": output_names } info.append(task_info) return info def _convert_to_func_name(self, name: str) -> str: """ Convert the task name to the function name. """ name = name.lower().strip() name = name.replace(' ', '_').replace('-', '_') name = ''.join(c for c in name if c.isalnum() or c == '_') # Replace multiple consecutive underscores with a single underscore name = regex.sub(r'_+', "_", name) # Remove leading/trailing underscores name = name.strip('_') return name def _convert_to_title(self, name: str) -> str: func_name = self._convert_to_func_name(name) words = func_name.split('_') return ' '.join(word.capitalize() for word in words) def get_workflow_python_repr(self) -> str: repr_info = self._get_workflow_repr_info() if not repr_info: return "" python_workflow_info = [] for task_info in repr_info: name = self._convert_to_func_name(task_info['task_name']) input_names = [f'{input_name}' for input_name in task_info['input_names']] output_names = [f'{output_name}' for output_name in task_info['output_names']] python_workflow_info.append( "{{'name': '{name}', 'args': {args}, 'outputs': {outputs}}}".format( name=name, args=input_names, outputs=output_names ) ) python_workflow_repr = "steps = [\n" + ",\n".join(python_workflow_info) + "\n]" # print(python_workflow_repr) return python_workflow_repr def get_workflow_yaml_repr(self) -> str: repr_info = self._get_workflow_repr_info() if not repr_info: return "" yaml_workflow_info = [] for task_info in repr_info: name = self._convert_to_func_name(task_info['task_name']) input_names = "\n".join([f' - {input_name}' for input_name in task_info['input_names']]) output_names = "\n".join([f' - {output_name}' for output_name in task_info['output_names']]) yaml_workflow_info.append( "- name: {name}\n args:\n{input_names}\n outputs:\n{output_names}".format( name=name, input_names=input_names, output_names=output_names ) ) yaml_workflow_repr = "\n\n".join(yaml_workflow_info) return yaml_workflow_repr def get_workflow_code_repr(self) -> str: repr_info = self._get_workflow_repr_info() if not repr_info: return "" workflow_lines = [] for task_info in repr_info: # Convert task name to snake_case name = self._convert_to_func_name(task_info['task_name']) # Format inputs and outputs inputs = ", ".join(task_info['input_names']) outputs = ", ".join(task_info['output_names']) # Create the line in format: task_name(inputs) -> outputs line = f"{name}({inputs}) -> {outputs}" workflow_lines.append(line) # Join all lines with newlines workflow_repr = "\n".join(workflow_lines) return workflow_repr def get_workflow_bpmn_repr(self) -> str: repr_info = self._get_workflow_repr_info() if not repr_info: return "" # Start the BPMN XML bpmn_lines = [ '', '', ' ' ] # Add tasks for i, task_info in enumerate(repr_info): task_name = self._convert_to_func_name(task_info['task_name']) task_title = self._convert_to_title(task_info['task_name']) bpmn_lines.append(f' ') bpmn_lines.append(' ') bpmn_lines.append('') bpmn_lines.append(' ') # Add sequence flows # First flow from start to first task if repr_info: first_task_id = self._convert_to_func_name(repr_info[0]['task_name']) bpmn_lines.append(f' ') # Flows between tasks for i in range(len(repr_info) - 1): source_id = self._convert_to_func_name(repr_info[i]['task_name']) target_id = self._convert_to_func_name(repr_info[i + 1]['task_name']) flow_num = i + 2 bpmn_lines.append(f' ') # Last flow from last task to end if repr_info: last_task_id = self._convert_to_func_name(repr_info[-1]['task_name']) flow_num = len(repr_info) + 1 bpmn_lines.append(f' ') # Close tags bpmn_lines.append('') bpmn_lines.append('') return '\n'.join(bpmn_lines) def get_workflow_core_repr(self) -> str: repr_info = self._get_workflow_repr_info() if not repr_info: return "" workflow_lines = [] for i, task_info in enumerate(repr_info, 1): # Convert task name to title case task_name = self._convert_to_title(task_info['task_name']) # Create the line with the specified format next_step = i + 1 line = f"Step {i}::: Process ::: {task_name}:::next::Step {next_step}" workflow_lines.append(line) # Add the terminal step last_step = len(repr_info) + 1 workflow_lines.append(f"Step {last_step}::: Terminal ::: End of Workflow:::") return "\n".join(workflow_lines) def _find_task_index(self, step: dict, graph_repr_info: List[dict]) -> int: """ Find the index of the task in the original workflow graph. If the task is not found, return -1. Args: step (dict): The step of the workflow. graph_repr_info (List[dict]): The information of the original workflow graph. Returns: int: The index of the task. """ def _is_task_name_match(task_name: str, another_name: str) -> bool: return self._convert_to_func_name(task_name) == self._convert_to_func_name(another_name) def _is_task_inputs_match(task_inputs: List[str], another_inputs: List[str]) -> bool: return len(set(task_inputs) & set(another_inputs)) == len(task_inputs) def _is_task_outputs_match(task_outputs: List[str], another_outputs: List[str]) -> bool: return len(set(task_outputs) & set(another_outputs)) == len(task_outputs) for i, task in enumerate(graph_repr_info): # if _is_task_name_match(task["task_name"], step["name"]) and _is_task_inputs_match(task["input_names"], step["args"]) and _is_task_outputs_match(task["output_names"], step["outputs"]): # return i if _is_task_name_match(task["task_name"], step["name"]) and _is_task_outputs_match(task["output_names"], step["outputs"]): return i return -1 def create_workflow_graph_from_steps( self, steps: List[dict] ) -> WorkFlowGraph: """ Create a new workflow graph from the steps. Since both the inputs and outputs are provided, new tasks will be created in the new workflow graph. It is used for the `python` `yaml` and `code` representations. Args: steps (List[dict]): The steps of the workflow. The steps are in the format of: [ { "name": str, "args": List[str], "outputs": List[str] } ] Returns: SequentialWorkFlowGraph: The new workflow graph. """ original_workflow_config = self.graph.get_graph_info() repr_info = self._get_workflow_repr_info() new_tasks = [] get_known_list = [] for step in repr_info: get_known_list.append(step) for step in steps: task_index = self._find_task_index(step=step, graph_repr_info=repr_info) if task_index == -1: # create a new task task_name = step["name"] most_known_step = find_closest_name(task_name, get_known_list) most_known_step['name'] = most_known_step['task_name'] most_known_step['args'] = most_known_step['input_names'] most_known_step['outputs'] = most_known_step['output_names'] task_index_new = self._find_task_index(step=most_known_step, graph_repr_info=repr_info) print(step) print(task_index) item_new = deepcopy(original_workflow_config["tasks"][task_index_new]) item_new["name"] = task_name +str(np.random.randint(0,10000)) item_new['task_name'] = task_name +str(np.random.randint(0,10000)) new_tasks.append(item_new) # description = f"Task to {task_name.lower()}. " # if step["args"]: # description += f"Takes {', '.join(step['args'])} as input. " # if step["outputs"]: # description += f"Produces {', '.join(step['outputs'])} as output." # new_task = { # "name": task_name, # "description": description, # "inputs": [ # { # "name": input_name, # "type": "str", # "description": f"Input parameter {input_name} for {task_name}" # } for input_name in step["args"] # ], # "outputs": [ # { # "name": output_name, # "type": "str", # "description": f"Output parameter {output_name} from {task_name}" # } for output_name in step["outputs"] # ], # "prompt": "to be updated", # # "llm_config": original_workflow_config["tasks"][0]["llm_config"], # "parse_mode": "str" # } # new_tasks.append(new_task) else: # copy the task from the original workflow graph if original_workflow_config["tasks"][task_index] not in new_tasks: new_tasks.append(deepcopy(original_workflow_config["tasks"][task_index])) # create new workflow configuration new_workflow_config = { "goal": original_workflow_config["goal"], "tasks": new_tasks } # create new workflow graph new_graph = SequentialWorkFlowGraph.from_dict(new_workflow_config) return new_graph def create_workflow_graph_from_task_names( self, task_names: Optional[List[str]] = None, task_titles: Optional[List[str]] = None ) -> SequentialWorkFlowGraph: """ Create a new workflow graph from the task names or titles. Since only the task names or titles are provided, the tasks in the new workflow graph will be copied from the original workflow graph. It is used for the `bpmn` and `core` representations. Args: task_names (Optional[List[str]]): The names of the tasks. task_titles (Optional[List[str]]): The titles of the tasks. Returns: SequentialWorkFlowGraph: The new workflow graph. """ if task_names: original_workflow_config = self.graph.get_graph_info() tasks = task_names original_tasks = {self._convert_to_func_name(task["name"]): task for task in original_workflow_config["tasks"]} elif task_titles: original_workflow_config = self.graph.get_graph_info() tasks = task_titles original_tasks = {self._convert_to_title(task["name"]): task for task in original_workflow_config["tasks"]} else: raise ValueError("No task names or titles provided.") new_tasks = [] for task in tasks: if task not in original_tasks: raise ValueError(f"Task {task} not found in the original workflow.") new_tasks.append(deepcopy(original_tasks[task])) # create new workflow configuration new_workflow_config = { "goal": original_workflow_config["goal"], "tasks": new_tasks } # create new workflow graph new_graph = WorkFlowGraph.from_dict(new_workflow_config) return new_graph def parse_workflow_python_repr(self, repr: str) -> WorkFlowGraph: """ Parse the workflow from the python representation. The input format is: steps = [ {"name": task_name, "args": [input1, input2, ...],"outputs": [output1, output2, ...]}, {"name": another_task_name, "args": [input1, input2, ...],"outputs": [output1, output2, ...]}, ... ] """ # try: # # extract ```python ``` # code_block = regex.search(r'```python\s*(.*?)\s*```', repr, regex.DOTALL) # if not code_block: # raise ValueError("No Python code block found in the representation") # code_block = code_block.group(1).strip() # # relevant_lines = [] # # for line in code_block.splitlines(): # # line = line.strip() # # if not line or line.startswith("#") or line.startswith("```"): # # continue # # if all(key in line for key in ["name", "args", "outputs"]): # # relevant_lines.append(line) # # steps_str = "[\n" + "\n".join(relevant_lines) + "\n]" # # steps = eval(steps_str) # steps = eval(code_block.replace("steps = ", "").strip()) # new_graph = self.create_workflow_graph_from_steps(steps=steps) # return new_graph # except Exception as e: # logger.warning(f"Failed to parse workflow string: {e}. Return the original workflow.") # extract ```python ``` code_block = regex.search(r'```python\s*(.*?)\s*```', repr, regex.DOTALL) if not code_block: raise ValueError("No Python code block found in the representation") code_block = code_block.group(1).strip() # relevant_lines = [] # for line in code_block.splitlines(): # line = line.strip() # if not line or line.startswith("#") or line.startswith("```"): # continue # if all(key in line for key in ["name", "args", "outputs"]): # relevant_lines.append(line) # steps_str = "[\n" + "\n".join(relevant_lines) + "\n]" # steps = eval(steps_str) steps = eval(code_block.replace("steps = ", "").strip()) # print(steps) new_graph = self.create_workflow_graph_from_steps(steps=steps) return new_graph def parse_workflow_yaml_repr(self, repr: str) -> WorkFlowGraph: """ Parse the workflow from the yaml representation. The input format is: - name: task_name args: - input1 - input2 outputs: - output1 """ try: # extract ```yaml ``` match = regex.search(r'```yaml\s*(.*?)\s*```', repr, regex.DOTALL) if not match: raise ValueError("No YAML code block found in the representation") yaml_block = match.group(1).strip() steps = yaml.safe_load(yaml_block) # relevant_lines = [] # in_step = False # for line in yaml_block.splitlines(): # stripped_line = line.strip() # if stripped_line.startswith('- name:'): # in_step = True # relevant_lines.append(line) # elif in_step and ( # stripped_line.startswith('args:') or # stripped_line.startswith('outputs:') or # stripped_line.startswith('- ') # ): # relevant_lines.append(line) # elif not stripped_line: # in_step = False # yaml_step = "\n".join(relevant_lines) # steps = yaml.safe_load(yaml_step) new_graph = self.create_workflow_graph_from_steps(steps=steps) return new_graph except Exception as e: logger.warning(f"Failed to parse workflow string: {e}. Return the original workflow.") return self.graph def parse_workflow_code_repr(self, repr: str) -> WorkFlowGraph: """ Parse the workflow from the code representation. The input format is: task_name(input1, input2, ...) -> output1, output2, ... another_task_name(input1, input2, ...) -> output1, output2, ... ... """ try: # extract ```code ``` match = regex.search(r'```code\s*(.*?)\s*```', repr, regex.DOTALL) if not match: raise ValueError("No code block found in the representation") code_block = match.group(1).strip() lines = [line.strip() for line in code_block.split("\n") if line.strip() and "->" in line] steps = [] for line in lines: # Remove any leading numbers and dots (e.g., "1. ") line = regex.sub(r'^\d+\.\s*', '', line) func_part, output_part = line.split('->') func_part = func_part.strip() name = func_part[:func_part.index('(')] args_str = func_part[func_part.index('(') + 1:func_part.rindex(')')] args = [arg.strip() for arg in args_str.split(',') if arg.strip()] outputs = [out.strip() for out in output_part.split(',') if out.strip()] step = {"name": name, "args": args, "outputs": outputs} steps.append(step) if not steps: raise ValueError("No steps found in the workflow.") new_graph = self.create_workflow_graph_from_steps(steps=steps) return new_graph except Exception as e: logger.warning(f"Failed to parse workflow string: {e}. Return the original workflow.") return self.graph def parse_workflow_bpmn_repr(self, repr: str) -> WorkFlowGraph: """ Parse the workflow from the BPMN XML representation. The input format is BPMN XML with: - task elements defining the tasks - sequenceFlow elements defining the order of tasks Will extract ordered task names from the sequence flows and create a workflow. """ try: # extract ```bpmn ``` match = regex.search(r'```bpmn\s*(.*?)\s*```', repr, regex.DOTALL) if not match: raise ValueError("No BPMN code block found in the representation") bpmn_block = match.group(1).strip() # Parse XML string root = ET.fromstring(bpmn_block) # Define namespace for BPMN XML ns = {'bpmn': 'http://www.omg.org/spec/BPMN/20100524/MODEL'} # Get process element process = root.find('bpmn:process', ns) or root.find('process') if process is None: raise ValueError("No process element found in BPMN XML") # Create a dictionary of all tasks tasks = {} # for task in process.findall('.//task', ns) or process.findall('.//task'): for task in process.findall("bpmn:task", ns): tasks[task.get('id')] = task.get('name') # Get sequence flows and order them flows = {} ordered_tasks = [] current_ref = 'start' # Create dictionary of source -> target # for flow in process.findall('.//sequenceFlow', ns) or process.findall('.//sequenceFlow'): for flow in process.findall("bpmn:sequenceFlow", ns): flows[flow.get('sourceRef')] = flow.get('targetRef') # Follow the sequence flows to get ordered tasks while current_ref in flows: next_ref = flows[current_ref] if next_ref in tasks: # Only add if it's a task (not end event) ordered_tasks.append(tasks[next_ref]) current_ref = next_ref # Create new workflow graph using the ordered task names new_graph = self.create_workflow_graph_from_task_names(task_titles=ordered_tasks) return new_graph except Exception as e: logger.warning(f"Failed to parse BPMN workflow string: {e}. Return the original workflow.") return self.graph def parse_workflow_core_repr(self, repr: str) -> WorkFlowGraph: """ Parse the workflow from the Core representation. The input format is: Step 1::: Process ::: Task Name:::next::Step 2 Step 2::: Process ::: Another Task:::next::Step 3 ... Step N::: Terminal ::: End of Workflow::: Will extract task names from Process steps and create a workflow. """ try: # extract ```core ``` match = regex.search(r'```core\s*(.*?)\s*```', repr, regex.DOTALL) if not match: raise ValueError("No core code block found in the representation") core_block = match.group(1).strip() # Split into lines and remove empty lines lines = [line.strip() for line in core_block.split('\n') if line.strip()] # Initialize flows and tasks dictionaries flows = {} # step -> next_step tasks = {} # step -> task_title # First pass: build flows and tasks mappings for line in lines: parts = line.split(':::') current_step = parts[0].strip() step_type = parts[1].strip() if step_type == 'Process': # Extract task title and next step task_title = parts[2].strip() tasks[current_step] = task_title if len(parts) > 3 and "next" in parts[3]: next_step = parts[3].split("::")[-1].strip() flows[current_step] = next_step elif step_type == 'Terminal': flows[current_step] = None # Second pass: follow flows to build ordered task list ordered_tasks = [] current_step = 'Step 1' while current_step in flows: if current_step in tasks: # Only add if it's a Process step ordered_tasks.append(tasks[current_step]) current_step = flows[current_step] # Create new workflow graph using the ordered task titles new_graph = self.create_workflow_graph_from_task_names(task_titles=ordered_tasks) return new_graph except Exception as e: logger.warning(f"Failed to parse Core workflow string: {e}. Return the original workflow.") return self.graph class SimplePromptBreeder: def __init__(self, llm: BaseLLM, evaluator: None, **kwargs): self.llm = llm self.evaluator = evaluator self.history_log = [] self.kwargs = kwargs def generate_mutation_prompt(self, task_description: str, **kwargs) -> str: """ Generate the mutation prompt for optimization. """ thinking_style = random.choice(thinking_styles) # hyper_mutation_prompt = thinking_style + "\n\nProblem Description: " + task_description + ".\n" + "Output: " hyper_mutation_prompt = "Please generate a improved prompts based on the following information. " + "\n\nProblem Description: " + task_description + ".\n" + "Output: " # print(">>>>>>>>>> Hyper mutation prompt: <<<<<<<<<<<\n", hyper_mutation_prompt) try: mutation_prompt = self.llm.generate( prompt=hyper_mutation_prompt, system_message="You are a helpful assistant. Do not generate harmful content. ", ).content except: mutation_prompt = self.llm.generate( prompt=hyper_mutation_prompt, system_message="You are a helpful assistant. Do not generate harmful content. ", ).content return mutation_prompt def get_mutation_prompt(self, task_description: str, order: Literal["zero-order", "first-order"], **kwargs) -> str: """ Get the mutation prompt for optimization. """ if order == "zero-order": mutation_prompt = self.generate_mutation_prompt(task_description=task_description) elif order == "first-order": mutation_prompt = random.choice(mutation_prompts) else: raise ValueError(f"Invalid order: {order}. The order should be either 'zero-order' or 'first-order'.") return mutation_prompt def generate_prompt(self, task_description: str, prompt: str, order: Literal["zero-order", "first-order"], **kwargs) -> str: """ Generate the prompt for optimization. Args: task_description (str): The description of the task, normally the goal of the workflow. prompt (str): The prompt to optimize. order (Literal["zero-order", "first-order"]): The order of the mutation prompt. Returns: str: The optimized prompt. """ mutation_prompt = self.get_mutation_prompt(task_description=task_description, order=order) prompt = mutation_prompt + "\n\nINSTRUCTION:\n\n" + prompt # print(">>>>>>>>>> Prompt: <<<<<<<<<<<\n", prompt) new_prompt = self.llm.generate( prompt=prompt, system_message="You are a helpful assistant", ).content return new_prompt def critic_and_update_prompt(self, task_description: str, prompt: str, order: Literal["zero-order", "first-order"], scorer=None, calltime=1, **kwargs) -> str: """ Generate the prompt for optimization. Args: task_description (str): The description of the task, normally the goal of the workflow. prompt (str): The prompt to optimize. order (Literal["zero-order", "first-order"]): The order of the mutation prompt. Returns: str: The optimized prompt. """ # print(self.evaluator._evaluation_records) problem_list = '''''' for item in self.evaluator._evaluation_records.keys(): problem_s = "Questions: " + self.evaluator._evaluation_records[item]['trajectory'][0].content['question']+'\n' prediction_s = "Predictions: " + self.evaluator._evaluation_records[item]['prediction']+'\n' solution_s = "Solutions: " + self.evaluator._evaluation_records[item]['label']['canonical_solution']+'\n' # if self.evaluator.dataname != "humanevalplus": # if 'test' in list(self.evaluator._evaluation_records[item]['label'].keys()): # test_s = "Unit tests: " + self.evaluator._evaluation_records[item]['label']['test'][0:1000] # elif 'tests' in list(self.evaluator._evaluation_records[item]['label'].keys()): # test_s = "Unit tests: " + self.evaluator._evaluation_records[item]['label']['tests'][0:1000] # else: # test_s = "Example solution: " + self.evaluator._evaluation_records[item]['label']["canonical_solution"] if 'test' in list(self.evaluator._evaluation_records[item]['label'].keys()): test_s = "Unit tests: " + self.evaluator._evaluation_records[item]['label']['test'][0:10000]+'\n' elif 'tests' in list(self.evaluator._evaluation_records[item]['label'].keys()): test_s = "Unit tests: " + self.evaluator._evaluation_records[item]['label']['tests'][0:10000]+'\n' metric_s = "Score: " + str(self.evaluator._evaluation_records[item]['metrics']['pass@1']) + "\n" if self.evaluator._evaluation_records[item]['metrics']['pass@1'] ==0: if "An error occurred:" == self.evaluator.error_list[item]: metric_s += "Error reason: Computation result is incorrect." else: erroreason = self.evaluator.error_list[item].replace("An error occurred:", "") metric_s += f"Error reason: {erroreason}" else: metric_s += "The solution is correct." joint_s = problem_s + prediction_s + solution_s +test_s + metric_s problem_list += joint_s # print(problem_list) if calltime==1: critic_prompt = "You will evaluate a coding solution as a workflow for code generation. You should analyze failures using execution history as well as workflow trajectory. Do not attempt to solve it yourself, do not give a solution, only identify errors and problems in the structure of this AI agent workflow and incorrect information in prompts. Only return the problems in one paragraph. Be super concise. The workflow is:" question_prompt = "The questions, solutions, unit tests, and evaluated metrics based on this workflow is: " + problem_list # question_prompt = "" critic_out = self.llm.generate( prompt=prompt+question_prompt, system_message=critic_prompt, ).content print(critic_out) else: critic_prompt_outlist = '''Please summarize the following problems in one paragraph. Be super concise.\n''' for item in range(calltime): critic_prompt = "You will evaluate a coding solution as a workflow for code generation. You should analyze failures using execution history of tested problems as well as workflow trajectory. Do not attempt to solve it yourself, do not give a solution, only identify errors and problems in the structure of this agent workflow. Be super concise. The workflow is:" question_prompt = "The questions, solutions, and evaluated metrics based on this workflow is: " + problem_list critic_out = self.llm.generate( prompt=prompt+question_prompt, system_message=critic_prompt, ).content critic_prompt_outlist = critic_prompt_outlist + f"Detected Issue {item+1}:" + critic_out +"\n" critic_out = self.llm.generate( prompt=critic_prompt_outlist, system_message="You are an expert in summarizing information and data.", ).content print(critic_out) # mutation_prompt = self.get_mutation_prompt(task_description=task_description, order=order) if scorer == None: prompt = "The detected issue is:" + critic_out+"You should always improve workflow by correcting the issue without changing the inputs and outputs of nodes in the workflow. You can remove redundant agents. You should keep the graph executable.\n" + "\n\nThe original workflow is:\n\n" + prompt else: prompt = "The detected issue is:" + critic_out + f"You also need to ensure the new workflow can increase the model performance score: {scorer}." + "\n\nThe original workflow is:\n\n" + prompt + "You should always improve workflow by correcting the issue without changing the inputs and outputs of nodes in the workflow. You can remove redundant agents. You should keep the graph executable.\nYour OUTPUT:" # print(">>>>>>>>>> Prompt: <<<<<<<<<<<\n", prompt) new_prompt = self.llm.generate( prompt=prompt, system_message="You are a Graph Optimization Agent. Your goal is to iteratively improve graph performance through systematic optimization. You need to modify the workflow and improve the structure by solving the issue. Only change the order of agents or recall the KNOWN agents in the workflow. DO NOT change the names and inputs of agents. You should keep the graph executable.", ).content print(new_prompt) return new_prompt def update_dev_set(dataset): import numpy as np permutation = np.random.permutation(len(dataset._dev_data_full)) # radnomly select 50 samples for dev and 100 samples for test (be consistent with other models) dev_data_task_ids = [dataset._dev_data_full[idx]["task_id"] for idx in permutation[:len(dataset._dev_data)]] full_data = dataset._dev_data_full dev_data = [example for example in full_data if example["task_id"] in dev_data_task_ids] return dev_data class STRUCTUREOptimizer(Optimizer): graph: Union[WorkFlowGraph, ActionGraph] = Field(description="The workflow to optimize.") repr_scheme: str = Field(default="python", description="The scheme to represent the workflow.") optimize_mode: Literal["all", "structure", "prompt"] = Field(default="all", description="The mode to optimize the workflow.") order: Literal["zero-order", "first-order"] = Field(default="zero-order", description="Whether to use zero-order (using hyper-mutation prompt) or first-order (using mutation prompt) optimization.") calltime: int = Field(default=1, description="Number of textgrad used for evaluation.") num_workers: int = Field(default=1, description="Number of textgrad used for evaluation.") def init_module(self, **kwargs): self._snapshot: List[dict] = [] self._prompt_breeder = SimplePromptBreeder(llm=self.llm, evaluator = self.evaluator) # generate prompt for optimization self._convergence_check_counter = 0 self._best_score = float("-inf") self._prompt_dict = {} if isinstance(self.graph, ActionGraph): if self.optimize_mode != "prompt": raise ValueError( f"{type(self).__name__} only support prompt optimization when `graph` is an `ActionGraph`. " f"The `optimize_mode` should be set to `prompt`, but got {self.optimize_mode}." ) def optimize(self, dataset: Benchmark, **kwargs): if isinstance(self.graph, WorkFlowGraph): logger.info(f"Optimizing the {type(self.graph).__name__} workflow with {self.repr_scheme} representation.") elif isinstance(self.graph, ActionGraph): logger.info(f"Optimizing the {type(self.graph).__name__} graph ...") graph: Union[WorkFlowGraph, ActionGraph] = self.graph logger.info("Run initial evaluation on the original workflow ...") with suppress_logger_info(): metrics = self.evaluate(dataset, eval_mode="dev", graph=graph) self._prompt_breeder = SimplePromptBreeder(llm=self.llm, evaluator = self.evaluator) # generate prompt for optimization logger.info(f"Initial metrics: {metrics}") self.log_snapshot(graph=graph, metrics=metrics) set_scorer = None if kwargs["provided_scorer"] == True: set_scorer = metrics for i in range(self.max_steps): # try: # # perform a step of optimization # graph = self.step(set_scorer=set_scorer) # # print(graph) # # evaluate the workflow # if (i + 1) % self.eval_every_n_steps == 0: # logger.info(f"Evaluate the workflow at step {i+1} ...") # with suppress_logger_info(): # metrics = self.evaluate(dataset, eval_mode="dev") # logger.info(f"Step {i+1} metrics: {metrics}") # self.log_snapshot(graph=graph, metrics=metrics) # except Exception as e: # logger.warning(f"Error in step {i}: {e}. Skip this step.") # continue # if self.convergence_check(): # logger.info(f"Convergence check passed at step {i+1}. Stop the optimization.") # break # perform a step of optimization graph = self.step(set_scorer=set_scorer, step=i) # print(graph) # evaluate the workflow if (i + 1) % self.eval_every_n_steps == 0: logger.info(f"Evaluate the workflow at step {i+1} ...") with suppress_logger_info(): metrics = self.evaluate(dataset, eval_mode="dev") logger.info(f"Step {i+1} metrics: {metrics}") self.log_snapshot(graph=graph, metrics=metrics) print("randomly update dataset") self.dataset._dev_data = update_dev_set(self.dataset) if i == self.max_steps - 1: logger.info(f"Reach the maximum number of steps {self.max_steps}. Stop the optimization.") # set self.graph to the best graph logger.info("Restore the best graph from the snapshot ...") self.restore_best_graph() def step(self, **kwargs) -> Union[WorkFlowGraph, ActionGraph]: """ Take a step of optimization and return the optimized graph. """ graph = self._select_graph_with_highest_score(return_metrics=False) if isinstance(graph, WorkFlowGraph): new_graph = self._workflow_graph_step(graph, kwargs["set_scorer"], kwargs["step"]) elif isinstance(graph, ActionGraph): new_graph = self._action_graph_step(graph, kwargs["set_scorer"], kwargs["step"]) else: raise ValueError(f"Invalid graph type: {type(graph)}. The graph should be an instance of `WorkFlowGraph` or `ActionGraph`.") return new_graph def evaluate( self, dataset: Benchmark, eval_mode: str = "test", graph: Optional[Union[WorkFlowGraph, ActionGraph]] = None, indices: Optional[List[int]] = None, sample_k: Optional[int] = None, **kwargs ) -> dict: """ Evaluate the workflow. If `graph` is provided, use the provided graph for evaluation. Otherwise, use the graph in the optimizer. Args: dataset (Benchmark): The dataset to evaluate the workflow on. eval_mode (str): The evaluation mode. Choices: ["test", "dev", "train"]. graph (Union[WorkFlowGraph, ActionGraph], optional): The graph to evaluate. If not provided, use the graph in the optimizer. indices (List[int], optional): The indices of the data to evaluate the workflow on. sample_k (int, optional): The number of data to evaluate the workflow on. If provided, a random sample of size `sample_k` will be used. Returns: dict: The metrics of the workflow evaluation. """ self.dataset = dataset graph = graph if graph is not None else self.graph agent_manager = self.evaluator.agent_manager agent_manager.add_agents_from_workflow(graph, llm_config=self.llm.config) # print(agent_manager) # obtain Evaluator self.evaluator = Evaluator(llm=self.llm, agent_manager=agent_manager, collate_func=self.collate_func, num_workers=self.num_workers, verbose=True) self.evaluator.dataname = self.dataset.dataname metrics_list = [] for i in range(self.eval_rounds): eval_info = [ f"[{type(graph).__name__}]", f"Evaluation round {i+1}/{self.eval_rounds}", f"Mode: {eval_mode}" ] if indices is not None: eval_info.append(f"Indices: {len(indices)} samples") if sample_k is not None: eval_info.append(f"Sample size: {sample_k}") logger.info(" | ".join(eval_info)) # if self.dataset.dataname == 'scicode': # metrics = await self.evaluator.async_evaluate( # graph=graph, # benchmark=dataset, # eval_mode=eval_mode, # indices=indices, # sample_k=sample_k, # **kwargs # ) # else: # metrics = self.evaluator.evaluate( # graph=graph, # benchmark=dataset, # eval_mode=eval_mode, # indices=indices, # sample_k=sample_k, # **kwargs # ) metrics = self.evaluator.evaluate( graph=graph, benchmark=dataset, eval_mode=eval_mode, indices=indices, sample_k=sample_k, **kwargs ) metrics_list.append(metrics) avg_metrics = self.evaluator._calculate_average_score(metrics_list) self.dataset = dataset self.evaluator.error_list = deepcopy(self.dataset.error_list) self.dataset.error_list = {} return avg_metrics def log_snapshot(self, graph: Union[WorkFlowGraph, ActionGraph], metrics: dict): if isinstance(graph, WorkFlowGraph): graph_info = graph.get_graph_info() elif isinstance(graph, ActionGraph): # TODO check if the action graph is valid graph_info = graph else: raise ValueError(f"Invalid graph type: {type(graph)}. The graph should be an instance of `SequentialWorkFlowGraph` or `ActionGraph`.") self._snapshot.append( { "index": len(self._snapshot), "graph": deepcopy(graph_info), "metrics": metrics, } ) def _select_graph_with_highest_score(self, return_metrics: bool = False) -> Union[SequentialWorkFlowGraph, ActionGraph]: if len(self._snapshot) == 0: return self.graph snapshot_scores = [np.mean(list(snapshot["metrics"].values())) for snapshot in self._snapshot] best_index = np.argmax(snapshot_scores) if isinstance(self.graph, WorkFlowGraph): graph = WorkFlowGraph.from_dict(self._snapshot[best_index]["graph"]) elif isinstance(self.graph, ActionGraph): # TODO check if the action graph is valid graph = self._snapshot[best_index]["graph"] else: raise ValueError(f"Invalid graph type: {type(self.graph)}. The graph should be an instance of `SequentialWorkFlowGraph` or `ActionGraph`.") if return_metrics: return graph, self._snapshot[best_index]["metrics"] return graph def restore_best_graph(self): best_graph, best_metrics = self._select_graph_with_highest_score(return_metrics=True) logger.info(f"Restore the best graph from snapshot with metrics {best_metrics} ...") self.graph = best_graph def _wfg_structure_optimization_step(self, graph: WorkFlowGraph, scorer, step) -> WorkFlowGraph: """ optinize the structure of the workflow graph and return the optimized graph. Args: graph (SequentialWorkFlowGraph): The workflow graph to optimize. Returns: SequentialWorkFlowGraph: The optimized workflow graph. """ graph_scheme = STRUCTUREWorkFlowScheme(graph=graph) graph_repr = graph_scheme.convert_to_scheme(scheme=self.repr_scheme) if self.repr_scheme == "python": output_format = "\n\nALWAYS wrap the refined workflow in ```python\n``` format and DON'T include any other text within the code block!" elif self.repr_scheme == "yaml": output_format = "\n\nALWAYS wrap the refined workflow in ```yaml\n``` format and DON'T include any other text within the code block!" elif self.repr_scheme == "code": output_format = "\n\nALWAYS wrap the refined workflow in ```code\n``` format and DON'T include any other text within the code block!" elif self.repr_scheme == "core": output_format = "\n\nALWAYS wrap the refined workflow in ```core\n``` format and DON'T include any other text within the code block!" elif self.repr_scheme == "bpmn": output_format = "\n\nALWAYS wrap the refined workflow in ```bpmn\n``` format and DON'T include any other text within the code block!" else: raise ValueError(f"Invalid representation scheme: {self.repr_scheme}. The scheme should be one of {VALID_SCHEMES}.") prompt = "Task Description: " + graph.goal + "\n\nWorkflow Steps: " + graph_repr + output_format # if step%5==0: # # print(prompt) # new_graph_repr = self._prompt_breeder.critic_and_update_prompt(task_description=graph.goal, prompt=prompt, order=self.order, scorer=scorer, calltime=self.calltime) # # print(new_graph_repr) # new_graph = graph_scheme.parse_from_scheme(scheme=self.repr_scheme, repr=new_graph_repr) # print(new_graph) # else: # new_graph = graph new_graph_repr = self._prompt_breeder.critic_and_update_prompt(task_description=graph.goal, prompt=prompt, order=self.order, scorer=scorer, calltime=self.calltime) new_graph = graph_scheme.parse_from_scheme(scheme=self.repr_scheme, repr=new_graph_repr) return new_graph def _wfg_prompt_optimization_step(self, graph: WorkFlowGraph, scorer=None) -> WorkFlowGraph: task_description = graph.goal graph_scheme = STRUCTUREWorkFlowScheme(graph=graph) graph_repr = graph_scheme.convert_to_scheme(scheme=self.repr_scheme) graph_info = graph.get_graph_info() problem_list = '''''' for item in self.evaluator._evaluation_records.keys(): problem_s = "Questions: " + self.evaluator._evaluation_records[item]['trajectory'][0].content['question']+'\n' prediction_s = "Predictions: " + self.evaluator._evaluation_records[item]['prediction']+'\n' solution_s = "Solutions: " + self.evaluator._evaluation_records[item]['label']['canonical_solution']+'\n' # if self.evaluator.dataname != "humanevalplus": # if 'test' in list(self.evaluator._evaluation_records[item]['label'].keys()): # test_s = "Unit tests: " + self.evaluator._evaluation_records[item]['label']['test'][0:1000] # elif 'tests' in list(self.evaluator._evaluation_records[item]['label'].keys()): # test_s = "Unit tests: " + self.evaluator._evaluation_records[item]['label']['tests'][0:1000] # else: # test_s = "Example solution: " + self.evaluator._evaluation_records[item]['label']["canonical_solution"] if 'test' in list(self.evaluator._evaluation_records[item]['label'].keys()): test_s = "Unit tests: " + self.evaluator._evaluation_records[item]['label']['test'][0:10000]+'\n' elif 'tests' in list(self.evaluator._evaluation_records[item]['label'].keys()): test_s = "Unit tests: " + self.evaluator._evaluation_records[item]['label']['tests'][0:10000]+'\n' metric_s = "Score: " + str(self.evaluator._evaluation_records[item]['metrics']['pass@1']) + "\n" if self.evaluator._evaluation_records[item]['metrics']['pass@1'] ==0: if "An error occurred: " == self.evaluator.error_list[item]: metric_s += "Error reason: Computation result is incorrect." else: erroreason = self.evaluator.error_list[item].replace("An error occurred: ", "") metric_s += f"Error reason: {erroreason}" else: metric_s += "The solution is correct." joint_s = problem_s + prediction_s + solution_s +test_s + metric_s problem_list += joint_s print(problem_list) for i, task in enumerate(graph_info["tasks"]): if task['name'] not in list(self._prompt_dict.keys()): self._prompt_dict[task['name']] = [] original_prompt = task["prompt"] optimization_prompt = "Task Description: " + task_description + "\n\nWorkflow Steps:\n" + graph_repr + f"\n\nINSTRUCTION for the {i+1}-th task:\n\"\"\"\n" + original_prompt + "\n\"\"\"" error_prompt = optimization_prompt + f"The name of this agent is {task['name']}" + "The questions, solutions, and evaluated metrics based on this workflow is: " + problem_list + "You should detect the issues in the original prompt by considering these questions, predictions, tests, solutions, and score!" critic_issues = self.llm.generate(error_prompt).content optimization_prompt += f"The new prompts should consider fixing the issues by adjusting the prompt content: {critic_issues}. You should not change the original role and task of the assigned agent." if self._prompt_dict[task['name']] != []: prev_prompt = "\n".join(self._prompt_dict[task['name']]) optimization_prompt += f"The previous prompts are: {prev_prompt}\nYou should also fix the problems in these prompts." optimization_prompt += f"\n\nGiven the above information, please refine the instruction for the {i+1}-th task.\n" optimization_prompt += r"Note that you must always use bracket (e.g. `{input_name}`, `{code}`, `{question}`) to wrap the inputs of the tasks in your refined instruction. You must ensure the prompts contain all inputs. You cannot change the name of functions.\n" ###new one optimization_prompt += "Your prompt should not change the function name and entry_point in the question. Only output the refined instruction and DON'T include any other text!" new_prompt = self._prompt_breeder.generate_prompt(task_description=task_description, prompt=optimization_prompt, order=self.order) graph_info["tasks"][i]["prompt"] = new_prompt # print("task name", task['name']) # print("detected issue", critic_issues) # print("renewed prompt", new_prompt) self._prompt_dict[task['name']].append(new_prompt) new_graph = SequentialWorkFlowGraph.from_dict(graph_info) return new_graph def _workflow_graph_step(self, graph: WorkFlowGraph, scorer, step) -> WorkFlowGraph: if self.optimize_mode == "structure" or self.optimize_mode == "all": # optimize the structure of the graph graph = self._wfg_structure_optimization_step(graph, scorer=scorer, step=step) if self.optimize_mode == "prompt" or self.optimize_mode == "all": # optimize the prompt of the graph graph = self._wfg_prompt_optimization_step(graph, scorer=scorer) return graph def _action_graph_prompt_optimization_step(self, graph: ActionGraph) -> ActionGraph: task_description = graph.description graph_info = graph.get_graph_info() graph_steps = inspect.getsource(getattr(graph, "execute")) for operator_name, operator_info in graph_info["operators"].items(): original_prompt = operator_info["prompt"] optimization_prompt = "Task Description: " + task_description + "\n\nWorkflow Steps:\n" + graph_steps + f"\n\nINSTRUCTION for the `{operator_name}` operator:\n\"\"\"\n" + original_prompt + "\n\"\"\"" optimization_prompt += "\n\nThe interface of the operator is as follows:\n" + operator_info["interface"] optimization_prompt += f"\n\nGiven the above information, please refine the instruction for the `{operator_name}` operator.\n" optimization_prompt += r"Note that you should always use bracket (e.g. `{input_name}`) to wrap the inputs of the operator in your refined instruction, " optimization_prompt += "and the input names should be EXACTLY the same as those defined in the interface. DON'T use bracket to wrap output names." optimization_prompt += "\nOnly output the refined instruction and DON'T include any other text!" new_prompt = self._prompt_breeder.generate_prompt(task_description=task_description, prompt=optimization_prompt, order=self.order) new_prompt = new_prompt.replace("\"", "").strip() graph_info["operators"][operator_name]["prompt"] = new_prompt new_graph = ActionGraph.from_dict(graph_info) return new_graph def _action_graph_step(self, graph: ActionGraph) -> ActionGraph: if self.optimize_mode == "prompt": graph = self._action_graph_prompt_optimization_step(graph) else: raise ValueError(f"{type(self).__name__} only support prompt optimization when `self.graph` is an `ActionGraph` instance. " f"The `optimize_mode` should be set to `prompt`, but got {self.optimize_mode}.") return graph def convergence_check(self, **kwargs) -> bool: if not self._snapshot: logger.warning("No snapshots available for convergence check") return False # Get scores from snapshots scores = [np.mean(list(snapshot["metrics"].values())) for snapshot in self._snapshot] current_score = scores[-1] if current_score > self._best_score: self._best_score = current_score self._convergence_check_counter = 0 else: self._convergence_check_counter += 1 if self._convergence_check_counter >= self.convergence_threshold: logger.info(f"Early stopping triggered: No improvement for {self.convergence_threshold} iterations") # logger.info(f"Score history: {scores[-self.convergence_threshold:]}") return True return False def save(self, path: str, ignore: List[str] = []): """ Save the (optimized) workflow graph to a file. Args: path (str): The path to save the workflow graph. ignore (List[str]): The keys to ignore when saving the workflow graph. """ self.graph.save_module(path, ignore=ignore)