optimizer works well, and it knows a cycle if it exists, and the data…

… and feedback will be present in sequential order

adalflow/adalflow/core/generator.py

@@ -773,7 +773,7 @@ def _backward_through_one_predecessor(
         )
         var_gradient.add_context(
             GradientContext(
-                context=conversation_str,
+                input_output=conversation_str,
                 response_desc=response.role_desc,
                 variable_desc=pred.role_desc,  # parameter_desc
             )

adalflow/adalflow/optim/grad_component.py

@@ -189,7 +189,7 @@ def backward(self, *, response: "Parameter", id: str = None, **kwargs):
                         variable_desc=pred.role_desc,
                         response_desc=response.name,
                         # context=f"""""",  # TODO: check the forward pass component trace
-                        context=f"""{response.component_trace}""",
+                        input_output=f"""{response.component_trace}""",
                     )
                 )
 

adalflow/adalflow/optim/parameter.py

@@ -30,13 +30,19 @@
 
 
 @dataclass
-class GradientContext:
+class GradientContext(DataClass):
+    """GradientContext is used to describe the component's function and trace its input and output.
+
+    To get the component's function desc, use GradientContext.to_yaml_signature()
+    To get the data: use instance.to_yaml()
+    """
+
     variable_desc: str = field(
         metadata={"desc": "The description of the target parameter"}
     )
     # from template LOSS_CONVERSATION_TEMPLATE_STRING
     # LLM_CONVERSATION_TEMPLATE from backward_engine_prompt
-    context: str = field(
+    input_output: str = field(
         metadata={
             "desc": "The context of the gradient in form of a conversation indicating \
                 the relation of the current parameter to the response parameter"
@@ -47,17 +53,6 @@ class GradientContext:
     )
 
 
-@dataclass(frozen=True)
-class ComponentNode(DataClass):
-    """Used to represent a node in the component graph."""
-
-    id: str = field(metadata={"desc": "The unique id of the component"})
-    name: str = field(metadata={"desc": "The name of the component"})
-    type: Literal["INPUT", "COMPONENT"] = field(
-        metadata={"desc": "The type of the node"}, default="COMPONENT"
-    )
-
-
 @dataclass
 class ComponentTrace(DataClass):
     name: str = field(metadata={"desc": "The name of the component"}, default=None)
@@ -92,48 +87,51 @@ class ScoreTrace:
     )
 
 
-# COMBINED_GRADIENTS_TEMPLATE = r"""
-# {% if combined_gradients %}
-# Batch size: {{ combined_gradients|length }}
-# {% endif %}
-# {% for g in combined_gradients %}
-# {% set gradient = g[0] %}
-# {% set gradient_context = g[1] %}
-
-# {% if gradient_context %}
-# {{loop.index}}.
-# <CONTEXT>{{gradient_context.context}}</CONTEXT>
-# {% endif %}
-
-# {% if gradient.data %}
-#   {% if gradient_context %}
-# {#The output is used as <{{gradient_context.response_desc}}>#}
-# <FEEDBACK>{{gradient.data}}</FEEDBACK>
-# {% else %}
-# <FEEDBACK>{{gradient.data}}</FEEDBACK>
-# {% endif %}
-# {% endif %}
-# {% endfor %}"""
+@dataclass(frozen=True)
+class ComponentNode(DataClass):
+    """Used to represent a node in the component graph."""
+
+    id: str = field(metadata={"desc": "The unique id of the component"})
+    name: str = field(metadata={"desc": "The name of the component"})
+    type: Literal["INPUT", "COMPONENT"] = field(
+        metadata={"desc": "The type of the node"}, default="COMPONENT"
+    )
 
 
 COMBINED_GRADIENTS_TEMPLATE = r"""
-{% if combined_gradients %}
-Batch size: {{ combined_gradients|length }}
+{% if component_schema %}
+<COMPONENT_SCHEMA>
+Gradients are from {{ component_schema | length }} components.
+{% for component_id, schema in component_schema.items() %}
+id: {{ component_id }}
+{{ schema }}
+{% endfor %}
+</COMPONENT_SCHEMA>
 {% endif %}
+<DESCRIPTION>
+If same data_id appears multiple times, it means this component/variable is called multiple times in the same order as it appears in the gradient list.
+Use this info to have more clarity while reasoning and proposing new variable
+</DESCRIPTION>
+{% if combined_gradients %}
 {% for g in combined_gradients %}
 {% set gradient = g %}
 
-{% if gradient.context %}
+{% if gradient['context'] %}
 {{loop.index}}.
-<CONTEXT>{{gradient.context}}</CONTEXT>
+data_id: {{gradient['data_id']}}
+INPUT_OUTPUT: {{gradient['context']}}
 {% endif %}
 
-{% if gradient.data %}
+{% if gradient['score'] is not none %}
 {#The output is used as <{{gradient_context.response_desc}}>#}
-<SCORE>{{gradient.score}}</SCORE>
-<FEEDBACK>{{gradient.data}}</FEEDBACK>
+<SCORE>{{gradient['score']}}</SCORE>
+<FEEDBACK>{{gradient['gradient']}}</FEEDBACK>
+{% endif %}
+{% endfor %}
 {% endif %}
-{% endfor %}"""
+"""
+
+# Batch size: {{ combined_gradients|length }}
 
 
 class Parameter(Generic[T]):
@@ -321,9 +319,74 @@ def get_gradient_and_context_text(self, skip_correct_sample: bool = False) -> st
             return ""
 
         # sore gradients by the _score from low to high
-        self.gradients = sorted(
-            self.gradients, key=lambda x: x.score if x.score is not None else 1
-        )
+        # self.gradients = sorted(
+        #     self.gradients, key=lambda x: x.score if x.score is not None else 1
+        # )
+        # print the score for the sorted gradients
+        lowest_score_gradients = []
+        for i, g in enumerate(self.gradients):
+            if skip_correct_sample:
+                if g.score > 0.5:
+                    continue
+            lowest_score_gradients.append(g)
+            print(f"{i} Score: {g.score} for {g.name}, {type(g.score)}")
+
+        gradient_context_combined_str = ""
+        if lowest_score_gradients and len(lowest_score_gradients) > 0:
+
+            # parse the gradients and context.
+            gradients_and_context: List[Dict[str, Any]] = (
+                []
+            )  # {gradient: data, context: GradientContext.input_output}
+            for g in lowest_score_gradients:
+                gradients_and_context.append(
+                    {
+                        "data_id": g.data_id,
+                        "gradient": g.data,
+                        "context": g.context.input_output,
+                        "score": g.score,
+                    }
+                )
+
+            gradient_context_combined_str = Prompt(
+                template=COMBINED_GRADIENTS_TEMPLATE,
+                prompt_kwargs={"combined_gradients": gradients_and_context},
+            )().strip()
+
+        # get component id: gradient
+        component_id_to_gradient: Dict[str, Gradient] = {}
+        for g in lowest_score_gradients:
+            component_id_to_gradient[g.from_response_component_id] = g
+
+        componend_id_to_schema: Dict[str, str] = {}
+        for id, g in component_id_to_gradient.items():
+            componend_id_to_schema[id] = g.context.to_yaml(exclude={"input_output"})
+
+        # if there are multiple successors, there will be multiple component schemas
+
+        return gradient_context_combined_str
+
+    def get_gradients_component_schema(self, skip_correct_sample: bool = False) -> str:
+        """Aggregates and returns:
+        1. the gradients
+        2. the context text for which the gradients are computed
+
+        Sort the gradients from the lowest score to the highest score.
+        Highlight the gradients with the lowest score to the optimizer.
+        """
+        from adalflow.core.prompt_builder import Prompt
+
+        # print(
+        #     f"len of gradients: {len(self.gradients)}, scores: {[g._score for g in self.gradients]} for {self.name}"
+        # )
+
+        if not self.gradients:
+            return ""
+
+        # sore gradients by the _score from low to high
+        # self.gradients = sorted(
+        #     self.gradients, key=lambda x: x.score if x.score is not None else 1
+        # )
         # print the score for the sorted gradients
         lowest_score_gradients = []
         for i, g in enumerate(self.gradients):
@@ -333,11 +396,39 @@ def get_gradient_and_context_text(self, skip_correct_sample: bool = False) -> st
             lowest_score_gradients.append(g)
             print(f"{i} Score: {g.score} for {g.name}, {type(g.score)}")
 
+        # get component id: gradient
+        component_id_to_gradient: Dict[str, Gradient] = {}
+        for g in lowest_score_gradients:
+            component_id_to_gradient[g.from_response_component_id] = g
+
+        componend_id_to_schema: Dict[str, str] = {}
+        for id, g in component_id_to_gradient.items():
+            componend_id_to_schema[id] = g.context.to_yaml(exclude=["input_output"])
+
+        # parse the gradients and context.
+        gradients_and_context: List[Dict[str, Any]] = (
+            []
+        )  # {gradient: data, context: GradientContext.input_output}
+        for g in lowest_score_gradients:
+            gradients_and_context.append(
+                {
+                    "data_id": g.data_id,
+                    "gradient": g.data,
+                    "context": g.context.input_output,
+                    "score": g.score,
+                }
+            )
+
         gradient_context_combined_str = Prompt(
             template=COMBINED_GRADIENTS_TEMPLATE,
-            prompt_kwargs={"combined_gradients": lowest_score_gradients},
+            prompt_kwargs={
+                "combined_gradients": gradients_and_context,
+                "component_schema": componend_id_to_schema,
+            },
         )().strip()
 
+        # if there are multiple successors, there will be multiple component schemas
+
         return gradient_context_combined_str
 
     def merge_gradients_for_cycle_components(self):
@@ -398,6 +489,7 @@ def set_peers(self, peers: List["Parameter"] = None):
     # Trace the tgd optimizer data
     ############################################################################################################
     def trace_optimizer(self, api_kwargs: Dict[str, Any], response: "TGDData"):
+        r"""Trace the inputs and output of a TGD optimizer."""
         from adalflow.optim.text_grad.tgd_optimizer import TGDOptimizerTrace
 
         self.tgd_optimizer_trace = TGDOptimizerTrace(
@@ -620,6 +712,11 @@ def generate_node_html(node: "Parameter", output_dir="node_pages"):
 
         gradients_json = json.dumps(gradients, indent=4, ensure_ascii=False)
 
+        optimizer_trace = None
+        if node.tgd_optimizer_trace:
+            optimizer_trace = node.tgd_optimizer_trace.to_json_obj()
+            optimizer_trace = json.dumps(optimizer_trace, indent=4, ensure_ascii=False)
+
         with open(filename, "w") as file:
             file.write(
                 f"""
@@ -648,6 +745,7 @@ def generate_node_html(node: "Parameter", output_dir="node_pages"):
                 <p><b>Requires Optimization:</b> {node.requires_opt}</p>
                 <p><b>Type:</b> {node.param_type.value} ({node.param_type.description})</p>
                 <pre><b>Gradients:</b>\n{gradients_json}</pre>
+                <pre><b>TGD Optimizer Trace:</b>\n{optimizer_trace}</pre>
 
             </body>
             </html>

adalflow/adalflow/optim/text_grad/text_loss_with_eval_fn.py

@@ -295,7 +295,7 @@ def _backward_through_one_predecessor(
         gradient_param.add_prompt(gradient_prompt)
         gradient_param.add_context(
             GradientContext(
-                context=conversation_str,
+                input_output=conversation_str,
                 response_desc=response.role_desc,
                 variable_desc=pred.role_desc,
             )

adalflow/adalflow/optim/text_grad/tgd_optimizer.py

@@ -120,28 +120,35 @@ class HistoryPrompt(DataClass):
 # """
 
 OPTIMIZER_SYSTEM_PROMPT = r"""
-You are part of an optimization system designed to refine variables based on feedback generated from a batch of input data.
+You are an optimizer and your task is to improve a variable based on feedback from a batch of input data points.
+
+The variable is either input or output of a functional component where the component schema will be provided.
 
 ### Your Responsibilities:
 1. **Address Feedback**: Resolve concerns raised in the feedback while preserving the positive aspects of the original variable.
-2. **Preserve Intent**: Ensure that the new variable maintains the same overall intent and purpose as the original.
-3. **Leverage Context**: Consider past performance patterns (when available) to retain good qualities in the variable.
-4. **Peer Awareness**:
-   - If a peer variable will be optimized separately, avoid overlapping its scope.
-   - If a peer variable is not being optimized, overlap is permitted when necessary to address the feedback effectively.
+2. **Peer Awareness**:
+   - If a peer will be optimized itself, do not overlap with its scope.
+   - Otherwise, you can overlap if it helps address the feedback effectively.
+3. Observe past performance patterns (when available) to retain good qualities in the variable.
 
 ### Notes:
-- In the absence of specific feedback, you may rephrase the initial variable to improve clarity or specificity without altering its core meaning.
-- When specific feedback is provided, you can either rephrase or refine the variable with more detailed instructions or adjustments to directly or indirectly address the feedback.
+1. You can eliminate unnecessary words or phrases to improve clarity.
+2. Add new elements or rephrase to address the feedback. When no feedback is provided(high batch performance), you rephrase the variable.
+3. Be creative. If adding new elements, be concise.
 
 {{output_format_str}}
 
 {% if instruction_to_optimizer %}
-5. **Additional User Instructions**: {{instruction_to_optimizer}}
+4. **Additional User Instructions**: {{instruction_to_optimizer}}
 {% endif %}
 """
 # 5. **Batch Consistency**: Do not optimize the variable to fit only one specific sample if the batch size is larger than 1. Ensure the variable remains applicable to the entire batch.
 
+# 2. **Preserve Intent**: Ensure that the new variable maintains the same overall intent and purpose as the original.
+
+# - In the absence of specific feedback, you may rephrase the initial variable to improve clarity or specificity without altering its core meaning.
+# - When specific feedback is provided, you can either rephrase or refine the variable with more detailed instructions or adjustments to directly or indirectly address the feedback.
+
 
 @dataclass
 class Instruction(DataClass):
@@ -165,7 +172,7 @@ class TGDData(DataClass):
 
 
 @dataclass
-class TGDOptimizerTrace:
+class TGDOptimizerTrace(DataClass):
     api_kwargs: Dict[str, Any] = field(
         metadata={
             "desc": "The api_kwargs for components like Generator and Retriever that pass to the model client"
@@ -231,7 +238,13 @@ def __init__(
             prompt_kwargs={
                 # "new_variable_start_tag": new_variable_tags[0],
                 # "new_variable_end_tag": new_variable_tags[1],
-                "output_format_str": self.output_parser.get_output_format_str(),
+                "output_format_str": """Your output should be formatted as a standard JSON instance with the following schema:
+```
+{
+    "reasoning": "Why the variable is proposed this way (str) (required)",
+    "proposed_variable": "The proposed variable (str) (required)"
+}
+```"""
             },
         )
         self.variable_and_peers_info = Prompt(
@@ -364,7 +377,8 @@ def _get_user_prompt_kwargs(self, param: Parameter) -> Dict[str, str]:
         )
 
         # variable_grad = param.get_gradients_str()
-        variable_grad = param.get_gradient_and_context_text()
+        # variable_grad = param.get_gradient_and_context_text(skip_correct_sample=False)
+        variable_grad = param.get_gradients_component_schema(skip_correct_sample=False)
 
         user_prompt_kwargs = {
             "variable_and_peers_info": variable_and_peer_info,

adalflow/adalflow/optim/trainer/trainer.py

@@ -1956,7 +1956,11 @@ def _text_grad_constraint_propose_step(
             )
             # check subset validation score
             val_score = val_output.avg_score
-            if val_score > subset_score:
+            if (
+                val_score == subset_score
+                and subset_score >= self.batch_val_score_threshold
+            ) or val_score > subset_score:  # allow perfect subset to pass
+
                 print(f"Pass subset check: {val_score} > {subset_score}")
                 self._track_effectiveness("subset", True)