fix: review comments

Author: srini047

haystack/components/builders/prompt_builder.py

@@ -8,7 +8,7 @@
 from jinja2.sandbox import SandboxedEnvironment
 
 from haystack import component, default_to_dict, logging
-from haystack.utils import Jinja2TimeExtension
+from haystack.utils import Jinja2TimeExtension, extract_declared_variables
 
 logger = logging.getLogger(__name__)
 
@@ -174,17 +174,9 @@ def __init__(
             self._env = SandboxedEnvironment()
 
         self.template = self._env.from_string(template)
-        if not variables:
-            # infer variables from template
-            ast = self._env.parse(template)
-            template_variables = meta.find_undeclared_variables(ast)
 
-            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 = extract_declared_variables(template, env=self._env)
 
-            variables = list(template_variables - assigned_variables)
         variables = variables or []
         self.variables = variables
 

haystack/utils/__init__.py

@@ -15,6 +15,7 @@
     "device": ["ComponentDevice", "Device", "DeviceMap", "DeviceType"],
     "deserialization": ["deserialize_document_store_in_init_params_inplace", "deserialize_chatgenerator_inplace"],
     "filters": ["document_matches_filter", "raise_on_invalid_filter_syntax"],
+    "jinja2": ["extract_declared_variables"],
     "jinja2_extensions": ["Jinja2TimeExtension"],
     "jupyter": ["is_in_jupyter"],
     "misc": ["expit", "expand_page_range"],
@@ -40,6 +41,7 @@
     from .device import DeviceType as DeviceType
     from .filters import document_matches_filter as document_matches_filter
     from .filters import raise_on_invalid_filter_syntax as raise_on_invalid_filter_syntax
+    from .jinja2 import extract_declared_variables as extract_declared_variables
     from .jinja2_extensions import Jinja2TimeExtension as Jinja2TimeExtension
     from .jupyter import is_in_jupyter as is_in_jupyter
     from .misc import expand_page_range as expand_page_range

haystack/utils/jinja2.py

@@ -0,0 +1,37 @@
+from jinja2 import Environment, meta, nodes
+
+
+def extract_declared_variables(template_str: str, env: Environment | None = None) -> list[str]:
+    """
+    Extract declared variables from a Jinja2 template string.
+
+    Args:
+        template_str (str): The Jinja2 template string to analyze.
+        env (Environment, optional): The Jinja2 Environment. Defaults to None.
+
+    Returns:
+        A list of variable names used in the template.
+    """
+    env = env or Environment()
+
+    try:
+        ast = env.parse(template_str)
+    except Exception as e:
+        raise RuntimeError(f"Failed to parse Jinja2 template: {e}")
+
+    # Find undeclared variables
+    template_variables = meta.find_undeclared_variables(ast)
+
+    # Collect all variables assigned inside the template via {% set %}
+    assigned_variables = set()
+
+    for node in ast.find_all(nodes.Assign):
+        if isinstance(node.target, nodes.Name):
+            assigned_variables.add(node.target.name)
+        elif isinstance(node.target, (nodes.List, nodes.Tuple)):
+            for name_node in node.target.items:
+                if isinstance(name_node, nodes.Name):
+                    assigned_variables.add(name_node.name)
+
+    variables = list(template_variables - assigned_variables)
+    return variables

test/components/builders/test_prompt_builder.py

@@ -352,10 +352,32 @@ def test_template_assigned_variables_from_required_inputs(self) -> None:
                    """
 
         builder = PromptBuilder(template=template, required_variables="*")
-        docs = [Document(content="Doc 1"), Document(content="Doc 2")]
 
-        res = builder.run(docs=docs, existing_documents=None)
+        builder = PromptBuilder(template=template, required_variables="*")
+        assert set(builder.variables) == {"docs", "existing_documents"}
+
+    def test_variables_correct_with_tuple_assignment(self):
+        template = """{% if existing_documents is not none %}
+{% set x, y = (existing_documents|length, 1) %}
+{% else %}
+{% set x, y = (0, 1) %}
+{% endif %}
+x={{ x }}, y={{ y }}
+"""
+        builder = PromptBuilder(template=template, required_variables="*")
+        assert set(builder.variables) == {"existing_documents"}
+        res = builder.run(existing_documents=None)
+        assert "x=0, y=1" in res["prompt"]
 
-        assert "<document reference=" in res["prompt"]
-        assert "Doc 1" in res["prompt"]
-        assert "Doc 2" in res["prompt"]
+    def test_variables_correct_with_list_assignment(self):
+        template = """{% if existing_documents is not none %}
+{% set x, y = [existing_documents|length, 1] %}
+{% else %}
+{% set x, y = [0, 1] %}
+{% endif %}
+x={{ x }}, y={{ y }}
+"""
+        builder = PromptBuilder(template=template, required_variables="*")
+        assert set(builder.variables) == {"existing_documents"}
+        res = builder.run(existing_documents=None)
+        assert "x=0, y=1" in res["prompt"]