Fix OptoPrimeV2 prompt masking for code, and set proper constraint tags

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</{value_tag}>\n</{node_tag}>\n")
             else:
-                constraint_expr = f"<constraint>\n{v[1]}\n</constraint>"
+                constraint_expr = f"<{constraint_tag}>\n{v[1]}\n</{constraint_tag}>"  # 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=(

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 "<guard>" in rendered, "Full code rendering must use custom constraint tags."
+    assert "<constraint>" 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")