fix(dspy): Fix template.extract to be more robust

dsp/templates/template_v2.py

@@ -159,7 +159,19 @@ def extract(
 
         import dspy
 
+        # This should point at the first output field.
+        # So we can start truncating raw_pred from this point.
         idx = min(idx, len(self.fields) - 1)
+
+        # Before that, truncate raw_pred to the output prefix - sometimes
+        # we see the input fields was also included in the completion.
+        # see test_predict.py:test_single_output_with_noise tests
+        next_field_regex = "^" + re.escape(self.fields[idx].name)
+        match = re.search(next_field_regex, raw_pred, flags=re.MULTILINE)
+        if match:
+            end_offset = match.end()
+            raw_pred = raw_pred[end_offset:].strip()
+
         while raw_pred != "" and idx < len(self.fields):
             if idx < len(self.fields) - 1:
                 next_field_name = "\n" + self.fields[idx + 1].name

tests/predict/test_predict.py

@@ -157,6 +157,31 @@ def test_config_management():
     assert "new_key" in config and config["new_key"] == "value"
 
 
+def test_single_output():
+    program = Predict("question -> answer")
+    dspy.settings.configure(lm=DummyLM(["my answer"]))
+    results = program(question="What is 1+1?")
+    assert results.completions.answer[0] == "my answer"
+
+
+def test_single_output_with_prefix():
+    """Text for when prefix for answer is included (e.g Gemini)."""
+    program = Predict("question -> answer")
+    dspy.settings.configure(lm=DummyLM(["Answer: my answer"]))
+    results = program(question="What is 1+1?")
+    assert results.completions.answer[0] == "my answer"
+
+
+def test_single_output_with_noise():
+    """Sometimes we see question was also listed as an output."""
+    program = Predict("question -> answer")
+    dspy.settings.configure(lm=DummyLM([
+        "Question: What is 1+1?\nAnswer: my answer"
+    ]))
+    results = program(question="What is 1+1?")
+    assert results.completions.answer[0] == "my answer"
+
+
 def test_multi_output():
     program = Predict("question -> answer", n=2)
     dspy.settings.configure(lm=DummyLM(["my first answer", "my second answer"]))