Update app.py

app.py

@@ -422,6 +422,17 @@ def process_gemini_output(output):
     return "\n".join(processed_lines)
     
 def generate_documentation_page():
+    # Sidebar with "Log Out" and "Back to Workspace" buttons
+    st.sidebar.title(f"Project: {st.session_state.current_project}")
+    if st.sidebar.button("Back to Workspace"):
+        st.session_state.page = "workspace"
+        st.rerun()
+    if st.sidebar.button("Log Out"):
+        st.session_state.authenticated = False
+        st.session_state.username = None
+        st.session_state.page = "login"
+        st.rerun()
+
     st.subheader(f"Generate Documentation for {st.session_state.current_project}")
     st.write("Enter the functionality or parts of the project for which you'd like to generate documentation.")
 
@@ -464,7 +475,7 @@ def generate_documentation_page():
         else:
             st.error("Please enter the functionality to analyze.")
 
-    # Add export buttons if documentation is available
+    # Add export/download buttons if documentation is available
     if "generated_documentation" in st.session_state and st.session_state.generated_documentation:
         documentation = st.session_state.generated_documentation
 
@@ -475,44 +486,48 @@ 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")
 
-        # Button to generate PDF
-        if st.button("Generate PDF"):
-            try:
-                generate_pdf(documentation, pdf_path)
-                st.success(f"PDF file generated successfully: {pdf_path}")
-            except Exception as e:
-                st.error(f"Failed to generate PDF: {e}")
+        # Track whether files have been generated
+        pdf_generated = os.path.exists(pdf_path)
+        markdown_generated = os.path.exists(markdown_path)
 
-        # Button to generate Markdown file
-        if st.button("Generate Markdown File"):
+        # Show appropriate buttons based on file generation status
+        if not pdf_generated:
+            if st.button("Generate PDF"):
+                try:
+                    generate_pdf(documentation, pdf_path)
+                    st.success(f"PDF file generated successfully: {pdf_path}")
+                except Exception as e:
+                    st.error(f"Failed to generate PDF: {e}")
+        else:
+            try:
+                with open(pdf_path, "rb") as pdf_file:
+                    st.download_button(
+                        label="Download PDF",
+                        data=pdf_file.read(),
+                        file_name=f"{st.session_state.current_project}_Documentation.pdf",
+                        mime="application/pdf",
+                    )
+            except FileNotFoundError:
+                st.write("Click 'Generate PDF' to be able to download and store the documentation as a PDF")
+
+        if not markdown_generated:
+            if st.button("Generate Markdown File"):
+                try:
+                    generate_markdown_file(documentation, markdown_path)
+                    st.success(f"Markdown file generated successfully: {markdown_path}")
+                except Exception as e:
+                    st.error(f"Failed to generate Markdown file: {e}")
+        else:
             try:
-                generate_markdown_file(documentation, markdown_path)
-                st.success(f"Markdown file generated successfully: {markdown_path}")
-            except Exception as e:
-                st.error(f"Failed to generate Markdown file: {e}")
-
-        # Provide download buttons
-        try:
-            with open(pdf_path, "rb") as pdf_file:
-                st.download_button(
-                    label="Download PDF",
-                    data=pdf_file.read(),
-                    file_name=f"{st.session_state.current_project}_Documentation.pdf",
-                    mime="application/pdf",
-                )
-        except FileNotFoundError:
-            st.error("PDF file not found. Please regenerate it.")
-
-        try:
-            with open(markdown_path, "rb") as markdown_file:
-                st.download_button(
-                    label="Download Markdown File",
-                    data=markdown_file.read(),
-                    file_name=f"{st.session_state.current_project}_Documentation.md",
-                    mime="text/markdown",
-                )
-        except FileNotFoundError:
-            st.error("Markdown file not found. Please regenerate it.")
+                with open(markdown_path, "rb") as markdown_file:
+                    st.download_button(
+                        label="Download Markdown File",
+                        data=markdown_file.read(),
+                        file_name=f"{st.session_state.current_project}_Documentation.md",
+                        mime="text/markdown",
+                    )
+            except FileNotFoundError:
+                st.write("Click 'Generate Markdown File' to be able to download and store the documentation as a markdown file")
 
 
 # Helper function to generate PDF