evoagentx/optimizers/textgrad_optimizer.py

# -----------------------------------------------------------------------------
# This file re-implements algorithms from the TextGrad project:
#   Repo: https://github.com/zou-group/textgrad
#   Paper: "TextGrad: Automatic 'Differentiation' via Text -- using large language models to backpropagate textual gradients."
#   Authors: Zou 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/
# -----------------------------------------------------------------------------

# ruff: noqa: E402
import os
import shutil
from copy import deepcopy
from typing import Any, Iterator, List, Literal, Optional, Tuple, Union

import numpy as np
from pydantic import Field, PositiveInt
from tqdm import tqdm

from ..agents import Agent, CustomizeAgent
from ..benchmark.benchmark import Benchmark, CodingBenchmark
from ..core.callbacks import suppress_logger_info
from ..core.logging import logger
from ..core.module import BaseModule
from ..evaluators import Evaluator
from ..models.base_model import BaseLLM
from ..prompts import PromptTemplate
from ..workflow.workflow_graph import WorkFlowGraph, WorkFlowNode

# Check if logs folder exists before importing textgrad
log_folder_exists = os.path.exists("./logs")

import textgrad as tg
from textgrad import Variable, EngineLM
from textgrad import logger as tg_logger
from textgrad import sh as tg_file_handler
from textgrad.autograd import StringBasedFunction
from textgrad.loss import MultiFieldEvaluation, TextLoss
from textgrad.optimizer import TextualGradientDescent

from ..prompts.optimizers.textgrad_optimizer import (
    CODE_LOSS_PROMPT,
    GENERAL_LOSS_PROMPT,
    NO_ANSWER_LOSS_PROMPT,
    OPTIMIZER_CONSTRAINTS,
    OPTIMIZER_SYSTEM_PROMPT,
    PERSONAL_FINANCE_ADVISOR_EXAMPLE,
    FITNESS_COACH_EXAMPLE,
    CODE_REVIEW_EXAMPLE,
)

tg_logger.removeHandler(tg_file_handler)
# remove the logs folder created by textgrad
if not log_folder_exists and os.path.exists("./logs"):
    shutil.rmtree("./logs")


class TextGradEngine(EngineLM):
    def __init__(self, llm: BaseLLM):
        self.llm = llm

    def generate(self, prompt: str, system_prompt: str = None, **kwargs):
        with suppress_logger_info():
            response = self.llm.generate(prompt, system_prompt=system_prompt, **kwargs)
            return response.content

    def __call__(self, prompt: str, **kwargs):
        return self.generate(prompt, **kwargs)


class CustomAgentCall:
    """A custom agent call with textgrad.Variable inputs and output."""
    def __init__(self, agent: Agent):
        self.agent = agent
        self.last_outputs: dict[str, str] = dict()
    
    def __call__(
        self, 
        instruction: Variable, 
        system_prompt: Variable,
        **inputs: Variable
    ) -> Variable:

        action = self.agent.actions[0]
        input_names = action.inputs_format.get_attrs()

        agent_inputs = {}

        for key, input_variable in inputs.items():
            if key in input_names:
                agent_inputs[key] = input_variable.value
            else:
                parsed_inputs: dict[str, str] = {k:v for k,v in input_variable.parsed_outputs.items() if k in input_names}
                agent_inputs.update(parsed_inputs)

        with suppress_logger_info():
            outputs = self.agent.execute(action_name=action.name, action_input_data=agent_inputs).content

        self.last_outputs = outputs.to_dict()

        return outputs.content


class TextGradAgent:
    """An agent that takes textgrad.Variable inputs and returns a textgrad.Variable response.
    This class is used to replace EvoAgentX Agent in WorkFlowGraph to allow TextGrad optimization.
    """
    def __init__(
        self, 
        agent: Agent,
        optimize_mode: Literal["all", "system_prompt", "instruction"] = "all"
    ):
        self.name = agent.name

        require_grad = {
            "all": {"system_prompt": True, "instruction": True}, 
            "system_prompt": {"system_prompt": True, "instruction": False}, 
            "instruction": {"system_prompt": False, "instruction": True}
        }

        system_prompt_require_grad = require_grad[optimize_mode]["system_prompt"]
        instruction_require_grad = require_grad[optimize_mode]["instruction"]

        self.system_prompt = Variable(
            agent.system_prompt, 
            requires_grad=system_prompt_require_grad, 
            role_description=f"{self.name}'s system prompt"
        )
        
        self.instruction = Variable(
            agent.actions[0].prompt_template.instruction,
            requires_grad=instruction_require_grad, 
            role_description=f"{self.name}'s instruction prompt"
        )

        self._agent_call = CustomAgentCall(agent)
        self.forward = StringBasedFunction(self._agent_call, agent.description)
        self.output_description = " and ".join(agent.actions[0].outputs_format.get_attr_descriptions().values())
        self.last_output = None

    def __call__(self, inputs: dict[str, Variable]) -> Variable:
        """Given textgrad.Variable inputs, generates a textgrad.Variable output."""

        forward_inputs: dict[str, Variable] = {
            "instruction": self.instruction, 
            "system_prompt": self.system_prompt,
            **inputs
        }

        output_variable = self.forward(forward_inputs, self.output_description)
        output_variable.parsed_outputs = self._agent_call.last_outputs
        self.last_output = output_variable
        return output_variable



class TextGradOptimizer(BaseModule):
    """Uses TextGrad to optimize agents' system prompts and instructions in a multi-agent workflow.
    For more information on TextGrad, see https://github.com/zou-group/textgrad.
    """
    graph: WorkFlowGraph = Field(description="The workflow to optimize.")
    optimize_mode: Literal["all", "system_prompt", "instruction"] = Field(default="all", description="The mode to optimize the workflow. 'all' optimizes both system prompts and instructions, 'system_prompt' only optimizes system prompts, and 'instruction' only optimizes instructions.")
    executor_llm: BaseLLM = Field(default=None, description="The LLM to use for execution.")
    optimizer_llm: BaseLLM = Field(default=None, description="The LLM to use for optimization.")
    batch_size: PositiveInt = Field(default=1, description="The batch size for optimization.")
    max_steps: PositiveInt = Field(default=10, description="The maximum number of optimization steps.")
    evaluator: Evaluator = Field(default=None, description="The evaluator to perform evaluation during optimization.")
    eval_every_n_steps: Optional[PositiveInt] = Field(default=None, description="Evaluate the workflow every `eval_every_n_steps` steps.")
    eval_rounds: PositiveInt = Field(default=1, description="The number of times to evaluate the performance.")
    eval_config: dict = Field(default={}, description="The configuration for evaluation. The keys are the arguments of `TextGradOptimizer.evaluate`.")
    save_interval: Optional[PositiveInt] = Field(default=None, description="Save the workflow every `save_interval` steps.")
    save_path: str = Field(default="./", description="The path to save the optimized workflow.")
    rollback: bool = Field(default=True, description="Whether to rollback to the best graph after each evaluation during optimization.")
    constraints: List[str] = Field(default=[], description="The constraints for optimization. e.g. ['They system prompt must not exceed 100 words.']")


    def init_module(self, **kwargs):
        self._validate_graph_compatibility(self.graph)
        self._snapshot: List[dict] = []
        self.output_lookup = self._create_output_lookup()
        

    def _init_textgrad(self, dataset: Benchmark, use_answers: bool = True):
        # Disable TextGrad's short variable value to allow the optimizer to receive the full variable value
        def disable_short_variable_value(self, n_words_offset: int = 10):
            return self.value
        Variable.get_short_value = disable_short_variable_value
            
        # Textgrad engine
        self.optimizer_engine = TextGradEngine(self.optimizer_llm)

        # Textgrad loss
        if use_answers:
            if isinstance(dataset, CodingBenchmark):
                loss_prompt = CODE_LOSS_PROMPT
                role_descriptions = ["code snippet to evaluate", "the task, the test result of the code snippet, and the correct code"]
            else:
                loss_prompt = GENERAL_LOSS_PROMPT
                role_descriptions = ["response to evaluate", "correct answer"]

            evaluation_instruction = Variable(loss_prompt, requires_grad=False, role_description="evaluation instruction")
            self.loss_fn = MultiFieldEvaluation(evaluation_instruction, role_descriptions, self.optimizer_engine)
        else:
            loss_prompt = NO_ANSWER_LOSS_PROMPT
            evaluation_instruction = Variable(loss_prompt, requires_grad=False, role_description="evaluation instruction")
            self.loss_fn = TextLoss(evaluation_instruction, self.optimizer_engine)

        # Create textgrad agents
        self._create_textgrad_agents()
        
        # Textgrad optimizer
        if self.optimize_mode == "all":
            optimize_variables = self._get_all_system_prompts() + self._get_all_instructions()
        elif self.optimize_mode == "system_prompt":
            optimize_variables = self._get_all_system_prompts()
        elif self.optimize_mode == "instruction":
            optimize_variables = self._get_all_instructions()
        else:
            raise ValueError("Unsupported `optimize_mode`, should be one of 'all', 'system_prompt', 'instruction'.")

        OPTIMIZER_CONSTRAINTS.extend(self.constraints)
        
        self.textgrad_optimizer = TextualGradientDescent(
            parameters=optimize_variables, 
            engine=self.optimizer_engine,
            constraints=OPTIMIZER_CONSTRAINTS,
            optimizer_system_prompt=OPTIMIZER_SYSTEM_PROMPT,
            in_context_examples=[PERSONAL_FINANCE_ADVISOR_EXAMPLE, FITNESS_COACH_EXAMPLE, CODE_REVIEW_EXAMPLE]
        )


    def optimize(self, dataset: Benchmark, use_answers: bool = True, seed: Optional[int] = None) -> None:
        """Optimizes self.graph using `dataset`.
        
        Args:
            dataset (Benchmark): The dataset to use for optimization.
            use_answers (bool): Whether to use the answers (labels) in the training set for optimization.
                If False, `dataset`'s training set does not need to have answers.
                If `eval_every_n_steps` is set to None, we can optimize the workflow without any labeled data.
            seed (Optional[int]): The random seed to use for shuffling the data.
        """
        self._init_textgrad(dataset, use_answers)
 
        def iterator() -> Iterator[Tuple[List[dict[str, str]],  Optional[List[Union[str, dict[str, str]]]]]]:
            epoch = 0
            while True:
                # Shuffle train data every epoch
                effective_seed = seed + epoch if seed is not None else None
                train_data = dataset.get_train_data(sample_k=len(dataset._train_data), seed=effective_seed)
                for i in range(0, len(train_data), self.batch_size):
                    batch = train_data[i:i + self.batch_size]
                    inputs = [self.evaluator.collate_func(x) for x in batch]
                    if use_answers:
                        labels = dataset.get_labels(batch)
                    else:
                        labels = None
                    yield inputs, labels
                epoch += 1

        data_iterator = iterator()

        for step in tqdm(range(self.max_steps)):
            inputs, labels = next(data_iterator)
            self.step(inputs, labels, dataset, use_answers)

            if self.eval_every_n_steps is not None and (step + 1) % self.eval_every_n_steps == 0:
                logger.info(f"Evaluating the workflow at step {step+1} ...")
                with suppress_logger_info():
                    metrics = self.evaluate(dataset, **self.eval_config)
                self.log_snapshot(self.graph, metrics)
                logger.info(f"Step {step+1} metrics: {metrics}")

                # If rollback is enabled, keep track of the best snapshot
                if self.rollback:
                    if len(self._snapshot) == 1:
                        best_snapshot = self._snapshot[-1]
                        best_average_score = np.mean(list(metrics.values()))
                    else:
                        current_average_score = np.mean(list(metrics.values()))
                        
                        if current_average_score >= best_average_score:
                            # If the current average score is better than the best average score, update the best snapshot
                            best_snapshot = self._snapshot[-1]
                            best_average_score = current_average_score
                        else:
                            # If the current average score is worse than the best average score, roll back to the best snapshot
                            logger.info(f"Metrics are worse than the best snapshot which has {best_snapshot['metrics']}. Rolling back to the best snapshot.")
                            best_graph = WorkFlowGraph.from_dict(best_snapshot["graph"])
                            self.graph = best_graph
                            self._create_textgrad_agents()

            if self.save_interval is not None and (step + 1) % self.save_interval == 0:
                logger.info(f"Saving the workflow at step {step+1} ...")
                self.save(os.path.join(self.save_path, f"{dataset.name}_textgrad_step_{step+1}.json"))

        logger.info(f"Reached the maximum number of steps {self.max_steps}. Optimization has finished.")
        self.save(os.path.join(self.save_path, f"{dataset.name}_textgrad_final.json"))

        # Saves the best graph
        if len(self._snapshot) > 0:
            best_graph = self._select_graph_with_highest_score()
            self.save(os.path.join(self.save_path, f"{dataset.name}_textgrad_best.json"), graph=best_graph)

        
    def step(
        self, 
        inputs: list[dict[str, str]], 
        labels: Optional[list[Union[str, dict[str, str]]]], 
        dataset: Benchmark, 
        use_answers: bool = True
    ) -> None:
        """Performs one optimization step using a batch of data."""

        losses = []
        logger.info("Executing workflow...")

        if use_answers:
            if labels is None:
                raise ValueError("Labels must be provided if `use_answers` is True.")

            for input, label in zip(inputs, labels, strict=True):
                output = self.forward(input)
                if isinstance(label, str):
                    label = Variable(label, requires_grad=False, role_description="correct answer for the query")
                elif isinstance(label, dict):
                    if not isinstance(dataset, CodingBenchmark):
                        raise ValueError("Label must be a string for non-coding benchmarks.")
                    end_node_name = self.graph.find_end_nodes()[0]
                    end_node = self.graph.get_node(end_node_name)
                    output_name = end_node.outputs[0].name
                    code = output.parsed_outputs[output_name]
                    label = self._format_code_label(code, label, dataset)
                    label = Variable(label, requires_grad=False, role_description="the task, the test result, and the correct code")
                loss = self.loss_fn([output, label])
                losses.append(loss)
        else:
            for input in inputs:
                output = self.forward(input)
                loss = self.loss_fn(output)
                losses.append(loss)

        total_loss = tg.sum(losses)
        logger.info("Computing gradients...")
        total_loss.backward(self.optimizer_engine)
        logger.info("Updating agents...")
        self.textgrad_optimizer.step()
        self.textgrad_optimizer.zero_grad()
        self._update_workflow_graph()
        logger.info("Agents updated")


    def forward(self, inputs: dict[str, str]) -> Variable:
        """Returns the final output from the workflow."""
        self._visited_nodes = set()
        end_node = self.graph.find_end_nodes()[0]
        input_variables = self._initial_inputs_to_variables(inputs)
        output = self._compute_node(end_node, input_variables)
        return output


    def evaluate(
        self, 
        dataset: Benchmark, 
        eval_mode: str = "dev", 
        graph: Optional[WorkFlowGraph] = 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 (WorkFlowGraph, 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.
        """
        if graph is None:
            graph = self.graph

        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))
            metrics = self.evaluator.evaluate(
                graph=graph, 
                benchmark=dataset, 
                eval_mode=eval_mode, 
                indices=indices, 
                sample_k=sample_k,
                update_agents=True, 
                **kwargs
            )
            metrics_list.append(metrics)
        avg_metrics = self.evaluator._calculate_average_score(metrics_list)
        
        return avg_metrics


    def save(self, path: str, graph: Optional[WorkFlowGraph] = None, ignore: List[str] = []) -> None:
        """Save the workflow graph containing the optimized prompts to a file. 

        Args:
            path (str): The path to save the workflow graph.
            graph (WorkFlowGraph, optional): The graph to save. If not provided, use the graph in the optimizer.
            ignore (List[str]): The keys to ignore when saving the workflow graph.
        """
        if graph is None:
            graph = self.graph
        graph.save_module(path, ignore=ignore)


    def log_snapshot(self, graph: WorkFlowGraph, metrics: dict) -> None:
        """Log the snapshot of the workflow."""
        self._snapshot.append(
            {
                "index": len(self._snapshot),
                "graph": deepcopy(graph.get_config()),
                "metrics": metrics,
            }
        )


    def restore_best_graph(self) -> None:
        """Restore the best graph from the snapshot and set it to `self.graph`."""
        if len(self._snapshot) == 0:
            logger.info("No snapshot found. No graph to restore.")
            return

        best_graph, best_metrics = self._select_graph_with_highest_score(return_metrics=True)
        self.graph = best_graph
        logger.info(f"Restored the best graph from snapshot with metrics {best_metrics}")


    def _format_code_label(self, code: str, label:dict[str, str], dataset: CodingBenchmark) -> str:
        """Formats the label for coding tasks to include the task, the test result, and the correct code.

        Args:
            code: The code to evaluate.
            label: A dictionary with keys "task_id", "test", "entry_point", and "canonical_solution".
            dataset: A CodingBenchmark instance with `check_solution` method.
        
        Returns:
            The formatted label which includes the task, the test result, and the correct code.
        """

        task_id = label["task_id"]
        prompt = dataset.get_example_by_id(task_id)["prompt"]
        if 'test' not in label.keys():
            label['test'] = label['tests']
        test = label["test"]
        entry_point = label["entry_point"]

        state, message = dataset.check_solution(
            task_id=task_id,
            solution=prompt + "\n" + code,
            test=test,
            entry_point=entry_point
        )

        if state != dataset.SUCCESS:
            message = message.replace("Solution", "Failed Code")

        formatted_label = f"## Task:\n{prompt}\n\n## Result on test:\n{message}\n\n## Correct Solution:\n{label['canonical_solution']}"
        return formatted_label


    def _initial_inputs_to_variables(self, initial_inputs: dict[str, str]) -> dict[str, Variable]:
        """Converts inputs to the initial nodes to textgrad variables."""
        variables = {}
        initial_nodes = self.graph.find_initial_nodes()
        for initial_node in initial_nodes:
            for key, value in initial_inputs.items():
                for input in self.graph.get_node(initial_node).inputs:
                    if input.name == key:
                        initial_input_variable = Variable(
                            value,
                            requires_grad=False,
                            role_description=input.description,
                        )
                        variables[key] = initial_input_variable
                        if len(variables) == len(initial_inputs):
                            return variables
        missing_inputs = set(initial_inputs.keys()) - set(variables.keys())
        raise ValueError(f"Initial inputs do not match the inputs of the initial nodes. Missing inputs: {missing_inputs}")



    def _compute_node(self, node: Union[str, WorkFlowNode], initial_inputs: dict[str, Variable]) -> Variable:
        """Computes the output of a node in the workflow graph by recursively computing the required inputs.

        Args:
            node: The node to compute the output of.
            initial_inputs: The initial inputs to the workflow that are not from any node in the workflow (e.g., user query).

        Returns:
            The output of the node as a textgrad.Variable.
        """
        if isinstance(node, str):
            node = self.graph.get_node(node)

        if node.name in self._visited_nodes:
            return node.textgrad_agent.last_output

        input_variables: dict[str, Variable] = {}    # inputs to TextGradAgent
        input_node_names: set[str] = set()           # which nodes we need to compute the output of      

        for input in node.inputs:
            if input.name in initial_inputs:
                input_variables[input.name] = initial_inputs[input.name]
            else:
                input_node_names.add(self.output_lookup[input.name])
        
        # if the input is from another node, compute the output of that node
        for node_name in input_node_names:
            input_variables[node_name] = self._compute_node(node_name, initial_inputs)

        output_variable = node.textgrad_agent(input_variables)
        self._visited_nodes.add(node.name)
        return output_variable
        
   
    def _create_textgrad_agent(self, node: Union[str, WorkFlowNode]) -> TextGradAgent:
        """Creates a textgrad agent for a given node in a WorkFlowGraph."""
        if isinstance(node, str):
            node = self.graph.get_node(node)

        if isinstance(node.agents[0], dict):
            agent_llm = node.agents[0].get("llm")
            agent_llm_config = node.agents[0].get("llm_config")
            if agent_llm is None and agent_llm_config is None:
                node.agents[0]["llm"] = self.executor_llm
            # CustomizeAgent.from_dict creates a CustomizeAgent if dict follows CustomizeAgent format
            # creates an Agent if dict follows Agent format
            agent: Union[CustomizeAgent, Agent] = CustomizeAgent.from_dict(node.agents[0])
        else:
            raise ValueError(f"Unsupported agent type {type(node.agents[0])}. Expected 'dict'.")
            
        textgrad_agent = TextGradAgent(agent, self.optimize_mode)
        return textgrad_agent


    def _create_textgrad_agents(self):
        """Creates textgrad agents for all nodes in the workflow graph."""
        for node in self.graph.nodes:
            node.textgrad_agent = self._create_textgrad_agent(node)


    def _update_agent_prompts(self, agent_dict: dict[str, Any], system_prompt: str, instruction: str) -> dict[str, Any]:
        agent_dict["system_prompt"] = system_prompt
        if "actions" in agent_dict:
            # Agent has actions in its dict
            agent_dict["actions"][0]["prompt_template"] = self._update_agent_instructions(
                agent_dict["actions"][0]["prompt_template"], 
                instruction
            )
        else:
            # CustomizeAgent does not have actions in its dict
            agent_dict["prompt_template"] = self._update_agent_instructions(
                agent_dict["prompt_template"], 
                instruction
            )
        return agent_dict


    def _update_agent_instructions(
        self, 
        prompt_template: Union[PromptTemplate, dict[str, str]], 
        instruction: str
    ) -> Union[PromptTemplate, dict[str, str]]:

        if isinstance(prompt_template, PromptTemplate):
            prompt_template.set_instruction(instruction)
        elif isinstance(prompt_template, dict):
            prompt_template["instruction"] = instruction
        else:
            raise ValueError(f"Unsupported prompt template type {type(prompt_template)}. Expected 'PromptTemplate' or 'dict'.")
        return prompt_template


    def _update_workflow_graph(self):
        """Updates the workflow graph with the latest prompts from the textgrad optimization."""
        for node in self.graph.nodes:
            if isinstance(node.agents[0], dict):
                node.agents[0] = self._update_agent_prompts(
                    node.agents[0], 
                    node.textgrad_agent.system_prompt.value, 
                    node.textgrad_agent.instruction.value
                )
            else:
                raise ValueError(f"Unsupported agent type {type(node.agents[0])}. Expected 'dict'.")


    def _select_graph_with_highest_score(self, return_metrics: bool = False) -> Union[WorkFlowGraph, tuple[WorkFlowGraph, Optional[dict]]]:
        """Select the graph in `self._snapshot` with the highest score."""
        if len(self._snapshot) == 0:
            if return_metrics:
                return self.graph, None
            return self.graph
            
        snapshot_scores = [np.mean(list(snapshot["metrics"].values())) for snapshot in self._snapshot]
        best_index = np.argmax(snapshot_scores)

        graph = WorkFlowGraph.from_dict(self._snapshot[best_index]["graph"])

        if return_metrics:
            return graph, self._snapshot[best_index]["metrics"]
        return graph


    def _get_all_system_prompts(self) -> List[Variable]:
        """Gets all system prompts from the textgrad agents."""
        system_prompts = []
        for node in self.graph.nodes:
            system_prompts.append(node.textgrad_agent.system_prompt)
        return system_prompts

    
    def _get_all_instructions(self) -> List[Variable]:
        """Gets all prompt templates from the textgrad agents."""
        instructions = []
        for node in self.graph.nodes:
            instructions.append(node.textgrad_agent.instruction)
        return instructions


    def _create_output_lookup(self) -> dict[str, str]:
        """Creates a lookup table for output names to node names."""
        output_name_to_node_name = {}
        for node in self.graph.nodes:
            for output in node.outputs:
                output_name_to_node_name[output.name] = node.name
        return output_name_to_node_name


    def _validate_graph_compatibility(self, graph: WorkFlowGraph) -> None:
        """Checks if the graph is compatible with the textgrad optimizer."""
        for node in graph.nodes:
            if len(node.agents) > 1:
                raise ValueError("TextGrad optimizer 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"TextGrad optimizer 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]:
                            raise ValueError(f"Please provide a PromptTemplate for {agent['name']}.")
                    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']}.")