Update app.py

app.py

@@ -425,10 +425,6 @@ def generate_documentation_page():
     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.")
 
-    # Get the path of the current project
-    user_folder = os.path.join("user_projects", st.session_state.username)
-    project_folder = os.path.join(user_folder, st.session_state.current_project)
-
     # Prompt user for functionality description
     functionality = st.text_area(
         "Describe the functionality",
@@ -440,6 +436,10 @@ def generate_documentation_page():
         if functionality.strip():
             st.write("Analyzing project files... Please wait.")
 
+            # Get the path of the current project
+            user_folder = os.path.join("user_projects", st.session_state.username)
+            project_folder = os.path.join(user_folder, st.session_state.current_project)
+
             if os.path.exists(project_folder):
                 try:
                     # Gather all file paths in the project directory
@@ -451,7 +451,7 @@ def generate_documentation_page():
                     # Generate documentation using Gemini
                     documentation = generate_detailed_documentation(file_contents, functionality)
 
-                    # Save the documentation in session state for exporting
+                    # Save the documentation in session state for exporting and viewing
                     st.session_state.generated_documentation = documentation
 
                     # Display the final documentation
@@ -468,23 +468,37 @@ def generate_documentation_page():
     if "generated_documentation" in st.session_state and st.session_state.generated_documentation:
         documentation = st.session_state.generated_documentation
 
+        # Define paths for PDF and Markdown files
+        user_folder = os.path.join("user_projects", st.session_state.username)
+        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"):
-            pdf_file_path = os.path.join(project_folder, f"{st.session_state.current_project} Documentation.pdf")
-            generate_pdf(documentation, pdf_file_path)
-
-            # Save the path in session state for later viewing
-            st.session_state.pdf_file_path = pdf_file_path
-            st.success("PDF file generated successfully! It will now appear in the 'View Documentation' page.")
+            generate_pdf(documentation, pdf_path)
+            st.success(f"PDF file generated successfully: {pdf_path}")
 
         # Button to generate Markdown file
         if st.button("Generate Markdown File"):
-            markdown_file_path = os.path.join(project_folder, f"{st.session_state.current_project} Documentation.md")
-            generate_markdown_file(documentation, markdown_file_path)
+            generate_markdown_file(documentation, markdown_path)
+            st.success(f"Markdown file generated successfully: {markdown_path}")
 
-            # Save the path in session state for later viewing
-            st.session_state.markdown_file_path = markdown_file_path
-            st.success("Markdown file generated successfully! It will now appear in the 'View Documentation' page.")
+        # Provide download buttons
+        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",
+            )
+
+        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",
+            )
 
 
 # Helper function to generate PDF