fix(UI): improve dropdowns and update UI

app.py

@@ -4,6 +4,8 @@ from utils.data_generator import generate_examples
 from llm_client import ask_llm
 from utils.evaluation import compute_similarity, validate_answer
 import asyncio
+import os
+
 
 print("Scanning for all available Engchain templates...")
 ALL_TEMPLATES = discover_templates()
@@ -109,29 +111,31 @@ with gr.Blocks(theme=gr.themes.Soft(), css="""
         background: var(--panel-background-fill);
     }
 """) as demo:
+
     gr.Markdown(
         """
         # 📈 Engchain Template Playground  
-        Select a domain, a file, and a template to generate 10 unique Q&A examples.  
+        Select a branch, domain, file, and template to generate 10 unique Q&A examples.  
         Compare your solutions against a live LLM and see similarity + correctness checks.
         """
     )
 
     with gr.Group(elem_classes="dropdown-box"):
         with gr.Row():
-            domain_dropdown = gr.Dropdown(
-                label="1. Select Engineering Domain",
-                choices=[(format_name(d), d) for d in sorted(list(ALL_TEMPLATES.keys()))],
+            branch_dropdown = gr.Dropdown(
+                label="1. Select Branch",
+                choices=[(format_name(b), b) for b in sorted(list(ALL_TEMPLATES.keys()))],
             )
-            file_dropdown = gr.Dropdown(label="2. Select Specific Area / File", interactive=False)
-            template_dropdown = gr.Dropdown(label="3. Select Specific Tool / Template", interactive=False)
+            domain_dropdown = gr.Dropdown(label="2. Select Engineering Domain", interactive=False)
+            file_dropdown = gr.Dropdown(label="3. Select Specific Area / File", interactive=False)
+            template_dropdown = gr.Dropdown(label="4. Select Specific Tool / Template", interactive=False)
 
     generate_button = gr.Button("Generate 10 Examples", variant="primary")
 
     # Pre-create a slot for template heading
     template_heading = gr.Markdown("", visible=False, elem_classes="template-heading")
 
-    # Pre-create 10 slots
+    # Pre-create 10 example slots
     example_blocks = []
     for i in range(10):
         with gr.Group(visible=False, elem_classes="example-card") as grp:
@@ -145,57 +149,70 @@ with gr.Blocks(theme=gr.themes.Soft(), css="""
 
             example_blocks.append((grp, q_md, s_md, btn, llm_group, eval_md, llm_response_md))
 
-            # Bind button → LLM response + evaluation
             async def reveal_llm(q, s, llm_group=llm_group, eval_md=eval_md, llm_response_md=llm_response_md):
-                # temporary loading message
-                await asyncio.sleep(0)  # ensures Gradio sees async
+                await asyncio.sleep(0)
                 yield gr.update(visible=True), gr.update(value=""), gr.update(value="⏳ Fetching LLM response...")
-
-                # call the LLM asynchronously
-                response = await asyncio.to_thread(call_llm, q)  # run blocking LLM call in thread
-                
+                response = await asyncio.to_thread(call_llm, q)
                 sim = compute_similarity(s, response)
                 match_flag = validate_answer(s, response)
-
                 eval_html = (
                     f"<div class='badge badge-blue'><span>Textual Similarity:</span> {sim}%</div>"
                     f"<div class='badge {'badge-green' if match_flag else 'badge-red'}'>"
                     f"<span>Final Answer Validation:</span> {'✅ Yes' if match_flag else '❌ No'}</div>"
                 )
-
                 yield gr.update(visible=True), gr.update(value=eval_html), gr.update(value=response)
 
             btn.click(reveal_llm, inputs=[q_md, s_md], outputs=[llm_group, eval_md, llm_response_md])
 
     # Dropdown update logic
-    def update_file_dropdown(domain):
-        if not domain:
+    def update_domain_dropdown(branch):
+        if not branch:
+            return gr.update(choices=[], value=None, interactive=False), \
+                   gr.update(choices=[], value=None, interactive=False), \
+                   gr.update(choices=[], value=None, interactive=False)
+        domains = sorted(list(ALL_TEMPLATES.get(branch, {}).keys()))
+        formatted_domains = [(format_name(d), d) for d in domains]
+        return gr.update(choices=formatted_domains, value=None, interactive=True), \
+               gr.update(choices=[], value=None, interactive=False), \
+               gr.update(choices=[], value=None, interactive=False)
+
+        
+    def update_file_dropdown(branch, domain):
+        if not branch or not domain:
             return (
                 gr.update(choices=[], value=None, interactive=False),
                 gr.update(choices=[], value=None, interactive=False),
             )
-        files = sorted(list(ALL_TEMPLATES.get(domain, {}).keys()))
-        formatted_files = [(format_name(f), f) for f in files]
+        # get all files for this branch+domain
+        files_dict = ALL_TEMPLATES.get(branch, {}).get(domain, {})
+        files = sorted(list(files_dict.keys()))
+
+        # Use only the last part of the path as label
+        formatted_files = [(format_name(os.path.basename(f)), f) for f in files]
+
         return (
             gr.update(choices=formatted_files, value=None, interactive=True),
             gr.update(choices=[], value=None, interactive=False),
         )
 
-    def update_template_dropdown(domain, filename):
-        if not domain or not filename:
+
+    def update_template_dropdown(branch, domain, filename):
+        if not branch or not domain or not filename:
             return gr.update(choices=[], value=None, interactive=False)
-        templates = sorted(ALL_TEMPLATES.get(domain, {}).get(filename, []))
+        templates = sorted(ALL_TEMPLATES.get(branch, {}).get(domain, {}).get(filename, []))
         formatted_templates = [(format_name(t), t) for t in templates]
         return gr.update(choices=formatted_templates, value=None, interactive=True)
 
-    # Populate the pre-created slots
-    def render_examples(domain, filename, template_name):
-        # If any selection is missing, clear screen and show friendly error
-        if not domain or not filename or not template_name:
-            heading_text = "<div style='margin-top:24px; font-size: 20px; color: red;'>⚠️ Please select a domain, file, and template/tool before generating examples.</div>"
+    branch_dropdown.change(update_domain_dropdown, branch_dropdown, [domain_dropdown, file_dropdown, template_dropdown])
+    domain_dropdown.change(update_file_dropdown, [branch_dropdown, domain_dropdown], [file_dropdown, template_dropdown])
+    file_dropdown.change(update_template_dropdown, [branch_dropdown, domain_dropdown, file_dropdown], template_dropdown)
+
+    # Populate example slots
+    def render_examples(branch, domain, filename, template_name):
+        if not branch or not domain or not filename or not template_name:
+            heading_text = "<div style='margin-top:24px; font-size: 20px; color: red;'>⚠️ Please select a branch, domain, file, and template/tool before generating examples.</div>"
             template_heading_update = gr.update(value=heading_text, visible=True)
 
-            # Clear all example slots
             updates = []
             for grp, q_md, s_md, btn, llm_group, eval_md, llm_response_md in example_blocks:
                 updates.extend([
@@ -209,9 +226,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css="""
                 ])
             return [template_heading_update] + updates
 
-        # Normal behavior: generate examples
-        examples = generate_examples(domain, filename, template_name)
-
+        examples = generate_examples(branch, domain, filename, template_name)
         heading_text = f"<div style='margin-top:24px; font-size: 24px;'>Generated Examples for Template: <b>{format_name(template_name)}</b></div>"
         template_heading_update = gr.update(value=heading_text, visible=True)
 
@@ -240,12 +255,9 @@ with gr.Blocks(theme=gr.themes.Soft(), css="""
                 ])
         return [template_heading_update] + updates
 
-
-    domain_dropdown.change(update_file_dropdown, domain_dropdown, [file_dropdown, template_dropdown])
-    file_dropdown.change(update_template_dropdown, [domain_dropdown, file_dropdown], template_dropdown)
     generate_button.click(
         render_examples,
-        [domain_dropdown, file_dropdown, template_dropdown],
+        [branch_dropdown, domain_dropdown, file_dropdown, template_dropdown],
         [template_heading] + sum(([grp, q_md, s_md, btn, llm_group, eval_md, llm_response_md]
              for (grp, q_md, s_md, btn, llm_group, eval_md, llm_response_md) in example_blocks), [])
     )
@@ -254,3 +266,4 @@ with gr.Blocks(theme=gr.themes.Soft(), css="""
 if __name__ == "__main__":
     print("Launching Gradio app...")
     demo.launch()
+

utils/data_generator.py

@@ -61,21 +61,12 @@ def discover_templates(root_dir="data/templates/branches"):
 
     return discovered
 
-
-def generate_examples(domain, filename, template_name):
+def generate_examples(branch, domain, filename, template_name):
     """
     Dynamically imports and runs a selected template function multiple times
     to generate a list of {question, solution} objects.
-
-    Args:
-        domain (str): The name of the domain (e.g. 'chemical_engineering').
-        filename (str): Relative path under branches/ (e.g. 'chemical_engineering/reaction_kinetics/mole_balances.py').
-        template_name (str): The template function to call (e.g. 'template_mole_balance').
-
-    Returns:
-        list[dict]: [{"question": ..., "solution": ...}, ...]
     """
-    if not all([domain, filename, template_name]):
+    if not all([branch, domain, filename, template_name]):
         print("Missing arguments to generate_examples()")
         return []
 
@@ -84,7 +75,6 @@ def generate_examples(domain, filename, template_name):
         filename_no_ext = filename.replace(".py", "").replace(os.sep, ".")
         module_path = f"data.templates.branches.{filename_no_ext}"
 
-        # Import module dynamically
         module = importlib.import_module(module_path)
         template_function = getattr(module, template_name)
 
@@ -96,48 +86,8 @@ def generate_examples(domain, filename, template_name):
                 examples.append({"question": question, "solution": solution})
             else:
                 print(f"Warning: {template_name} did not return (question, solution) tuple.")
-
         return examples
 
-    except ImportError as e:
-        print(f"ImportError in {domain}/{filename}: {e}")
-        return []
-    except AttributeError as e:
-        print(f"AttributeError in {domain}/{filename}: {e}")
-        return []
     except Exception as e:
-        print(f"Error running template '{template_name}' from {filename}: {e}")
+        print(f"Error running template '{template_name}' from {branch}/{domain}/{filename}: {e}")
         return []
-
-
-if __name__ == "__main__":
-    """
-    Run this file directly to test the discovery and generation system.
-    """
-    print("Discovering templates...\n")
-    templates = discover_templates()
-
-    if not templates:
-        print("No templates found.")
-    else:
-        for domain, files in templates.items():
-            print(f"\nDomain: {domain}")
-            for file_path, funcs in files.items():
-                print(f"  {file_path}")
-                for func in funcs:
-                    print(f"     {func}")
-
-        # Example test: run one of the templates
-        print("\n\n Testing template generation...")
-        # pick first available template
-        first_domain = next(iter(templates))
-        first_file, first_funcs = next(iter(templates[first_domain].items()))
-        first_func = first_funcs[0]
-
-        examples = generate_examples(first_domain, first_file, first_func)
-        print(f"\n Generated {len(examples)} examples from {first_func} in {first_file}:")
-        for ex in examples[:3]:
-            print("Q:", ex["question"])
-            print("A:", ex["solution"])
-            print("---")
-

utils/template_discovery.py

@@ -4,7 +4,15 @@ import ast
 def discover_templates(root_dir="data/templates/branches"):
     """
     Recursively scans directory structure to find all template functions
-    under each domain (supports nested folders like branches/chemical_engineering/...).
+    under each branch/domain/file.
+    Returns:
+        dict: {
+            branch: {
+                domain: {
+                    file_path: [template_function_names]
+                }
+            }
+        }
     """
     discovered = {}
 
@@ -12,35 +20,40 @@ def discover_templates(root_dir="data/templates/branches"):
         print(f"Error: Directory '{root_dir}' not found.")
         return discovered
 
-    # Walk one level down (e.g., chemical_engineering, mechanical_engineering, etc.)
-    for domain in sorted(os.listdir(root_dir)):
-        domain_path = os.path.join(root_dir, domain)
-        if os.path.isdir(domain_path):
-            domain_templates = {}
-
-            # Recursively walk through all .py files under this domain
-            for dirpath, _, filenames in os.walk(domain_path):
-                for filename in sorted(filenames):
-                    if filename.endswith(".py"):
-                        file_path = os.path.join(dirpath, filename)
-                        try:
-                            with open(file_path, "r", encoding="utf-8") as f:
-                                file_content = f.read()
-
-                            tree = ast.parse(file_content)
-                            template_functions = [
-                                node.name for node in ast.walk(tree)
-                                if isinstance(node, ast.FunctionDef) and node.name.startswith("template_")
-                            ]
-
-                            if template_functions:
-                                rel_path = os.path.relpath(file_path, root_dir)
-                                domain_templates[rel_path] = sorted(template_functions)
-                        except Exception as e:
-                            print(f"Error parsing {file_path}: {e}")
-
-            if domain_templates:
-                discovered[domain] = domain_templates
+    # Top-level folders = branches
+    for branch in sorted(os.listdir(root_dir)):
+        branch_path = os.path.join(root_dir, branch)
+        if os.path.isdir(branch_path):
+            branch_dict = {}
+
+            # Domains under branch
+            for domain in sorted(os.listdir(branch_path)):
+                domain_path = os.path.join(branch_path, domain)
+                if os.path.isdir(domain_path):
+                    domain_templates = {}
+
+                    # Walk all Python files under domain
+                    for dirpath, _, filenames in os.walk(domain_path):
+                        for filename in sorted(filenames):
+                            if filename.endswith(".py"):
+                                file_path = os.path.join(dirpath, filename)
+                                try:
+                                    with open(file_path, "r", encoding="utf-8") as f:
+                                        tree = ast.parse(f.read())
+                                    template_functions = [
+                                        node.name for node in ast.walk(tree)
+                                        if isinstance(node, ast.FunctionDef) and node.name.startswith("template_")
+                                    ]
+                                    if template_functions:
+                                        rel_path = os.path.relpath(file_path, root_dir)
+                                        domain_templates[rel_path] = sorted(template_functions)
+                                except Exception as e:
+                                    print(f"Error parsing {file_path}: {e}")
+
+                    if domain_templates:
+                        branch_dict[domain] = domain_templates
+
+            if branch_dict:
+                discovered[branch] = branch_dict
 
     return discovered
-