microsoft · ultmaster · Jun 30, 2025 · Jun 30, 2025 · Jun 30, 2025
diff --git a/python/poml/prompt.py b/python/poml/prompt.py
@@ -176,8 +176,17 @@ def dump_xml(self) -> str:
         return self._generate_xml_string(pretty=True)
 
     def __enter__(self):
-        """Initializes Prompt for a new XML construction session within a 'with' block."""
-        self.root_elements = []
+        """Enter a context for building a prompt.
+
+        The Prompt instance can be reused across multiple ``with`` blocks. On
+        each entry we simply reset the stack of currently open elements while
+        preserving any previously created root elements so that additional tags
+        can be appended in subsequent sessions.
+        """
+
+        # Reset the stack of open elements for this new session but leave any
+        # existing root elements intact so the prompt can be extended across
+        # multiple ``with`` blocks.
         self.current_parent_stack = []
         return self
 
@@ -191,10 +200,11 @@ def __exit__(self, exc_type, exc_val, exc_tb):
                 "This may indicate nested tag context managers were not properly closed before the Prompt context ended."
             )
 
-        # Consistent with original behavior: clear internal state on exit.
-        # This means results should typically be obtained (via render/dump_xml)
-        # before the Prompt's 'with' block finishes if Prompt itself is a context manager.
-        self.root_elements.clear()
+        # Clear any open elements from the stack.  Previously the entire state
+        # was discarded on exit which meant ``dump_xml`` and ``render`` could
+        # only be called while inside the ``with`` block.  By keeping the root
+        # elements around we allow callers to finalize or extend the prompt
+        # after the block has exited.
         self.current_parent_stack.clear()
 
     def tag(self, tag_name: str, **attrs) -> _ImplicitDualTagHandler:

diff --git a/python/tests/assets/pdf_latex_image.pdf b/python/tests/assets/pdf_latex_image.pdf
diff --git a/python/tests/test_basic.py b/python/tests/test_basic.py
@@ -21,7 +21,6 @@ def test_prompt():
                 p.text("This is a paragraph in the document.")
 
         xml_output = p.dump_xml()
-        print(xml_output)
         assert xml_output == (
             """<Task caption="My Task">This is a task description.<Paragraph>
     <Header>Subheading</Header>This is a paragraph in the document.</Paragraph>
@@ -36,6 +35,51 @@ def test_prompt():
         ]
 
 
+def test_document():
+    with open(Path(__file__).parent / "assets" / "pdf_latex_image.pdf", "rb") as f:
+        pdf_contents = f.read()
+    with Prompt() as p:
+        with p.document(buffer=pdf_contents, parser="pdf"):
+            p.text("This is a PDF document.")
+        xml_output = p.dump_xml()
+        assert "<Document" in xml_output
+        assert "base64=" in xml_output
+
+        result = p.render()
+        assert isinstance(result, list)
+        assert "Lorem ipsum" in result[0]["content"]
+
+
+def test_prompt_reuse_appendable():
+    """Ensure a Prompt instance can be reused across multiple context blocks."""
+
+    prompt = Prompt()
+
+    with prompt:
+        with prompt.role():
+            prompt.text("You are helpful.")
+        with prompt.task(caption="First"):
+            prompt.text("Describe A.")
+
+    # Results should be available outside the context
+    first_xml = prompt.dump_xml()
+    assert first_xml.count("<Task") == 1
+    first_render = prompt.render()
+    assert "Describe A." in first_render[0]["content"]
+
+    # Re-enter the context and append more content
+    with prompt:
+        with prompt.task(caption="Second"):
+            prompt.text("Describe B.")
+
+    second_xml = prompt.dump_xml()
+    # Both tasks should be present
+    assert second_xml.count("<Task") == 2
+    second_render = prompt.render()
+    assert "Describe A." in second_render[0]["content"]
+    assert "Describe B." in second_render[0]["content"]
+
+
 def test_trace():
     clear_trace()
     set_trace(True)