Update app.py

app.py

@@ -1117,9 +1117,7 @@ with st.sidebar:
 st.title("Template Generator")
 
 # Create tabs for workflow
-tab1, tab2, tab3, tab4 = st.tabs(
-    ["Setup", "Edit Template", "Use Template", "Generate Data"]
-)
+tab1, tab2, tab3 = st.tabs(["Setup", "Edit and Use Template", "Generate Data"])
 
 with tab1:
     st.header("Project Setup")
@@ -1269,13 +1267,23 @@ with tab2:
     if st.session_state.show_template_editor and st.session_state.template_spec:
         st.header("Template Editor")
 
-        # Initialize suggested variables in session state if not present
+        # Initialize session state variables
         if "suggested_variables" not in st.session_state:
             st.session_state.suggested_variables = []
-
-        # Initialize a tracking variable for added suggestions
         if "added_suggestions" not in st.session_state:
             st.session_state.added_suggestions = set()
+        if (
+            "last_template" not in st.session_state
+            or st.session_state.last_template != st.session_state.template_spec
+        ):
+            st.session_state.user_inputs = {}
+            st.session_state.last_template = st.session_state.template_spec
+        if "show_variable_editor" not in st.session_state:
+            st.session_state.show_variable_editor = None
+        if "show_output_editor" not in st.session_state:
+            st.session_state.show_output_editor = None
+        if "show_suggested_vars" not in st.session_state:
+            st.session_state.show_suggested_vars = False
 
         # Basic template information
         with st.expander("Template Information", expanded=True):
@@ -1295,13 +1303,44 @@ with tab2:
                 height=100,
             )
 
-        # Knowledge Base Analysis
-        with st.expander("Knowledge Base Analysis", expanded=True):
-            if st.session_state.knowledge_base:
+        # Prompt Template Section
+        with st.expander("Prompt Template", expanded=True):
+            st.info("Use {variable_name} to refer to input variables in your template")
+
+            # Add buttons for prompt management
+            col1, col2 = st.columns([1, 1])
+            with col1:
+                rewrite_prompt = st.button("AI Rewrite Prompt")
+            with col2:
+                reroll_prompt = st.button("Reroll Prompt Variation")
+
+            # Handle prompt rewriting
+            if rewrite_prompt or reroll_prompt:
+                with st.spinner("Generating improved prompt template..."):
+                    improved_template = generate_improved_prompt_template(
+                        st.session_state.template_spec, st.session_state.knowledge_base
+                    )
+                    # Only update if we got a valid result back
+                    if improved_template and len(improved_template) > 10:
+                        st.session_state.template_spec["prompt"] = improved_template
+                        st.success("Prompt template updated!")
+
+            # Display the prompt template
+            prompt_template = st.text_area(
+                "Edit the prompt template",
+                value=st.session_state.template_spec["prompt"],
+                height=200,
+            )
+            st.session_state.template_spec["prompt"] = prompt_template
+
+        # Knowledge Base Analysis Section - Moved outside of Input Variables expander
+        if st.session_state.knowledge_base:
+            with st.expander("Knowledge Base Analysis", expanded=False):
                 st.info("Analyze the knowledge base to suggest variables and values")
 
                 if st.button(
-                    "Analyze Knowledge Base for Variables", key="analyze_kb_button"
+                    "Analyze Knowledge Base for Variables",
+                    key="analyze_kb_button_input",
                 ):
                     client = get_openai_client()
                     if not client:
@@ -1315,6 +1354,7 @@ with tab2:
                             )
                             if suggested_vars:
                                 st.session_state.suggested_variables = suggested_vars
+                                st.session_state.show_suggested_vars = True
                                 st.success(
                                     f"Found {len(suggested_vars)} potential variables in the knowledge base"
                                 )
@@ -1324,10 +1364,12 @@ with tab2:
                                 )
 
                 # Display suggested variables if they exist
-                if st.session_state.suggested_variables:
+                if (
+                    st.session_state.suggested_variables
+                    and st.session_state.show_suggested_vars
+                ):
                     st.subheader("Suggested Variables")
 
-                    # Create a container for the variables
                     for i, var in enumerate(st.session_state.suggested_variables):
                         # Generate a unique ID for this variable
                         var_id = f"{var['name']}_{i}"
@@ -1336,81 +1378,39 @@ with tab2:
                         if var_id in st.session_state.added_suggestions:
                             continue
 
-                        with st.container():
-                            col1, col2 = st.columns([3, 1])
-                            with col1:
-                                st.markdown(
-                                    f"**{var['name']}** ({var['type']}): {var['description']}"
-                                )
+                        col1, col2 = st.columns([4, 1])
+                        with col1:
+                            st.markdown(
+                                f"**{var['name']}** ({var['type']}): {var['description']}"
+                            )
+                            if var.get("options"):
+                                st.markdown(f"Options: {', '.join(var['options'])}")
+                        with col2:
+                            if st.button("Add", key=f"add_suggested_{var_id}"):
+                                # Add this variable to the template
+                                new_var = {
+                                    "name": var["name"],
+                                    "description": var["description"],
+                                    "type": var["type"],
+                                }
                                 if var.get("options"):
-                                    st.markdown(f"Options: {', '.join(var['options'])}")
-                            with col2:
-                                # Use a unique key for each button
-                                if st.button("Add", key=f"add_suggested_{var_id}"):
-                                    # Add this variable to the template
-                                    new_var = {
-                                        "name": var["name"],
-                                        "description": var["description"],
-                                        "type": var["type"],
-                                    }
-                                    if var.get("options"):
-                                        new_var["options"] = var["options"]
-                                    if var["type"] in ["string", "int", "float"]:
-                                        new_var["min"] = 1
-                                        new_var["max"] = 100
-
-                                    # Add to input variables
-                                    st.session_state.template_spec["input"].append(
-                                        new_var
-                                    )
-
-                                    # Mark this variable as added
-                                    st.session_state.added_suggestions.add(var_id)
+                                    new_var["options"] = var["options"]
+                                if var["type"] in ["string", "int", "float"]:
+                                    new_var["min"] = 1
+                                    new_var["max"] = 100
 
-                                    # Show success message
-                                    st.success(
-                                        f"Added {var['name']} to input variables!"
-                                    )
-            else:
-                st.warning(
-                    "No knowledge base available. Please upload documents in the Setup tab first."
-                )
+                                # Add to input variables
+                                st.session_state.template_spec["input"].append(new_var)
 
-        # Prompt Template Section
-        with st.expander("Prompt Template", expanded=True):
-            st.info("Use {variable_name} to refer to input variables in your template")
-
-            # Add buttons for prompt management
-            col1, col2 = st.columns([1, 1])
-            with col1:
-                rewrite_prompt = st.button("AI Rewrite Prompt")
-            with col2:
-                reroll_prompt = st.button("Reroll Prompt Variation")
-
-            # Handle prompt rewriting
-            if rewrite_prompt or reroll_prompt:
-                with st.spinner("Generating improved prompt template..."):
-                    improved_template = generate_improved_prompt_template(
-                        st.session_state.template_spec, st.session_state.knowledge_base
-                    )
-                    # Only update if we got a valid result back
-                    if improved_template and len(improved_template) > 10:
-                        st.session_state.template_spec["prompt"] = improved_template
-                        st.success("Prompt template updated!")
+                                # Mark this variable as added
+                                st.session_state.added_suggestions.add(var_id)
 
-            # Display the prompt template
-            prompt_template = st.text_area(
-                "Edit the prompt template",
-                value=st.session_state.template_spec["prompt"],
-                height=200,
-            )
-            st.session_state.template_spec["prompt"] = prompt_template
+                                # Show success message
+                                st.success(f"Added {var['name']} to input variables!")
 
-        # Input Variables Editor
+        # Input Variables Section
         with st.expander("Input Variables", expanded=True):
-            st.subheader("Input Variables")
-
-            # Add input variable button with smart functionality
+            # Add input variable button
             col1, col2 = st.columns([3, 1])
             with col1:
                 new_input_name = st.text_input(
@@ -1431,132 +1431,223 @@ with tab2:
                     }
                     st.session_state.template_spec["input"].append(new_var)
 
-            # Display input variables
+            # Display input variables with integrated input fields
+            st.subheader("Input Variables")
+
+            # Create a container for the variables
             for i, input_var in enumerate(st.session_state.template_spec["input"]):
+                var_name = input_var["name"]
+                var_type = input_var["type"]
+                var_desc = input_var["description"]
+
                 with st.container():
-                    st.markdown(f"##### {input_var['name']}")
+                    # Variable header with description
+                    st.markdown(f"##### {var_desc}")
 
-                    col1, col2, col3 = st.columns([2, 2, 1])
+                    # Create columns for the variable controls
+                    col1, col2, col3 = st.columns([3, 1, 1])
 
                     with col1:
-                        input_var["name"] = st.text_input(
-                            "Name", value=input_var["name"], key=f"input_name_{i}"
-                        )
-
-                        input_var["description"] = st.text_input(
-                            "Description",
-                            value=input_var["description"],
-                            key=f"input_desc_{i}",
-                        )
+                        # Create the appropriate input field based on variable type
+                        if var_type == "string":
+                            st.session_state.user_inputs[var_name] = st.text_input(
+                                f"Enter value for {var_name}", key=f"use_{var_name}"
+                            )
+                        elif var_type == "int":
+                            st.session_state.user_inputs[var_name] = st.number_input(
+                                f"Enter value for {var_name}",
+                                min_value=input_var.get("min", None),
+                                max_value=input_var.get("max", None),
+                                step=1,
+                                key=f"use_{var_name}",
+                            )
+                        elif var_type == "float":
+                            st.session_state.user_inputs[var_name] = st.number_input(
+                                f"Enter value for {var_name}",
+                                min_value=float(input_var.get("min", 0)),
+                                max_value=float(input_var.get("max", 100)),
+                                key=f"use_{var_name}",
+                            )
+                        elif var_type == "bool":
+                            st.session_state.user_inputs[var_name] = st.checkbox(
+                                f"Select value for {var_name}", key=f"use_{var_name}"
+                            )
+                        elif var_type == "categorical":
+                            options = input_var.get("options", [])
+                            min_selections = input_var.get("min", 1)
+                            max_selections = input_var.get("max", 1)
+
+                            if options:
+                                if min_selections == 1 and max_selections == 1:
+                                    # Single selection
+                                    st.session_state.user_inputs[var_name] = (
+                                        st.selectbox(
+                                            f"Select value for {var_name}",
+                                            options=options,
+                                            key=f"use_{var_name}",
+                                        )
+                                    )
+                                else:
+                                    # Multi-selection
+                                    st.session_state.user_inputs[var_name] = (
+                                        st.multiselect(
+                                            f"Select {min_selections}-{max_selections} values for {var_name}",
+                                            options=options,
+                                            default=(
+                                                options[:min_selections]
+                                                if len(options) >= min_selections
+                                                else options
+                                            ),
+                                            key=f"use_{var_name}",
+                                        )
+                                    )
+                            else:
+                                st.warning(f"No options defined for {var_name}")
 
                     with col2:
-                        var_type = st.selectbox(
-                            "Type",
-                            options=["string", "int", "float", "bool", "categorical"],
-                            index=[
-                                "string",
-                                "int",
-                                "float",
-                                "bool",
-                                "categorical",
-                            ].index(input_var["type"]),
-                            key=f"input_type_{i}",
-                        )
-                        input_var["type"] = var_type
+                        # Button to edit this variable
+                        if st.button("Edit Settings", key=f"edit_input_{i}"):
+                            st.session_state.show_variable_editor = i
 
-                        if var_type in ["string", "int", "float"]:
-                            col_min, col_max = st.columns(2)
-                            with col_min:
-                                input_var["min"] = st.number_input(
-                                    "Min",
-                                    value=int(input_var.get("min", 0)),
-                                    key=f"input_min_{i}",
-                                )
-                            with col_max:
-                                input_var["max"] = st.number_input(
-                                    "Max",
-                                    value=int(input_var.get("max", 100)),
-                                    key=f"input_max_{i}",
-                                )
+                    with col3:
+                        # Button to remove this variable
+                        if st.button("Remove", key=f"remove_input_{i}"):
+                            st.session_state.template_spec["input"].pop(i)
+                            st.rerun()
 
-                        if var_type == "categorical":
-                            # Add a button to suggest options from knowledge base
-                            kb_button_key = f"suggest_input_{i}_{input_var['name']}"
-                            if st.button("Suggest Options from KB", key=kb_button_key):
-                                client = get_openai_client()
-                                if not client:
-                                    st.error(
-                                        "Please provide an OpenAI API key to suggest options."
+                    # Show editor if this variable is selected
+                    if st.session_state.show_variable_editor == i:
+                        with st.container():
+                            st.markdown("---")
+                            st.markdown(f"##### Variable Settings: {input_var['name']}")
+
+                            # Name and description
+                            input_var["name"] = st.text_input(
+                                "Name", value=input_var["name"], key=f"input_name_{i}"
+                            )
+                            input_var["description"] = st.text_input(
+                                "Description",
+                                value=input_var["description"],
+                                key=f"input_desc_{i}",
+                            )
+
+                            # Type selection
+                            var_type = st.selectbox(
+                                "Type",
+                                options=[
+                                    "string",
+                                    "int",
+                                    "float",
+                                    "bool",
+                                    "categorical",
+                                ],
+                                index=[
+                                    "string",
+                                    "int",
+                                    "float",
+                                    "bool",
+                                    "categorical",
+                                ].index(input_var["type"]),
+                                key=f"input_type_{i}",
+                            )
+                            input_var["type"] = var_type
+
+                            # Type-specific settings
+                            if var_type in ["string", "int", "float"]:
+                                col1, col2 = st.columns(2)
+                                with col1:
+                                    input_var["min"] = st.number_input(
+                                        "Min",
+                                        value=int(input_var.get("min", 0)),
+                                        key=f"input_min_{i}",
                                     )
-                                elif not st.session_state.knowledge_base:
-                                    st.warning(
-                                        "No knowledge base available. Please upload documents first."
+                                with col2:
+                                    input_var["max"] = st.number_input(
+                                        "Max",
+                                        value=int(input_var.get("max", 100)),
+                                        key=f"input_max_{i}",
                                     )
-                                else:
-                                    with st.spinner(
-                                        f"Suggesting options for {input_var['name']}..."
-                                    ):
-                                        suggestions = suggest_variable_values_from_kb(
-                                            input_var["name"],
-                                            "categorical",
-                                            st.session_state.knowledge_base,
-                                            client,
+
+                            if var_type == "categorical":
+                                # Suggest options from KB button
+                                if st.button(
+                                    "Suggest Options from KB", key=f"suggest_input_{i}"
+                                ):
+                                    client = get_openai_client()
+                                    if not client:
+                                        st.error(
+                                            "Please provide an OpenAI API key to suggest options."
                                         )
-                                        if suggestions and "options" in suggestions:
-                                            # Update the options
-                                            input_var["options"] = suggestions[
-                                                "options"
-                                            ]
-                                            st.success(
-                                                f"Found {len(suggestions['options'])} options"
-                                            )
-                                        else:
-                                            st.warning(
-                                                "Could not find suitable options in the knowledge base"
+                                    elif not st.session_state.knowledge_base:
+                                        st.warning(
+                                            "No knowledge base available. Please upload documents first."
+                                        )
+                                    else:
+                                        with st.spinner(
+                                            f"Suggesting options for {input_var['name']}..."
+                                        ):
+                                            suggestions = (
+                                                suggest_variable_values_from_kb(
+                                                    input_var["name"],
+                                                    "categorical",
+                                                    st.session_state.knowledge_base,
+                                                    client,
+                                                )
                                             )
-
-                            # Display and edit options
-                            options = input_var.get("options", [])
-                            options_str = st.text_area(
-                                "Options (one per line)",
-                                value="\n".join(options),
-                                key=f"input_options_{i}",
-                            )
-                            input_var["options"] = [
-                                opt.strip()
-                                for opt in options_str.split("\n")
-                                if opt.strip()
-                            ]
-
-                            # Add min and max for categorical variables
-                            col_min, col_max = st.columns(2)
-                            with col_min:
-                                input_var["min"] = st.number_input(
-                                    "Min selections",
-                                    value=int(input_var.get("min", 1)),
-                                    min_value=0,
-                                    key=f"input_cat_min_{i}",
-                                )
-                            with col_max:
-                                input_var["max"] = st.number_input(
-                                    "Max selections",
-                                    value=int(input_var.get("max", 1)),
-                                    min_value=1,
-                                    key=f"input_cat_max_{i}",
+                                            if suggestions and "options" in suggestions:
+                                                input_var["options"] = suggestions[
+                                                    "options"
+                                                ]
+                                                st.success(
+                                                    f"Found {len(suggestions['options'])} options"
+                                                )
+                                            else:
+                                                st.warning(
+                                                    "Could not find suitable options in the knowledge base"
+                                                )
+
+                                # Options editor
+                                options = input_var.get("options", [])
+                                options_str = st.text_area(
+                                    "Options (one per line)",
+                                    value="\n".join(options),
+                                    key=f"input_options_{i}",
                                 )
+                                input_var["options"] = [
+                                    opt.strip()
+                                    for opt in options_str.split("\n")
+                                    if opt.strip()
+                                ]
+
+                                # Min/max selections
+                                col1, col2 = st.columns(2)
+                                with col1:
+                                    input_var["min"] = st.number_input(
+                                        "Min selections",
+                                        value=int(input_var.get("min", 1)),
+                                        min_value=0,
+                                        key=f"input_cat_min_{i}",
+                                    )
+                                with col2:
+                                    input_var["max"] = st.number_input(
+                                        "Max selections",
+                                        value=int(input_var.get("max", 1)),
+                                        min_value=1,
+                                        key=f"input_cat_max_{i}",
+                                    )
 
-                    with col3:
-                        if st.button("Remove", key=f"remove_input_{i}"):
-                            st.session_state.template_spec["input"].pop(i)
-                            st.rerun()  # Only use rerun for removal
+                            # Close editor button
+                            if st.button("Done Editing", key=f"done_input_{i}"):
+                                st.session_state.show_variable_editor = None
+                                st.rerun()
+
+                            st.markdown("---")
 
                     st.divider()
 
-        # Output Variables Editor
+        # Output Variables Section
         with st.expander("Output Variables", expanded=True):
-            st.subheader("Output Variables")
-
-            # Add output variable button with smart functionality
+            # Add output variable button
             col1, col2 = st.columns([3, 1])
             with col1:
                 new_output_name = st.text_input(
@@ -1577,25 +1668,46 @@ with tab2:
                     }
                     st.session_state.template_spec["output"].append(new_var)
 
-            # Display output variables
+            # Display output variables in a table-like format
+            st.subheader("Output Variables")
+
+            # Create a container for the variables
             for i, output_var in enumerate(st.session_state.template_spec["output"]):
-                with st.container():
-                    st.markdown(f"##### {output_var['name']}")
+                col1, col2, col3 = st.columns([3, 1, 1])
 
-                    col1, col2, col3 = st.columns([2, 2, 1])
+                with col1:
+                    st.markdown(
+                        f"**{output_var['name']}** - {output_var['description']}"
+                    )
 
-                    with col1:
+                with col2:
+                    # Button to edit this variable
+                    if st.button("Edit", key=f"edit_output_{i}"):
+                        st.session_state.show_output_editor = i
+
+                with col3:
+                    # Button to remove this variable
+                    if st.button("Remove", key=f"remove_output_{i}"):
+                        st.session_state.template_spec["output"].pop(i)
+                        st.rerun()
+
+                # Show editor if this variable is selected
+                if st.session_state.show_output_editor == i:
+                    with st.container():
+                        st.markdown("---")
+                        st.markdown(f"##### Edit Output Variable: {output_var['name']}")
+
+                        # Name and description
                         output_var["name"] = st.text_input(
                             "Name", value=output_var["name"], key=f"output_name_{i}"
                         )
-
                         output_var["description"] = st.text_input(
                             "Description",
                             value=output_var["description"],
                             key=f"output_desc_{i}",
                         )
 
-                    with col2:
+                        # Type selection
                         var_type = st.selectbox(
                             "Type",
                             options=["string", "int", "float", "bool", "categorical"],
@@ -1610,15 +1722,16 @@ with tab2:
                         )
                         output_var["type"] = var_type
 
+                        # Type-specific settings
                         if var_type in ["string", "int", "float"]:
-                            col_min, col_max = st.columns(2)
-                            with col_min:
+                            col1, col2 = st.columns(2)
+                            with col1:
                                 output_var["min"] = st.number_input(
                                     "Min",
                                     value=int(output_var.get("min", 0)),
                                     key=f"output_min_{i}",
                                 )
-                            with col_max:
+                            with col2:
                                 output_var["max"] = st.number_input(
                                     "Max",
                                     value=int(output_var.get("max", 100)),
@@ -1626,9 +1739,10 @@ with tab2:
                                 )
 
                         if var_type == "categorical":
-                            # Add a button to suggest options from knowledge base
-                            kb_button_key = f"suggest_output_{i}_{output_var['name']}"
-                            if st.button("Suggest Options from KB", key=kb_button_key):
+                            # Suggest options from KB button
+                            if st.button(
+                                "Suggest Options from KB", key=f"suggest_output_{i}"
+                            ):
                                 client = get_openai_client()
                                 if not client:
                                     st.error(
@@ -1649,7 +1763,6 @@ with tab2:
                                             client,
                                         )
                                         if suggestions and "options" in suggestions:
-                                            # Update the options
                                             output_var["options"] = suggestions[
                                                 "options"
                                             ]
@@ -1661,7 +1774,7 @@ with tab2:
                                                 "Could not find suitable options in the knowledge base"
                                             )
 
-                            # Display and edit options
+                            # Options editor
                             options = output_var.get("options", [])
                             options_str = st.text_area(
                                 "Options (one per line)",
@@ -1674,16 +1787,16 @@ with tab2:
                                 if opt.strip()
                             ]
 
-                            # Add min and max for categorical variables
-                            col_min, col_max = st.columns(2)
-                            with col_min:
+                            # Min/max selections
+                            col1, col2 = st.columns(2)
+                            with col1:
                                 output_var["min"] = st.number_input(
                                     "Min selections",
                                     value=int(output_var.get("min", 1)),
                                     min_value=0,
                                     key=f"output_cat_min_{i}",
                                 )
-                            with col_max:
+                            with col2:
                                 output_var["max"] = st.number_input(
                                     "Max selections",
                                     value=int(output_var.get("max", 1)),
@@ -1691,127 +1804,46 @@ with tab2:
                                     key=f"output_cat_max_{i}",
                                 )
 
-                    with col3:
-                        if st.button("Remove", key=f"remove_output_{i}"):
-                            st.session_state.template_spec["output"].pop(i)
-                            st.rerun()  # Only use rerun for removal
-
-                    st.divider()
-
-        # Template Specification and Download Section
-        with st.expander("Template JSON", expanded=False):
-            st.json(st.session_state.template_spec)
-
-            # Download button
-            template_json = json.dumps(st.session_state.template_spec, indent=2)
-            st.download_button(
-                label="Download Template JSON",
-                data=template_json,
-                file_name="template_spec.json",
-                mime="application/json",
-            )
-    else:
-        st.info(
-            "No template has been generated yet. Go to the 'Setup' tab to create one."
-        )
-
-with tab3:
-    if st.session_state.show_template_editor and st.session_state.template_spec:
-        st.header("Use Template")
-
-        # Reset user inputs when template changes
-        if (
-            "last_template" not in st.session_state
-            or st.session_state.last_template != st.session_state.template_spec
-        ):
-            st.session_state.user_inputs = {}
-            st.session_state.last_template = st.session_state.template_spec
+                        # Close editor button
+                        if st.button("Done Editing", key=f"done_output_{i}"):
+                            st.session_state.show_output_editor = None
+                            st.rerun()
 
-        # Create input fields based on the template specification
-        for input_var in st.session_state.template_spec["input"]:
-            var_name = input_var["name"]
-            var_type = input_var["type"]
-            var_desc = input_var["description"]
-
-            st.markdown(f"##### {var_desc}")
-
-            if var_type == "string":
-                st.session_state.user_inputs[var_name] = st.text_input(
-                    f"Enter value for {var_name}", key=f"use_{var_name}"
-                )
-
-            elif var_type == "int":
-                st.session_state.user_inputs[var_name] = st.number_input(
-                    f"Enter value for {var_name}",
-                    min_value=input_var.get("min", None),
-                    max_value=input_var.get("max", None),
-                    step=1,
-                    key=f"use_{var_name}",
-                )
-
-            elif var_type == "float":
-                st.session_state.user_inputs[var_name] = st.number_input(
-                    f"Enter value for {var_name}",
-                    min_value=float(input_var.get("min", 0)),
-                    max_value=float(input_var.get("max", 100)),
-                    key=f"use_{var_name}",
-                )
-
-            elif var_type == "bool":
-                st.session_state.user_inputs[var_name] = st.checkbox(
-                    f"Select value for {var_name}", key=f"use_{var_name}"
-                )
-
-            elif var_type == "categorical":
-                options = input_var.get("options", [])
-                min_selections = input_var.get("min", 1)
-                max_selections = input_var.get("max", 1)
-
-                if options:
-                    if min_selections == 1 and max_selections == 1:
-                        # Single selection
-                        st.session_state.user_inputs[var_name] = st.selectbox(
-                            f"Select value for {var_name}",
-                            options=options,
-                            key=f"use_{var_name}",
-                        )
-                    else:
-                        # Multi-selection
-                        st.session_state.user_inputs[var_name] = st.multiselect(
-                            f"Select {min_selections}-{max_selections} values for {var_name}",
-                            options=options,
-                            default=(
-                                options[:min_selections]
-                                if len(options) >= min_selections
-                                else options
-                            ),
-                            key=f"use_{var_name}",
-                        )
-                else:
-                    st.warning(f"No options defined for {var_name}")
+                        st.markdown("---")
 
         # Handle the lore/knowledge base as a special variable
         prompt_template = st.session_state.template_spec["prompt"]
         if "{lore}" in prompt_template:
-            st.markdown("##### Document Knowledge Base")
-
-            # Display info about the knowledge base
-            if st.session_state.knowledge_base:
-                st.success(
-                    f"Using content from {len(st.session_state.uploaded_filenames) if 'uploaded_filenames' in st.session_state else 'uploaded'} documents as knowledge base"
-                )
+            with st.expander("Document Knowledge Base", expanded=True):
+                st.markdown("##### Document Knowledge Base")
 
-                with st.expander("View knowledge base content"):
-                    st.text_area(
-                        "Knowledge base content",
-                        value=st.session_state.knowledge_base[:2000]
-                        + (
-                            "..." if len(st.session_state.knowledge_base) > 2000 else ""
-                        ),
-                        height=200,
-                        disabled=True,
+                # Display info about the knowledge base
+                if st.session_state.knowledge_base:
+                    st.success(
+                        f"Using content from {len(st.session_state.uploaded_filenames) if 'uploaded_filenames' in st.session_state else 'uploaded'} documents as knowledge base"
                     )
 
+                    # Use a button to toggle knowledge base content view instead of an expander
+                    if st.button(
+                        "View/Hide Knowledge Base Content", key="toggle_kb_view"
+                    ):
+                        st.session_state.show_kb_content = not st.session_state.get(
+                            "show_kb_content", False
+                        )
+
+                    if st.session_state.get("show_kb_content", False):
+                        st.text_area(
+                            "Knowledge base content",
+                            value=st.session_state.knowledge_base[:2000]
+                            + (
+                                "..."
+                                if len(st.session_state.knowledge_base) > 2000
+                                else ""
+                            ),
+                            height=200,
+                            disabled=True,
+                        )
+
                     # Add option to edit if needed
                     use_edited_lore = st.checkbox("Edit knowledge base content")
                     if use_edited_lore:
@@ -1824,13 +1856,31 @@ with tab3:
                         st.session_state.user_inputs["lore"] = (
                             st.session_state.knowledge_base
                         )
-            else:
-                st.warning("No documents uploaded. You can provide custom lore below.")
-                st.session_state.user_inputs["lore"] = st.text_area(
-                    "Enter background information or context",
-                    placeholder="Enter custom lore or background information here...",
-                    height=150,
-                )
+                else:
+                    st.warning(
+                        "No documents uploaded. You can provide custom lore below."
+                    )
+                    st.session_state.user_inputs["lore"] = st.text_area(
+                        "Enter background information or context",
+                        placeholder="Enter custom lore or background information here...",
+                        height=150,
+                    )
+
+        # Template JSON
+        with st.expander("Template JSON", expanded=False):
+            st.json(st.session_state.template_spec)
+
+            # Download button
+            template_json = json.dumps(st.session_state.template_spec, indent=2)
+            st.download_button(
+                label="Download Template JSON",
+                data=template_json,
+                file_name="template_spec.json",
+                mime="application/json",
+            )
+
+        # Generate Output Section
+        st.header("Generate Output")
 
         # Generate Output button
         if st.button("Generate Output", key="generate_button"):
@@ -1863,9 +1913,8 @@ with tab3:
                     st.session_state.generated_output = generated_output
 
         # Display generated output
-        if st.session_state.generated_output:
+        if "generated_output" in st.session_state and st.session_state.generated_output:
             st.header("Generated Output")
-            st.markdown("### Result")
 
             # Check if the output is a dictionary (JSON)
             if isinstance(st.session_state.generated_output, dict):
@@ -1896,7 +1945,7 @@ with tab3:
             "No template has been generated yet. Go to the 'Setup' tab to create one."
         )
 
-with tab4:
+with tab3:
     if st.session_state.show_template_editor and st.session_state.template_spec:
         st.header("Generate Synthetic Data")