fix: prompt consumes same iterator twice (#768)

Author: ds-sebastianchwilczynski

packages/ragbits-core/CHANGELOG.md

@@ -2,6 +2,9 @@
 
 ## Unreleased
 
+- Add tool_choice parameter to LLM interface (#738)
+- Fix Prompt consumes same iterator twice leading to no data added to chat (#768)
+
 ## 1.2.1 (2025-08-04)
 
 ## 1.2.0 (2025-08-01)

packages/ragbits-core/src/ragbits/core/prompt/prompt.py

@@ -188,7 +188,6 @@ def __init__(self, input_data: PromptInputT | None = None, history: ChatFormat |
         self.rendered_system_prompt = (
             self._render_template(self.system_prompt_template, input_data) if self.system_prompt_template else None
         )
-        self.rendered_user_prompt = self._render_template(self.user_prompt_template, input_data)
         self.attachments = self._get_attachments_from_input_data(input_data)
 
         # Additional few shot examples that can be added dynamically using methods
@@ -197,7 +196,9 @@ def __init__(self, input_data: PromptInputT | None = None, history: ChatFormat |
 
         # Additional conversation history that can be added dynamically using methods
         self._conversation_history: list[dict[str, Any]] = history or []
-        self.add_user_message(input_data if input_data else self.rendered_user_prompt)
+
+        self.add_user_message(input_data or self._render_template(self.user_prompt_template, input_data))
+        self.rendered_user_prompt = self.chat[-1]["content"]
         super().__init__()
 
     @property

packages/ragbits-core/tests/unit/prompts/test_prompt.py

@@ -1,3 +1,4 @@
+from collections.abc import Iterator
 from pathlib import Path
 from typing import Any
 
@@ -19,6 +20,15 @@ class _PromptInput(pydantic.BaseModel):
     age: int
 
 
+class _PromptInputWithIterator(pydantic.BaseModel):
+    """
+    Input format for the TestPromptWithIterator, which can be consumed only once.
+    """
+
+    model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
+    context: Iterator[str]
+
+
 class _SingleAttachmentPromptInput(pydantic.BaseModel):
     """
     Single input format for the TestAttachmentPrompt.
@@ -204,6 +214,9 @@ class TestAttachmentPrompt(Prompt):
     assert chat[0]["role"] == "user"
     assert chat[0]["content"][0]["text"] == "What is in this image?"
     assert chat[0]["content"][1]["type"] == "image_url"
+    assert len(prompt.rendered_user_prompt) == 2
+    assert prompt.rendered_user_prompt[0] == {"type": "text", "text": "What is in this image?"}
+    assert prompt.rendered_user_prompt[1]["type"] == "image_url"
 
 
 def test_image_prompt_encoding():
@@ -270,6 +283,9 @@ class TestAttachmentPrompt(Prompt):
     assert chat[0]["role"] == "user"
     assert chat[0]["content"][0]["text"] == "What is in this PDF?"
     assert chat[0]["content"][1]["type"] == "file"
+    assert len(prompt.rendered_user_prompt) == 2
+    assert prompt.rendered_user_prompt[0] == {"type": "text", "text": "What is in this PDF?"}
+    assert prompt.rendered_user_prompt[1]["type"] == "file"
 
 
 def test_pdf_prompt_encoding():
@@ -337,6 +353,21 @@ class TestPrompt(Prompt[_PromptInput, str]):  # pylint: disable=unused-variable
     ]
 
 
+def test_prompt_with_iterator_input():
+    """Test that a prompt can be created with an iterator input."""
+
+    class TestPrompt(Prompt[_PromptInputWithIterator, str]):  # pylint: disable=unused-variable
+        """A test prompt"""
+
+        user_prompt = "Context: {% for chunk in context %}{{ chunk }}{%- endfor %}"
+
+    prompt = TestPrompt(_PromptInputWithIterator(context=iter(["Hello", " World"])))
+    assert prompt.rendered_user_prompt == "Context: Hello World"
+    assert prompt.chat == [
+        {"role": "user", "content": "Context: Hello World"},
+    ]
+
+
 def test_input_type_must_be_pydantic_model():
     """Test that an error is raised when the input type is not a Pydantic model."""
     with pytest.raises(AssertionError):