Update app.py

app.py

@@ -380,9 +380,34 @@ def generate_detailed_documentation(file_contents, functionality_description):
         prompt += f"\nFile: {os.path.basename(file_path)}\n{content}\n"
 
     response = model.generate_content(prompt)
-    return response.text.strip()
-
+    return process_gemini_output(response.text.strip())
 
+def process_gemini_output(output):
+    """
+    Processes the raw output from Gemini to format it appropriately.
+    
+    - Removes ``` at the top and bottom of the output.
+    - Removes all pairs of **.
+    - Replaces ` used as a quote with '.
+    
+    Args:
+        output (str): The raw Gemini output.
+    
+    Returns:
+        str: The processed and cleaned output.
+    """
+    # Remove leading and trailing ```
+    if output.startswith("```") and output.endswith("```"):
+        output = output[3:-3].strip()
+    
+    # Remove all pairs of **
+    output = output.replace("**", "")
+    
+    # Replace ` with '
+    output = output.replace("`", "'")
+    
+    return output
+    
 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.")
@@ -473,21 +498,18 @@ def generate_pdf(documentation):
     pdf = FPDF()
     pdf.set_auto_page_break(auto=True, margin=15)
     pdf.add_page()
-    pdf.set_font("Courier", size=12)  # Monospaced font for IDE-like appearance
+    pdf.set_font("Arial", size=12)
 
-    # Process and format content
+    # Add headers and content
     for line in documentation.splitlines():
-        # Remove asterisks and backticks
-        line = line.replace("*", "").replace("`", "'")
-
-        # Detect headers and adjust formatting
         if line.startswith("Project Summary:") or line.startswith("Functionality Summary:") or \
            line.startswith("Functionality Flow:") or line.startswith("Function Documentation:"):
-            pdf.set_font("Courier", style="B", size=14)
-            pdf.cell(0, 10, line, ln=True)
-            pdf.set_font("Courier", size=12)  # Reset font for content
+            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.multi_cell(0, 10, line.strip())  # Handle regular content
+            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")
@@ -500,35 +522,15 @@ def generate_pdf(documentation):
     )
     os.unlink(pdf_file.name)
 
+
 # Helper function to generate Markdown file
 def generate_markdown_file(documentation):
-    """
-    Generate a Markdown file from the documentation with IDE-like formatting.
-    Adjustments:
-    - Removes asterisks (*)
-    - Replaces backticks (`) with single quotes (')
-    - Removes extra spaces between section headers and their content
-    - Removes Markdown-specific delimiters (e.g., ``` at the top and bottom)
-    """
-    # Process and format content
-    formatted_documentation = ""
-    for line in documentation.splitlines():
-        # Remove asterisks and backticks
-        line = line.replace("*", "").replace("`", "'")
-
-        # Add a newline only for headers for better structure in Markdown
-        if line.startswith("Project Summary:") or line.startswith("Functionality Summary:") or \
-           line.startswith("Functionality Flow:") or line.startswith("Function Documentation:"):
-            formatted_documentation += f"## {line.strip()}\n"
-        else:
-            formatted_documentation += f"{line.strip()}\n"
-
-    # Save to a temporary Markdown file
+    # 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)
+        f.write(documentation)
 
-    # Allow the user to download the Markdown file
+    # Download the Markdown file
     st.download_button(
         label="Download Markdown File",
         data=open(markdown_file.name, "rb").read(),