Update app.py

app.py

@@ -401,18 +401,19 @@ def generate_documentation_page():
 
             if os.path.exists(project_folder):
                 try:
-                    # Gather all file paths in the project directory
-                    file_paths = read_project_files(project_folder)
-
-                    # Read file contents
-                    file_contents = read_files(file_paths)
-
-                    # Generate documentation using Gemini
-                    documentation = generate_detailed_documentation(file_contents, functionality)
+                    # Call Gemini to identify required functions
+                    gemini_result = identify_required_functions(project_folder, functionality)
 
                     # Display the final documentation
                     st.success("Documentation generated successfully!")
-                    st.text_area("Generated Documentation", documentation, height=600)
+                    st.text_area("Generated Documentation", gemini_result, height=600)
+
+                    # Buttons to generate PDF and Markdown files
+                    if st.button("Generate PDF"):
+                        generate_pdf(gemini_result)
+                    if st.button("Generate Markdown File"):
+                        generate_markdown_file(gemini_result)
+
                 except Exception as e:
                     st.error(f"An error occurred: {e}")
             else:
@@ -420,6 +421,55 @@ def generate_documentation_page():
         else:
             st.error("Please enter the functionality to analyze.")
 
+
+# Helper function to generate PDF
+def generate_pdf(documentation):
+    pdf = FPDF()
+    pdf.set_auto_page_break(auto=True, margin=15)
+    pdf.add_page()
+    pdf.set_font("Arial", size=12)
+
+    # Add headers and content
+    for line in documentation.splitlines():
+        if line.startswith("Project Summary:") or line.startswith("Functionality Summary:") or \
+           line.startswith("Functionality Flow:") or line.startswith("Function Documentation:"):
+            pdf.set_font("Arial", style="B", size=14)
+        elif line.startswith("*") or line.startswith("- **Function**"):
+            pdf.set_font("Arial", style="B", size=12)
+        else:
+            pdf.set_font("Arial", size=12)
+        pdf.multi_cell(0, 10, line)
+
+    # Save and download the PDF
+    pdf_file = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf")
+    pdf.output(pdf_file.name)
+    st.download_button(
+        label="Download PDF",
+        data=open(pdf_file.name, "rb").read(),
+        file_name="documentation.pdf",
+        mime="application/pdf",
+    )
+    os.unlink(pdf_file.name)
+
+# Helper function to generate Markdown file
+def generate_markdown_file(documentation):
+    # Format the documentation
+    formatted_documentation = documentation.replace("**", "**").replace(":", ":\n")
+
+    # Save as a temporary Markdown file
+    markdown_file = tempfile.NamedTemporaryFile(delete=False, suffix=".md")
+    with open(markdown_file.name, "w") as f:
+        f.write(formatted_documentation)
+
+    # Download the Markdown file
+    st.download_button(
+        label="Download Markdown File",
+        data=open(markdown_file.name, "rb").read(),
+        file_name="documentation.md",
+        mime="text/markdown",
+    )
+    os.unlink(markdown_file.name)
+
 #------------------------------------------------------------------------------------------------------------------------------------------------------------------------