fix: prompt-builder - jinja2 template set vars still shows required

Author: srini047

haystack/components/builders/prompt_builder.py

@@ -4,7 +4,7 @@
 
 from typing import Any, Literal, Optional, Union
 
-from jinja2 import meta
+from jinja2 import meta, nodes
 from jinja2.sandbox import SandboxedEnvironment
 
 from haystack import component, default_to_dict, logging
@@ -178,7 +178,13 @@ def __init__(
             # infer variables from template
             ast = self._env.parse(template)
             template_variables = meta.find_undeclared_variables(ast)
-            variables = list(template_variables)
+
+            assigned_variables = set()
+            for node in ast.find_all((nodes.Assign, nodes.For)):
+                if hasattr(node.target, "name"):
+                    assigned_variables.add(node.target.name)
+
+            variables = list(template_variables - assigned_variables)
         variables = variables or []
         self.variables = variables
 

test/components/builders/test_prompt_builder.py

@@ -337,3 +337,25 @@ def test_warning_no_required_variables(self, caplog):
         with caplog.at_level(logging.WARNING):
             _ = PromptBuilder(template="This is a {{ variable }}")
             assert "but `required_variables` is not set." in caplog.text
+
+    def test_template_assigned_variables_from_required_inputs(self) -> None:
+        template = """{% if existing_documents is not none %}
+                   {% set existing_doc_len = existing_documents|length %}
+                   {% else %}
+                   {% set existing_doc_len = 0 %}
+                   {% endif %}
+                   {% for doc in docs %}
+                   <document reference="{{loop.index + existing_doc_len}}">
+                    {{ doc.content }}
+                    </document>
+                   {% endfor %}
+                   """
+
+        builder = PromptBuilder(template=template, required_variables="*")
+        docs = [Document(content="Doc 1"), Document(content="Doc 2")]
+
+        res = builder.run(docs=docs, existing_documents=None)
+
+        assert "<document reference=" in res["prompt"]
+        assert "Doc 1" in res["prompt"]
+        assert "Doc 2" in res["prompt"]