Update src/streamlit_app.py

src/streamlit_app.py

@@ -3,8 +3,8 @@ import os
 import json
 import streamlit as st
 from app_assets import app_logo
-from app_utils import parse_data_dictionary
-
+import io
+from app_utils import parse_data_dictionary, display_notebook, style_css
 from workflow_nodes import (
     think_sections_node,
     enhance_plan_node,
@@ -14,37 +14,40 @@ from workflow_nodes import (
     correct_notebook_node,
     write_insights_node
 )
-
 from workflow_nodes import InteractiveCaseStudyState
 
-
 st.set_page_config(layout="wide")
 
+# Apply CSS styling
+# This will apply the styles defined in style.css to the Streamlit app
+style_css("style.css")
+
 st.title("📓 NBforge Demo")
 
-st.logo(app_logo,size='large',link='https://www.mygreatlearning.com/')
+st.logo(app_logo, size='large', link='https://www.mygreatlearning.com/')
+
 
 # ─── Session State Initialization ─────────────────────────────────────────────
 if "state" not in st.session_state:
     st.session_state.state = InteractiveCaseStudyState()
-    st.session_state.plan_ready          = False
-    st.session_state.plan_approved       = None
-    st.session_state.notebook_ready      = False
-    st.session_state.notebook_approved   = None
-    st.session_state.exec_ready          = False
-    st.session_state.execution_approved  = None
+    st.session_state.plan_ready = False
+    st.session_state.plan_approved = None
+    st.session_state.notebook_ready = False
+    st.session_state.notebook_approved = None
+    st.session_state.exec_ready = False
+    st.session_state.execution_approved = None
 
 s = st.session_state.state
 
 # ─── 1) Collect Inputs ────────────────────────────────────────────────────────
 st.header("1. 🔧 Provide Your Inputs")
-s["domain"]                = st.selectbox("📂 Select Domain", [
-    "EDA","Machine Learning","Deep Learning","Natural Language Processing",
-    "Computer Vision","Generative AI Prompt Engineering","Generative AI RAG"
+s["domain"] = st.selectbox("📂 Select Domain", [
+    "EDA", "Machine Learning", "Deep Learning", "Natural Language Processing",
+    "Computer Vision", "Generative AI Prompt Engineering", "Generative AI RAG"
 ])
-s["topics_and_subtopics"]  = st.text_area("📚 Weekly Topics & Subtopics", height=80)
-s["problem_statement"]     = st.text_area("🧩 Problem Statement",       height=80)
-s["dataset_type"]          = st.selectbox("📁 Dataset Type", ["csv","pdf","images","json"])
+s["topics_and_subtopics"] = st.text_area("📚 Weekly Topics & Subtopics", height=80)
+s["problem_statement"] = st.text_area("🧩 Problem Statement", height=80)
+s["dataset_type"] = st.selectbox("📁 Dataset Type", ["csv", "pdf", "images", "json"])
 
 uploaded = st.file_uploader("⬆️ Upload Dataset", accept_multiple_files=False)
 if uploaded:
@@ -55,47 +58,44 @@ if uploaded:
     s["dataset_file_path"] = path
     st.success(f"Dataset saved to `{path}`")
 
-# Data dictionary input
 s["data_dictionary"] = {}
 dd = st.text_area("📖 Data Dictionary")
 s["data_dictionary"] = parse_data_dictionary(dd)
-# try:
-#     s["data_dictionary"] = json.loads(dd) if dd.strip() else {}
-# except json.JSONDecodeError:
-#     st.error("Invalid JSON in Data Dictionary")
 
-s["additional_instructions"] = st.text_area(
-    "📝 Additional Instructions (optional)", height=68
-)
+s["additional_instructions"] = st.text_area("📝 Additional Instructions (optional)", height=60)
+
 
 # Kick off plan generation
-if st.session_state.plan_ready==False:
+if not st.session_state.plan_ready:
     if st.button("🚀 Generate Plan") and uploaded:
         with st.spinner("Designing notebook plan…"):
             s = think_sections_node(s)
-        st.session_state.plan_ready         = True
-        st.session_state.plan_approved       = None
-        st.session_state.notebook_ready      = False
-        st.session_state.exec_ready          = False
+        st.session_state.plan_ready = True
+        st.session_state.plan_approved = None
+        st.session_state.notebook_ready = False
+        st.session_state.exec_ready = False
 
 # ─── 2) Review & Enhance Plan ────────────────────────────────────────────────
 if st.session_state.plan_ready:
     st.subheader("2. 🔍 Review Notebook Plan")
-    st.json(s["plan"])
-
-    #ask approval once
+    
+    with st.expander("📑 View Generated Plan"):
+        st.json(s["plan"])
+    
+    # ask approval once
     if st.session_state.plan_approved is None:
-        st.session_state.plan_approved = st.radio("Approve this plan?", index=None, options=["YES","NO"])
+        st.session_state.plan_approved = st.radio("Approve this plan?", index=None, options=["YES", "NO"])
         st.write(f"Plan approved: {st.session_state.plan_approved}")
 
-
+    
     # on approval, let user move on
     if st.session_state.plan_approved == "YES":
         if st.button("➡️ Next: Write Code"):
             with st.spinner("Writing notebook skeleton…"):
                 s = write_code_node(s)
-            st.session_state.notebook_ready    = True
+            st.session_state.notebook_ready = True
             st.session_state.notebook_approved = None
+
     # on rejection, collect feedback & loop back
     elif st.session_state.plan_approved == "NO":
         st.text_area("What should be changed in the plan?", key="plan_fb")
@@ -104,6 +104,7 @@ if st.session_state.plan_ready:
             s["plan_feedback"] = st.session_state.plan_fb
             with st.spinner("Modifying notebook plan…"):
                 s = enhance_plan_node(s)
+            
             # reset approval so we review again
             st.session_state.plan_approved = None
             st.rerun()
@@ -111,19 +112,29 @@ if st.session_state.plan_ready:
 # ─── 3) Review & Modify Notebook ─────────────────────────────────────────────
 if st.session_state.notebook_ready:
     st.subheader("3. 💻 Review Notebook Skeleton")
-    st.json(s["raw_notebook"])
+    with st.expander("📄 View Notebook Skeleton"):
+        display_notebook(io.StringIO(json.dumps(s["raw_notebook"], indent=2)))
+        
+        st.info("----------------------------------------------------------------")
+        st.write("**NOTE**: If it’s hard to read here, download and open in Colab:")    
+        st.download_button(
+            "Download raw_notebook.ipynb",
+            data=json.dumps(s["raw_notebook"], indent=2),
+            file_name="raw_notebook.ipynb",
+            mime="application/json",
+        )
 
     if st.session_state.notebook_approved is None:
-        st.session_state.notebook_approved = st.radio("Looks good?", index=None, options=["YES","NO"])
+        st.session_state.notebook_approved = st.radio("Looks good?", index=None, options=["YES", "NO"])
         st.write(f"notebook approved: {st.session_state.notebook_approved}")
 
     if st.session_state.notebook_approved == "YES":
         if st.button("➡️ Next: Execute Notebook"):
             with st.spinner("Executing notebook…"):
                 s = execute_notebook_node(s)
-            st.session_state.exec_ready         = True
+            st.session_state.exec_ready = True
             st.session_state.execution_approved = None
-    
+
     elif st.session_state.notebook_approved == "NO":
         st.write("🛠 Please provide your feedback on the notebook skeleton:")
         st.text_area("What needs changing in the notebook skeleton?", key="nb_fb")
@@ -135,10 +146,22 @@ if st.session_state.notebook_ready:
             st.session_state.notebook_approved = None
             st.rerun()
 
+
+
 # ─── 4) Review & Correct Execution ───────────────────────────────────────────
-if st.session_state.exec_ready: 
+if st.session_state.exec_ready:
     st.subheader("4. 🚀 Execution Results")
-    st.json(s.get("executed_notebook"))
+    with st.expander("📄 View Executed Notebook"):
+        display_notebook(io.StringIO(s["executed_notebook"]))
+
+        st.info("----------------------------------------------------------------")
+        st.write("**NOTE**: If it’s hard to read here, download and open in Colab:")    
+        st.download_button(
+            "Download executed_notebook.ipynb",
+            data=s["executed_notebook"],
+            file_name="executed_notebook.ipynb",
+            mime="application/json",
+        )
 
     if s.get("execution_error"):
         # Always reset approval until they choose Download As-Is
@@ -157,11 +180,7 @@ if st.session_state.exec_ready:
             st.session_state.execution_approved = True
 
         elif st.session_state.exec_action == "Provide Feedback":
-            fb = st.text_area(
-                "📝 Describe how I should fix the error:",
-                height=80,
-                key="exec_feedback"
-            )
+            fb = st.text_area("📝 Describe how I should fix the error:", height=80, key="exec_feedback")
             st.write(f"Execution feedback: {st.session_state.exec_feedback}")
             if st.button("🛠 Apply Feedback"):
                 s["execution_feedback"] = fb
@@ -173,7 +192,7 @@ if st.session_state.exec_ready:
                 st.session_state.execution_approved = None
                 st.rerun()
 
-        elif st.session_state.exec_action=="Auto-Correct & Retry":  # Auto-Correct & Retry
+        elif st.session_state.exec_action == "Auto-Correct & Retry":
             if st.button("🔧 Auto-Correct & Retry"):
                 s["execution_feedback"] = ""
                 with st.spinner("Auto-correcting and re-executing…"):
@@ -187,27 +206,25 @@ if st.session_state.exec_ready:
     # If execution succeeded OR user chose “Download As-Is”, show downloads
     if not s.get("execution_error") or st.session_state.execution_approved:
         st.success("✅ Ready to download your artifacts!")
-        st.download_button(
-            "📥 Download plan.json",
-            data=json.dumps(s["plan"], indent=2),
-            file_name="plan.json",
-            mime="application/json",
-        )
-        st.download_button(
-            "📥 Download raw_notebook.ipynb",
-            data=json.dumps(s["raw_notebook"], indent=2),
-            file_name="raw_notebook.ipynb",
-            mime="application/json",
-        )
-        st.download_button(
-            "📥 Download executed_notebook.ipynb",
-            data=s["executed_notebook"],
-            file_name="executed_notebook.ipynb",
-            mime="application/json",
-        )
-
-
-# ─── 5) Optionally Write Insights ─────────────────────────────────────────────
+        st.download_button("📥 Download plan.json", 
+                           data=json.dumps(s["plan"], indent=2), 
+                           file_name="plan.json", 
+                           mime="application/json"
+                           )
+        st.download_button("📥 Download raw_notebook.ipynb", 
+                           data=json.dumps(s["raw_notebook"], indent=2), 
+                           file_name="raw_notebook.ipynb", 
+                           mime="application/json"
+                           )
+        st.download_button("📥 Download executed_notebook.ipynb", 
+                           data=s["executed_notebook"], 
+                           file_name="executed_notebook.ipynb", 
+                           mime="application/json"
+                           )
+
+
+
+# ─── 5) Generate Insights (Optional) ──────────────────────────────────────────
 # Only show this once the user has downloaded (or chosen) the executed notebook
 if st.session_state.exec_ready and (
     not s.get("execution_error") or st.session_state.execution_approved
@@ -221,22 +238,25 @@ if st.session_state.exec_ready and (
     else:
         st.success("Notebook executed end-to-end without errors!")
         st.info("You can now generate insights and conclusions based on the executed notebook.")
-
+    
     # Ask permission
     st.session_state.insights_opt = st.radio(
         "Shall I proceed to write the insights and conclusion for this notebook?", index=None,
         options=["YES", "NO"]
-        )
+    )
 
     if st.session_state.insights_opt == "YES":
         if st.button("📝 Generate Insights Now"):
             with st.spinner("Writing observations & final recommendations…"):
                 s = write_insights_node(s)
             st.success("🎉 Insights generated!")
+            
             # Render the final notebook with insights inline
             st.markdown("#### Final Notebook with Insights")
-            st.json(s["final_notebook"])
-            # Offer download
+            with st.expander("📄 View Final Notebook with Insights"):
+                display_notebook(io.StringIO(json.dumps(s["final_notebook"], indent=2)))
+            
+            # Offer download     
             st.download_button(
                 "📥 Download Final Notebook (.ipynb)",
                 data=json.dumps(s["final_notebook"], indent=2),