evoagentx/optimizers/aflow_optimizer.py

# Acknowledgement: Modified from AFlow (https://github.com/geekan/MetaGPT/blob/main/metagpt/ext/aflow/scripts/optimizer.py) under MIT License 

import os 
import re 
import shutil
import asyncio
import numpy as np
from tqdm import tqdm
from typing import List, Any
from pydantic import Field

from ..core.logging import logger
from ..core.module import BaseModule
from ..models.base_model import BaseLLM, LLMOutputParser
from ..benchmark.benchmark import Benchmark
from ..utils.aflow_utils.data_utils import DataUtils
from ..utils.aflow_utils.experience_utils import ExperienceUtils
from ..utils.aflow_utils.evaluation_utils import EvaluationUtils
from ..utils.aflow_utils.graph_utils import GraphUtils, OPERATOR_MAP
from ..utils.aflow_utils.convergence_utils import ConvergenceUtils


class GraphOptimizeOutput(LLMOutputParser):

    modification: str = Field(default="", description="modification")
    graph: str = Field(default="", description="graph")
    prompt: str = Field(default="", description="prompt")


class AFlowOptimizer(BaseModule):

    """
    AFlow Optimizer for workflow optimization.
    
    This optimizer iteratively improves workflows through multiple rounds of optimization
    using large language models. It evaluates workflow performance, identifies improvement
    opportunities, and applies optimizations based on experience and convergence metrics.
    
    Attributes:
        question_type: Type of task to optimize for (e.g., qa, match, code)
        graph_path: Path to the workflow graph directory (must contain graph.py and prompt.py)
        optimized_path: Path to save optimized workflows (defaults to graph_path)
        initial_round: Starting round number for optimization
        optimizer_llm: LLM used for generating optimizations
        executor_llm: LLM used for executing the workflow
        operators: List of operators available for optimization
        sample: Number of rounds to sample from for optimization
        max_rounds: Maximum number of optimization rounds to perform
        validation_rounds: Number of validation runs per optimization round
        eval_rounds: Number of evaluation runs for test mode
        check_convergence: Whether to check for optimization convergence
    """
    question_type: str = Field(description="The type of question to optimize the workflow for, e.g., qa, match, code, etc.")
    graph_path: str = Field(description="The folder of the workflow graph. This folder must contain a `graph.py` file that defines the workflow structure, and a `prompt.py` file that defines the prompt for the workflow.")
    optimized_path: str = Field(default=None, description="The path to save the optimized workflow. If not provided, the optimized path will be the same as the graph path.")
    initial_round: int = Field(default=0, description="The round number to start or continue optimization from. If not provided, will start from round 0 using the `graph.py` file in `graph_path`.")
    optimizer_llm: BaseLLM = Field(default=None, description="The LLM to use for optimization.")
    executor_llm: BaseLLM = Field(default=None, description="The LLM to use for execution.")

    operators: List[str] = Field(default_factory=lambda: list(OPERATOR_MAP.keys()), description="The operators to use for optimization. If not provided, will use all operators in OPERATOR_MAP.")
    sample: int = Field(default=4, description="The number of rounds to sample from the top scores.")
    max_rounds: int = Field(default=20, description="The maximum number of rounds to optimize the workflow.")
    validation_rounds: int = Field(default=5, description="Run the workflow for `validation_rounds` times to evaluate the performance on the validation set.")
    eval_rounds: int = Field(default=3, description="Run the workflow for `eval_rounds` times to evaluate the performance on the test set.")
    check_convergence: bool = Field(default=True, description="Whether to check for convergence.")

    def init_module(self, **kwargs):

        self.root_path = self.optimized_path or self.graph_path
        os.makedirs(self.root_path, exist_ok=True)

        # Initialize utilities
        self.graph_utils = GraphUtils(self.root_path)
        self.data_utils = DataUtils(self.root_path)
        self.evaluation_utils = EvaluationUtils(self.root_path)
        self.experience_utils = ExperienceUtils(self.root_path)
        self.convergence_utils = ConvergenceUtils(self.root_path)

        self.graph = None
        self.round = self.initial_round
        if self.round == 0:
            round_zero_path = os.path.join(self.root_path, f"round_{self.round}")
            os.makedirs(round_zero_path, exist_ok=True)
            shutil.copy2(os.path.join(self.graph_path, "graph.py"), os.path.join(round_zero_path, "graph.py"))
            shutil.copy2(os.path.join(self.graph_path, "prompt.py"), os.path.join(round_zero_path, "prompt.py"))
            self.graph_utils.update_prompt_import(os.path.join(round_zero_path, "graph.py"), round_zero_path)
        
        if not os.path.exists(os.path.join(self.root_path, f"round_{self.round}")):
            raise ValueError(f"Round {self.round} does not exist in {self.root_path}")
        
        if self.optimizer_llm is None:
            raise ValueError("optimizer_llm is not provided") 
        if self.executor_llm is None:
            self.executor_llm = self.optimizer_llm

    def optimize(self, benchmark: Benchmark):
        """Run the optimization process on the workflow.
        
        Performs multiple rounds of optimization, evaluating each round against
        the benchmark and checking for convergence. Continues until convergence
        is detected or the maximum number of rounds is reached.
        
        Args:
            benchmark: The benchmark to evaluate the workflow against
        """
        self.benchmark = benchmark
        for _ in range(self.max_rounds):
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            score = loop.run_until_complete(self._execute_with_retry(self._optimize_graph))
            self.round += 1
            logger.info(f"Score for round {self.round}: {score}")
            if self._check_convergence():
                break
            if self.round >= self.max_rounds:
                logger.info(f"Max rounds reached: {self.max_rounds}, stopping optimization.")
                break
    
    def test(self, benchmark: Benchmark, test_rounds: List[int] = None):
        """Run the test evaluation on optimized workflows.
        
        Evaluates specified rounds (or the best round if none specified) against
        the benchmark multiple times and logs the results.
        
        Args:
            benchmark: The benchmark to evaluate against
            test_rounds: Specific round numbers to test, or None to use the best round
        """
        self.benchmark = benchmark
        if test_rounds is None:
            best_round = self._load_best_round()
            logger.info(f"No test rounds provided, using best round: {best_round}")
            test_rounds = [best_round]
        for _ in tqdm(range(self.eval_rounds)):
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            loop.run_until_complete(self._run_test(test_rounds))
    
    async def _execute_with_retry(self, func: callable, max_retries: int = 3) -> Any:

        retry_count = 0
        while retry_count < max_retries:
            try:
                return await func()
            except Exception as e:
                retry_count += 1
                logger.info(f"Error occurred: {e}. Retrying... (Attempt {retry_count}/{max_retries})")
                if retry_count == max_retries:
                    logger.info("Max retries reached.")
                    return None
                await asyncio.sleep(5 * retry_count)
        
        return None
    
    def _check_convergence(self) -> bool:
        if not self.check_convergence:
            return False

        converged, convergence_round, final_round = self.convergence_utils.check_convergence(top_k=3)
        if converged:
            logger.info(
                f"Convergence detected, occurred in round {convergence_round}, final round is {final_round}"
            )
            self.convergence_utils.print_results()
            return True
        return False

    async def _optimize_graph(self) -> float:
        """Optimize the graph for one round"""
        validation_n = self.validation_rounds
        graph_path = self.root_path
        data = self.data_utils.load_results(graph_path)

        if self.round == 0:
            self.avg_score = await self._handle_initial_round(graph_path, validation_n, data)
        
        return await self._handle_optimization_round(graph_path, validation_n, data)

    async def _handle_initial_round(self, graph_path: str, validation_n: int, data: list) -> float:
        """Handle the initial round of optimization"""
        self.graph_utils.create_round_directory(graph_path, self.round)
        self.graph = self.graph_utils.load_graph(self.round, graph_path)
        return await self.evaluation_utils.evaluate_graph_async(self, validation_n, data, initial=True)

    async def _handle_optimization_round(self, graph_path: str, validation_n: int, data: list) -> float:

        directory = self.graph_utils.create_round_directory(graph_path, self.round + 1)

        while True:
            sample = self._get_optimization_sample()
            prompt, graph_load = self.graph_utils.read_graph_files(sample["round"], graph_path)
            graph = self.graph_utils.extract_solve_graph(graph_load)
            processed_experience = self.experience_utils.load_experience()
            experience = self.experience_utils.format_experience(processed_experience, sample["round"])
            operator_description = self.graph_utils.load_operators_description(self.operators, self.optimizer_llm)
            log_data = self.data_utils.load_log(sample["round"])
            graph_optimize_prompt = self.graph_utils.create_graph_optimize_prompt(
                experience, sample["score"], graph[0], prompt, operator_description, self.question_type, log_data
            )
            # response = await self.optimizer_llm.async_generate(prompt=graph_optimize_prompt, parser=GraphOptimizeOutput, parse_mode="xml")
            # response = response.get_structured_data()
            response = await self.optimizer_llm.async_generate(prompt=graph_optimize_prompt, parse_mode="str")
            print(response.content)
            try:
                parsed_response = GraphOptimizeOutput.parse(response.content, parse_mode="xml")
                response = parsed_response.get_structured_data()
            except Exception:
                response = self._parse_optimizer_llm_output(response.content, orig_graph=graph[0], orig_prompt=prompt)

            if self.experience_utils.check_modification(processed_experience, response['modification'], sample["round"]):
                break
        
        # Save and evaluate results
        avg_score = await self._evaluate_and_save_optimization_results(directory, response, sample, data, validation_n)
        return avg_score
    
    def _get_optimization_sample(self) -> dict:

        top_rounds = self.data_utils.get_top_rounds(self.sample)
        return self.data_utils.select_round(top_rounds)

    def _parse_optimizer_llm_output(self, content: str, orig_graph: str, orig_prompt: str) -> dict:

        response = {"modification": "", "graph": "", "prompt": ""}

        # Extract content between <modification> tags
        modification_pattern = r'<modification>(.*?)</modification>'
        modification_match = re.search(modification_pattern, content, re.DOTALL)
        if modification_match:
            response["modification"] = modification_match.group(1).strip()
        
        # extract code block
        code_block_pattern = r'```(?:python)?(.*?)```'
        code_blocks = re.finditer(code_block_pattern, content, re.DOTALL)

        # Process found code blocks
        for block in code_blocks:
            code = block.group(1).strip()
            # If code contains graph-related content, store in graph
            if 'class' in code or 'workflow' in code.lower():
                response["graph"] = code
            # If code contains prompt-related content, store in prompt
            # elif 'PROMPT' in code or 'prompt' in code.lower():
            #     response["prompt"] = code
            else:
                response["prompt"] = code
        
        if not response["graph"] and not response["prompt"]:
            response["modification"] = "No modification due to error in LLM output"
            response["graph"] = orig_graph
            response["prompt"] = orig_prompt 
        
        return response
    
    async def _evaluate_and_save_optimization_results(self, directory: str, response: dict, sample: dict, data: list, validation_n: int):

        # Write optimized files
        self.graph_utils.write_graph_files(directory, response)

        experience = self.experience_utils.create_experience_data(sample, response['modification'])

        self.graph = self.graph_utils.load_graph(self.round + 1, self.root_path)

        # evaluate the graph 
        avg_score = await self.evaluation_utils.evaluate_graph_async(self, validation_n, data, initial=False)
        self.experience_utils.update_experience(directory, experience, avg_score)

        return avg_score

    def _load_best_round(self) -> int:
        """Load the best round"""
        ranked_scores = self.data_utils._load_scores()
        return ranked_scores[0]["round"]

    async def _run_test(self, test_rounds: List[int]):
        """Run test evaluation"""

        logger.info("Running test evaluation...")

        graph_path = self.root_path
        data = self.data_utils.load_results(graph_path)
        json_file_path = self.data_utils.get_results_file_path(graph_path)
        scores = [] 

        # for round in tqdm(test_rounds, desc="Testing"):
        for round in test_rounds:

            logger.info(f"Running test for round {round}...")

            self.graph = self.graph_utils.load_graph(round, graph_path)

            score, avg_cost, total_cost = await self.evaluation_utils.evaluate_graph_test_async(self)
            scores.append(score)

            new_data = self.data_utils.create_result_data(round, score, avg_cost, total_cost)
            data.append(new_data)

            logger.info(f"Test round {round} score: {score}, avg_cost: {avg_cost}, total_cost: {total_cost}")

            self.data_utils.save_results(json_file_path, data)

        logger.info(f"Test round {round} avg_score: {np.mean(scores)}")
        return np.mean(scores)