diff --git a/opto/optimizers/optoprime_v2.py b/opto/optimizers/optoprime_v2.py index a512af8e..0d6f3267 100644 --- a/opto/optimizers/optoprime_v2.py +++ b/opto/optimizers/optoprime_v2.py @@ -492,7 +492,7 @@ def repr_node_value(self, node_dict, node_tag="node", temp_list.append( f"<{node_tag} name=\"{k}\" type=\"{type(v[0]).__name__}\">\n<{value_tag}>\n{v[0]}\n\n\n") else: - constraint_expr = f"\n{v[1]}\n" + constraint_expr = f"<{constraint_tag}>\n{v[1]}\n" # Respect custom constraint tags for code variables. signature = v[1].replace("The code should start with:\n", "") func_body = v[0].replace(signature, "") temp_list.append( @@ -572,7 +572,7 @@ def problem_instance(self, summary, mask=None): instruction=self.objective if "#Instruction" not in mask else "", code=( "\n".join([v for k, v in sorted(summary.graph)]) - if self.optimizer_prompt_symbol_set.inputs_section_title not in mask + if self.optimizer_prompt_symbol_set.code_section_title not in mask # Use the code section mask for the code section. else "" ), documentation=( diff --git a/tests/llm_optimizers_tests/test_optoprime_v2.py b/tests/llm_optimizers_tests/test_optoprime_v2.py index b1032f28..49e3bb98 100644 --- a/tests/llm_optimizers_tests/test_optoprime_v2.py +++ b/tests/llm_optimizers_tests/test_optoprime_v2.py @@ -10,7 +10,7 @@ from opto import trace from opto.trace import node, bundle -from opto.optimizers.optoprime_v2 import OptoPrimeV2, OptimizerPromptSymbolSet2 +from opto.optimizers.optoprime_v2 import OptoPrimeV2, OptimizerPromptSymbolSet, OptimizerPromptSymbolSet2 # You can override for temporarly testing a specific optimizer ALL_OPTIMIZERS = [TextGrad] # [OptoPrimeMulti] ALL_OPTIMIZERS = [OptoPrime] @@ -34,6 +34,53 @@ def test_response_extraction(): pass +class CustomConstraintSymbolSet(OptimizerPromptSymbolSet): + constraint_tag = "guard" + + +def test_code_section_mask_only_hides_code_section(): + num = node(1, trainable=True) + result = num + 1 + optimizer = OptoPrimeV2([num], use_json_object_format=False) + + optimizer.zero_feedback() + optimizer.backward(result, "make this number bigger") + + summary = optimizer.summarize() + + code_masked = optimizer.problem_instance( + summary, + mask=[optimizer.optimizer_prompt_symbol_set.code_section_title], + ) + inputs_masked = optimizer.problem_instance( + summary, + mask=[optimizer.optimizer_prompt_symbol_set.inputs_section_title], + ) + + assert code_masked.code == "", "# Code must hide the code section." + assert inputs_masked.code != "", "# Inputs must not hide the code section." + assert inputs_masked.inputs == "", "# Inputs must still hide the inputs section." + + +def test_repr_node_value_respects_custom_constraint_tag_for_code_variables(): + num = node(1, trainable=True) + optimizer = OptoPrimeV2( + [num], + use_json_object_format=False, + optimizer_prompt_symbol_set=CustomConstraintSymbolSet(), + ) + + rendered = optimizer.repr_node_value( + {"__code0": ("def f(x):\n return x", "The code should start with:\ndef f(x):")}, + node_tag=optimizer.optimizer_prompt_symbol_set.variable_tag, + value_tag=optimizer.optimizer_prompt_symbol_set.value_tag, + constraint_tag=optimizer.optimizer_prompt_symbol_set.constraint_tag, + ) + + assert "" in rendered, "Full code rendering must use custom constraint tags." + assert "" not in rendered, "Full code rendering must not hardcode the default tag." + + def test_tag_template_change(): num_1 = node(1, trainable=True) num_2 = node(2, trainable=True, description="<=5")