Update app.py

app.py

@@ -11,9 +11,6 @@ import requests
 from fpdf import FPDF
 import markdown
 import tempfile
-import re
-import spacy
-from collections import defaultdict
 
 
 # Database setup
@@ -81,8 +78,8 @@ def main():
         project_view_page()
     elif st.session_state.page == "generate_documentation":
         generate_documentation_page()
-    elif st.session_state.page == "saved_documentation":
-        saved_documentation_page()
+    elif st.session_state.page == "view_documentation":
+        view_documentation_page()
 
 def login_page():
     st.subheader("Please Log In or Register to Continue")
@@ -101,7 +98,7 @@ def login_page():
                 st.session_state.page = "workspace"
             else:
                 st.error("Invalid username or password. Please try again.")
-    
+
     elif auth_mode == "Register":
         st.subheader("Register")
         username = st.text_input("Create Username", key="register_username")
@@ -208,6 +205,8 @@ def workspace_page():
                 except Exception as e:
                     st.error(f"Failed to clone repository: {e}")
 
+#------------------------------------------------------------------------------------------------------------------------------------------------------------------------
+
 
 # Configure Gemini API
 gemini = os.getenv("GEMINI")
@@ -249,15 +248,14 @@ def generate_prompt(file_contents, functionality_description):
 
     for file_path, content in file_contents.items():
         prompt += f"File: {os.path.basename(file_path)}\n{content}\n\n"
-    
+
     prompt += "For each relevant function, provide:\n"
     prompt += "1. Which file the function is found in.\n"
     prompt += "2. The function name.\n"
     prompt += "3. Dependencies on other functions or modules.\n"
 
-    prompt += """ 
+    prompt += """
             Return your output in the following format providing no commentary:
-
             Project Summary:
                 <A summary of the project>
             Functionality:
@@ -275,16 +273,16 @@ def identify_required_functions(project_path, functionality_description):
     """Identifies required functions for a specified functionality."""
     # Gather all file paths in the project directory
     file_paths = read_project_files(project_path)
-    
+
     # Read file contents
     file_contents = read_files(file_paths)
-    
+
     # Generate a refined prompt for Gemini
     prompt = generate_prompt(file_contents, functionality_description)
-    
+
     # Call the Gemini model
     response = model.generate_content(prompt)
-    
+
     # Process and return the response
     return response.text
 
@@ -339,9 +337,7 @@ def generate_detailed_documentation(file_contents, functionality_description):
     """
     prompt = f"""
     The following code files are provided. Analyze their contents and generate comprehensive documentation.
-
     Functionality description: '{functionality_description}'
-
     Tasks:
     1. Generate a project summary:
         '
@@ -375,7 +371,6 @@ def generate_detailed_documentation(file_contents, functionality_description):
           - Example Usage: <Example demonstrating usage>
         '
     Please return only the required documentation in the specified format.
-
     Code files:
     """
     for file_path, content in file_contents.items():
@@ -387,28 +382,28 @@ def generate_detailed_documentation(file_contents, functionality_description):
 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 '.
     - Replaces leading '*' with '-' for lists.
-    
+
     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("`", "'")
-    
+
     # Replace leading '*' with '-' for list items
     lines = output.splitlines()
     processed_lines = []
@@ -419,9 +414,9 @@ def process_gemini_output(output):
             processed_lines.append(line.replace("*", "-", 1))
         else:
             processed_lines.append(line)
-    
+
     return "\n".join(processed_lines)
-    
+
 def generate_documentation_page():
     # Sidebar with "Log Out" and "Back to Project" buttons
     st.sidebar.title(f"Project: {st.session_state.current_project}")
@@ -469,26 +464,6 @@ def generate_documentation_page():
                     # Display the final documentation
                     st.success("Documentation generated successfully!")
                     st.text_area("Generated Documentation", documentation, height=600)
-
-                    #---------------------------------------------------------------------------
-                    # Example Gemini output for testing
-                    test_gemini_output = """
-                    Functionality Summary:
-                    Authenticate the user by verifying their credentials. Once authenticated, fetch user-specific data such as preferences and usage history. Handle errors during both authentication and data retrieval.
-
-                    Project Summary:
-                    <Some project details>
-
-                    Functionality Flow:
-                    <Flow details>
-                    """
-
-                    # Streamlit testing
-                    if "page" not in st.session_state or st.session_state.page == "test_functionality":
-                        st.title("Test Functionality Extraction")
-                        display_functionality_analysis(test_gemini_output)
-                    #---------------------------------------------------------------------------
-                    
                 except Exception as e:
                     st.error(f"An error occurred: {e}")
             else:
@@ -496,8 +471,6 @@ def generate_documentation_page():
         else:
             st.error("Please enter the functionality to analyze.")
 
-        
-
     # 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
@@ -588,7 +561,6 @@ def generate_pdf(documentation, pdf_path):
 def generate_markdown_file(documentation, output_path):
     """
     Generates a markdown file from the documentation with structured formatting.
-
     Args:
         documentation (str): The documentation content to format and save.
         output_path (str): Path to save the generated markdown file.
@@ -610,121 +582,47 @@ def generate_markdown_file(documentation, output_path):
     with open(output_path, "w") as f:
         f.write("\n".join(formatted_lines))
 
-#---------------------------------------------------------------------------------------------------------------------------------------------------------------------------
-
-# Load the NLP model (use the small model for speed)
-nlp = spacy.load("en_core_web_sm")
-
-def extract_functionality_and_relationships(gemini_output):
-    """
-    Extracts functionality descriptions and sub-functionality relationships from Gemini output.
-    Args:
-        gemini_output (str): The raw output from Gemini.
-    Returns:
-        dict: A dictionary with functionality and its relationships.
-    """
-    functionality_summary = None
-    relationships = defaultdict(list)
-
-    # Extract the 'Functionality Summary' section
-    functionality_match = re.search(r"Functionality Summary:\n(.*?)\n\n", gemini_output, re.DOTALL)
-    if functionality_match:
-        functionality_summary = functionality_match.group(1).strip()
-
-    # Parse the functionality summary with NLP
-    if functionality_summary:
-        doc = nlp(functionality_summary)
-
-        # Extract relationships between sub-functionalities
-        for sentence in doc.sents:
-            sub_functionalities = []
-            for token in sentence:
-                # Capture sub-functionalities using noun chunks or proper nouns
-                if token.dep_ in {"nsubj", "dobj", "pobj"} or token.ent_type_ in {"ORG", "PRODUCT"}:
-                    sub_functionalities.append(token.text)
-            if len(sub_functionalities) > 1:
-                # Create relationships between identified sub-functionalities
-                main_functionality = sub_functionalities[0]
-                for related in sub_functionalities[1:]:
-                    relationships[main_functionality].append(related)
-
-    return functionality_summary, dict(relationships)
-
-
-def display_functionality_analysis(gemini_output):
-    """
-    Extracts and displays functionality analysis for testing purposes.
-    Args:
-        gemini_output (str): The raw output from Gemini.
-    """
-    # Extract functionality and relationships
-    functionality, relationships = extract_functionality_and_relationships(gemini_output)
-
-    # Display extracted functionality
-    st.subheader("Extracted Functionality Summary")
-    if functionality:
-        st.text_area("Functionality Summary", functionality, height=200)
-    else:
-        st.error("No functionality summary found in the Gemini output.")
-
-    # Display extracted relationships
-    st.subheader("Extracted Relationships")
-    if relationships:
-        for main, related in relationships.items():
-            st.write(f"**{main}** is related to: {', '.join(related)}")
-    else:
-        st.info("No relationships between sub-functionalities identified.")
-
-
-
 
 
 #------------------------------------------------------------------------------------------------------------------------------------------------------------------------
-def saved_documentation_page():
-    # Sidebar with "Back to Project" and "Log Out" buttons
-    st.sidebar.title(f"Project: {st.session_state.current_project}")
-    if st.sidebar.button("Back to Project"):
+def view_documentation_page():
+    st.subheader(f"View Documentation for {st.session_state.current_project}")
+    st.write("This page displays the generated documentation for the selected project.")
+
+    # Check if files are available in session state
+    pdf_path = st.session_state.get("pdf_file_path")
+    markdown_path = st.session_state.get("markdown_file_path")
+
+    # Display PDF file if it exists
+    if pdf_path and os.path.exists(pdf_path):
+        st.write("### PDF Documentation")
+        with open(pdf_path, "rb") as pdf_file:
+            pdf_data = pdf_file.read()
+            st.download_button(
+                label="Download PDF",
+                data=pdf_data,
+                file_name=os.path.basename(pdf_path),
+                mime="application/pdf",
+            )
+        st.write("View PDF Documentation below:")
+        st.pdf_viewer(pdf_path)
+
+    # Display Markdown file if it exists
+    if markdown_path and os.path.exists(markdown_path):
+        st.write("### Markdown Documentation")
+        with open(markdown_path, "r") as md_file:
+            markdown_content = md_file.read()
+            st.download_button(
+                label="Download Markdown",
+                data=markdown_content,
+                file_name=os.path.basename(markdown_path),
+                mime="text/markdown",
+            )
+        st.markdown(f"```\n{markdown_content}\n```", unsafe_allow_html=True)
+
+    if st.button("Back to Project"):
         st.session_state.page = "project_view"
         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"Saved Documentation for {st.session_state.current_project}")
-    st.write("Below is a list of saved documentation for this project. Click to download.")
-
-    # Define paths for generated 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")
-
-    # Check if files exist
-    pdf_exists = os.path.exists(pdf_path)
-    markdown_exists = os.path.exists(markdown_path)
-
-    # Display available documentation
-    if not pdf_exists and not markdown_exists:
-        st.info("No documentation has been generated yet. Please go to the 'Generate Documentation' page.")
-    else:
-        if pdf_exists:
-            with open(pdf_path, "rb") as pdf_file:
-                st.download_button(
-                    label="Download PDF Documentation",
-                    data=pdf_file.read(),
-                    file_name=f"{st.session_state.current_project}_Documentation.pdf",
-                    mime="application/pdf",
-                )
-
-        if markdown_exists:
-            with open(markdown_path, "rb") as markdown_file:
-                st.download_button(
-                    label="Download Markdown Documentation",
-                    data=markdown_file.read(),
-                    file_name=f"{st.session_state.current_project}_Documentation.md",
-                    mime="text/markdown",
-                )
 
 
 
@@ -749,8 +647,8 @@ def project_view_page():
         st.session_state.page = "generate_documentation"
         st.rerun()
 
-    if st.button("Saved Documentation"):
-        st.session_state.page = "saved_documentation"
+    if st.button("View Documentation"):
+        st.session_state.page = "view_documentation"
         st.rerun()
 
     # Toggle file structure display (if required)