Fix role nested next node

guidance-ai · May 22, 2023 · 0bc0be1 · 0bc0be1
1 parent 39b3d6a
commit 0bc0be1
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/guidance/_program_executor.py b/guidance/_program_executor.py
@@ -385,6 +385,7 @@ def update_return_value(s):
                         "block_content": self.block_content[-1],
                         "partial_output": self.extend_prefix,
                         "parser_node": node,
+                        "block_close_node": node.children[-1],
                         "next_node": next_node,
                         "next_next_node": next_next_node,
                         "prev_node": node.children[0]

diff --git a/guidance/library/_role.py b/guidance/library/_role.py
@@ -13,9 +13,9 @@ async def role(name, hidden=False, _parser_context=None):
 
     out = await parser.visit(
         block_content[0],
-        next_node=_parser_context["next_node"],
+        next_node=_parser_context["block_close_node"],
         prev_node=_parser_context["prev_node"],
-        next_next_node=_parser_context["next_next_node"]
+        next_next_node=_parser_context["next_node"]
     )
 
     # send the role-end special tokens