Fixing bug 857 (regression from 0.1.14 to 0.1.15)

guidance/library/_capture.py

@@ -1,6 +1,9 @@
 from .._guidance import guidance
 from .._grammar import capture as grammar_capture, GrammarFunction
 
+# Adapted from active_role_end in _model.py, functionality should be shared probably?
+import re
+format_pattern = re.compile(r"<\|\|_.*?_\|\|>", flags=re.DOTALL)
 
 @guidance(stateless=lambda *args, **kwargs: isinstance(args[0], GrammarFunction))
 def capture(lm, value, name):
@@ -9,4 +12,11 @@ def capture(lm, value, name):
     else:
         start_len = len(lm)
         lm += value
-        return lm.set(name, str(lm)[start_len:])
+        # Adapted from active_role_end in _model.py
+        parts = ""
+        for _, role_end_str in lm.opened_blocks.values():
+            role_end_str = format_pattern.sub("", role_end_str)
+            if len(role_end_str) > 0 and not re.fullmatch(r"\s+", role_end_str):
+                parts += role_end_str
+
+        return lm.set(name, str(lm)[start_len-len(parts):].removesuffix(parts))

tests/library/test_capture.py

@@ -1,8 +1,18 @@
-from guidance import capture, models, one_or_more, select, guidance
+from guidance import capture, models, one_or_more, select, guidance, user
+import pytest
 
 from ..utils import get_model
 
 
+@pytest.fixture(scope="module")
+def instruct_model(selected_model, selected_model_name):
+    if selected_model_name in ["transformers_phi3cpu_mini_4k_instruct"]:
+        return selected_model
+    else:
+        pytest.skip("Requires Phi3 4k Instruct model")
+
+
+
 def test_capture():
     model = models.Mock()
     model += "This is" + capture(select(options=["bad", "quite bad"]), name="my_var")
@@ -44,4 +54,11 @@ def raw_fn(lm):
 
     assert str(lm_nocap).endswith("|the end")
     assert str(lm_cap_arg).endswith("|the end")
-    assert str(lm_cap_kwarg).endswith("|the end")
+    assert str(lm_cap_kwarg).endswith("|the end")
+
+def test_capture_within_role(instruct_model: guidance.models.Model):
+    lm = instruct_model
+    test_text = "This is some text in a role."
+    with user():
+        lm += capture(test_text, "test")
+    assert lm["test"] == test_text