Update app.py

app.py

@@ -422,10 +422,10 @@ def process_gemini_output(output):
     return "\n".join(processed_lines)
     
 def generate_documentation_page():
-    # Sidebar with "Log Out" and "Back to Workspace" buttons
+    # Sidebar with "Log Out" and "Back to Project" buttons
     st.sidebar.title(f"Project: {st.session_state.current_project}")
-    if st.sidebar.button("Back to Workspace"):
-        st.session_state.page = "workspace"
+    if st.sidebar.button("Back to Project"):
+        st.session_state.page = "project_view"
         st.rerun()
     if st.sidebar.button("Log Out"):
         st.session_state.authenticated = False
@@ -486,16 +486,17 @@ def generate_documentation_page():
         pdf_path = os.path.join(user_folder, f"{st.session_state.current_project}_Documentation.pdf")
         markdown_path = os.path.join(user_folder, f"{st.session_state.current_project}_Documentation.md")
 
-        # Track whether files have been generated
+        # Show appropriate buttons based on file generation status
         pdf_generated = os.path.exists(pdf_path)
         markdown_generated = os.path.exists(markdown_path)
 
-        # Show appropriate buttons based on file generation status
+        # Generate and replace buttons dynamically
         if not pdf_generated:
             if st.button("Generate PDF"):
                 try:
                     generate_pdf(documentation, pdf_path)
-                    st.success(f"PDF file generated successfully: {pdf_path}")
+                    st.session_state.pdf_generated = True  # Set a flag for PDF generation
+                    st.experimental_rerun()  # Trigger a rerun to show the download button
                 except Exception as e:
                     st.error(f"Failed to generate PDF: {e}")
         else:
@@ -514,7 +515,8 @@ def generate_documentation_page():
             if st.button("Generate Markdown File"):
                 try:
                     generate_markdown_file(documentation, markdown_path)
-                    st.success(f"Markdown file generated successfully: {markdown_path}")
+                    st.session_state.markdown_generated = True  # Set a flag for Markdown generation
+                    st.experimental_rerun()  # Trigger a rerun to show the download button
                 except Exception as e:
                     st.error(f"Failed to generate Markdown file: {e}")
         else: