Update app.py

app.py

@@ -54,58 +54,63 @@ def parse_home_depot(html):
 
 # === Streamlit App ===
 
-st.set_page_config(page_title="AI Architecture Agents", layout="centered")
-st.title("🏗️ AI Architecture Assistant")
-
-agent = st.selectbox("Choose an Agent:", ["Agent 1: Floor Plan Image Generator", "Agent 3: Material Explorer"])
-
-if agent == "Agent 1: Floor Plan Image Generator":
-    st.header("🖼️ Generate a Tiny Home Floor Plan")
-    
-    if "prompt" not in st.session_state:
-        st.session_state.prompt = ""
-
-    st.text_area("Enter your prompt:", key="prompt_input", value=st.session_state.prompt, height=200)
-
-    if st.button("Insert Sample Prompt"):
-        st.session_state.prompt = sample_prompt
-        st.rerun()
-
-    st.session_state.prompt = st.session_state.prompt_input
-    st.text_area("Current prompt being used:", st.session_state.prompt, height=200, disabled=True)
-
-    if st.button("Generate Image"):
-        if st.session_state.prompt.strip():
-            with st.spinner("Generating image..."):
-                try:
-                    response = client.images.generate(
-                        model="dall-e-3",
-                        prompt=st.session_state.prompt,
-                        size="1024x1024",
-                        quality="standard",
-                        n=1
-                    )
-                    image_url = response.data[0].url
-                    st.image(image_url, caption="Generated Image", use_column_width=True)
-                    st.markdown(f"[Download Image]({image_url})", unsafe_allow_html=True)
-                except Exception as e:
-                    st.error(f"Error: {e}")
-        else:
-            st.warning("Please enter a valid prompt.")
-
-elif agent == "Agent 3: Material Explorer":
-    st.header("📦 Building Material Browser")
-    st.markdown("Browse current **laminate flooring** options from Home Depot.")
-
-    if st.button("Get Laminate Flooring Options"):
-        with st.spinner("Scraping Home Depot via Zyte..."):
-            html = zyte_scrape(LAMINATE_FLOORING_URL)
-            if html:
-                data = parse_home_depot(html)
-                if data:
-                    st.success("Here are some laminate flooring options:")
-                    st.table(data)
-                else:
-                    st.warning("No products found. Parsing may need updating.")
+def main():
+    st.set_page_config(page_title="AI Architecture Assistant", layout="centered")
+    st.title("🏗️ AI Architecture Assistant")
+
+    # === Agent Selector ===
+    agent = st.selectbox("Choose an Agent:", ["Agent 1: Floor Plan Generator", "Agent 3: Material Explorer"])
+
+    if agent == "Agent 1: Floor Plan Generator":
+        st.header("🖼️ Generate a Tiny Home Floor Plan")
+
+        if "prompt" not in st.session_state:
+            st.session_state.prompt = ""
+
+        st.text_area("Enter your prompt:", key="prompt_input", value=st.session_state.prompt, height=200)
+
+        if st.button("Insert Sample Prompt"):
+            st.session_state.prompt = sample_prompt
+            st.rerun()
+
+        st.session_state.prompt = st.session_state.prompt_input
+        st.text_area("Current prompt being used:", st.session_state.prompt, height=200, disabled=True)
+
+        if st.button("Generate Image"):
+            if st.session_state.prompt.strip():
+                with st.spinner("Generating image..."):
+                    try:
+                        response = client.images.generate(
+                            model="dall-e-3",
+                            prompt=st.session_state.prompt,
+                            size="1024x1024",
+                            quality="standard",
+                            n=1
+                        )
+                        image_url = response.data[0].url
+                        st.image(image_url, caption="Generated Image", use_column_width=True)
+                        st.markdown(f"[Download Image]({image_url})", unsafe_allow_html=True)
+                    except Exception as e:
+                        st.error(f"Error: {e}")
             else:
-                st.error("Failed to scrape the page.")
+                st.warning("Please enter a valid prompt.")
+
+    elif agent == "Agent 3: Material Explorer":
+        st.header("📦 Building Material Browser")
+        st.markdown("Browse current **laminate flooring** options from Home Depot.")
+
+        if st.button("Get Laminate Flooring Options"):
+            with st.spinner("Scraping Home Depot via Zyte..."):
+                html = zyte_scrape(LAMINATE_FLOORING_URL)
+                if html:
+                    data = parse_home_depot(html)
+                    if data:
+                        st.success("Here are some laminate flooring options:")
+                        st.table(data)
+                    else:
+                        st.warning("No products found. Parsing may need updating.")
+                else:
+                    st.error("Failed to scrape the page.")
+
+if __name__ == "__main__":
+    main()