import streamlit as st import sqlite3 import hashlib import os import google.generativeai as genai import zipfile from git import Repo from transformers import AutoModelForCausalLM, AutoTokenizer import torch import requests from fpdf import FPDF import markdown import tempfile import re import spacy from collections import defaultdict # Database setup DB_FILE = "users.db" def create_user_table(): conn = sqlite3.connect(DB_FILE) cursor = conn.cursor() cursor.execute(""" CREATE TABLE IF NOT EXISTS users ( username TEXT PRIMARY KEY, password TEXT ) """) conn.commit() conn.close() def add_user(username, password): conn = sqlite3.connect(DB_FILE) cursor = conn.cursor() hashed_password = hashlib.sha256(password.encode()).hexdigest() try: cursor.execute("INSERT INTO users (username, password) VALUES (?, ?)", (username, hashed_password)) conn.commit() except sqlite3.IntegrityError: st.error("Username already exists. Please choose a different username.") conn.close() def authenticate_user(username, password): conn = sqlite3.connect(DB_FILE) cursor = conn.cursor() hashed_password = hashlib.sha256(password.encode()).hexdigest() cursor.execute("SELECT * FROM users WHERE username = ? AND password = ?", (username, hashed_password)) user = cursor.fetchone() conn.close() return user def initialize_session_state(): if "authenticated" not in st.session_state: st.session_state.authenticated = False if "username" not in st.session_state: st.session_state.username = None if "page" not in st.session_state: st.session_state.page = "login" if "current_project" not in st.session_state: st.session_state.current_project = None if "project_uploaded" not in st.session_state: st.session_state.project_uploaded = False def main(): st.title("SimplifAI") # Initialize session state initialize_session_state() # Initialize database create_user_table() # Page routing logic if st.session_state.page == "login": login_page() elif st.session_state.page == "workspace": workspace_page() elif st.session_state.page == "project_view": project_view_page() elif st.session_state.page == "generate_documentation": generate_documentation_page() elif st.session_state.page == "saved_documentation": saved_documentation_page() def login_page(): st.subheader("Please Log In or Register to Continue") auth_mode = st.radio("Choose an Option", ["Log In", "Register"], horizontal=True) if auth_mode == "Log In": st.subheader("Log In") username = st.text_input("Username", key="login_username") password = st.text_input("Password", type="password", key="login_password") # Handle single-click login if st.button("Log In"): if authenticate_user(username, password): st.session_state.authenticated = True st.session_state.username = username 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") password = st.text_input("Create Password", type="password", key="register_password") # Handle single-click registration if st.button("Register"): if username and password: add_user(username, password) st.success("Account created successfully! You can now log in.") else: st.error("Please fill in all fields.") def workspace_page(): # Sidebar with logout button st.sidebar.title(f"Hello, {st.session_state.username}!") if st.sidebar.button("Log Out"): st.session_state.authenticated = False st.session_state.username = None st.session_state.page = "login" # User's folder for projects user_folder = os.path.join("user_projects", st.session_state.username) os.makedirs(user_folder, exist_ok=True) # Refresh project list dynamically projects = [d for d in os.listdir(user_folder) if os.path.isdir(os.path.join(user_folder, d))] # Display "Projects" dropdown selected_project = st.sidebar.selectbox("Projects", ["Select a project"] + projects) if selected_project != "Select a project": st.session_state.current_project = selected_project st.session_state.page = "project_view" st.rerun() # Display success message if a project was uploaded if st.session_state.project_uploaded: st.success(f"Project '{st.session_state.current_project}' uploaded successfully!") st.session_state.project_uploaded = False # Reset after showing the message # Main content area st.subheader("Workspace") st.write("You can create a new project by uploading files or folders, or by cloning a GitHub repository.") # Existing upload functionality... # User action selection action = st.radio("Choose an action", ["Upload Files or Folders", "Clone GitHub Repository"], horizontal=True) project_name = st.text_input("Enter a project name") if action == "Upload Files or Folders": st.subheader("Upload Files or Folders") uploaded_files = st.file_uploader( "Upload one or more files or a .zip archive for folders", accept_multiple_files=True ) if uploaded_files and project_name: if st.button("Upload Project"): project_folder = os.path.join(user_folder, project_name) os.makedirs(project_folder, exist_ok=True) for uploaded_file in uploaded_files: # Save uploaded .zip files or regular files file_path = os.path.join(project_folder, uploaded_file.name) with open(file_path, "wb") as f: f.write(uploaded_file.getbuffer()) # If a .zip file is uploaded, extract its contents if uploaded_file.name.endswith(".zip"): try: with zipfile.ZipFile(file_path, "r") as zip_ref: zip_ref.extractall(project_folder) os.remove(file_path) # Remove the .zip file after extraction st.success(f"Folder from {uploaded_file.name} extracted successfully!") except zipfile.BadZipFile: st.error(f"File {uploaded_file.name} is not a valid .zip file.") else: st.success(f"File {uploaded_file.name} saved successfully!") # Update session state and trigger a rerun st.session_state.current_project = project_name st.session_state.project_uploaded = True st.rerun() elif action == "Clone GitHub Repository": st.subheader("Clone GitHub Repository") repo_url = st.text_input("Enter the GitHub repository URL") if repo_url and project_name: if st.button("Upload Project"): project_folder = os.path.join(user_folder, project_name) os.makedirs(project_folder, exist_ok=True) try: Repo.clone_from(repo_url, project_folder) # Update session state and trigger a rerun st.session_state.current_project = project_name st.session_state.project_uploaded = True st.rerun() except Exception as e: st.error(f"Failed to clone repository: {e}") # Configure Gemini API gemini = os.getenv("GEMINI") genai.configure(api_key=gemini) model = genai.GenerativeModel("gemini-1.5-flash") def read_project_files(project_path): """Reads all files in the project directory and its subdirectories.""" file_paths = [] for root, _, files in os.walk(project_path): for file in files: # Skip .git files or folders from GitHub clones if ".git" not in root: file_paths.append(os.path.join(root, file)) return file_paths def read_files(file_paths): """Reads content from a list of file paths.""" file_contents = {} for file_path in file_paths: if os.path.exists(file_path): try: # Attempt to read the file as UTF-8 with open(file_path, 'r', encoding='utf-8') as file: file_contents[file_path] = file.read() except UnicodeDecodeError: print(f"Skipping binary or non-UTF-8 file: {file_path}") else: print(f"File not found: {file_path}") return file_contents def generate_prompt(file_contents, functionality_description): """Generates a prompt for Gemini to analyze the files.""" prompt = "Analyze the following code files to identify all functions required to implement the functionality: " prompt += f"'{functionality_description}'.\n\n" 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 += """ Return your output in the following format providing no commentary: Project Summary: Functionality: Functions: : -: -Function Dependencies: """ return prompt 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 def extract_cleaned_gemini_output(gemini_output): """ Removes the 'Functions' section and any irrelevant content from the Gemini output. Args: gemini_output (str): The raw output returned by Gemini. Returns: str: Cleaned output without the 'Functions' section. """ lines = gemini_output.splitlines() cleaned_output = [] skip_section = False for line in lines: line = line.strip() # Detect the start of the "Functions:" section and skip it if line.startswith("Functions:"): skip_section = True # Detect the end of the "Functions:" section and resume processing if skip_section and (line.startswith("Tasks:") or line.startswith("Project Summary:")): skip_section = False # Append lines outside the skipped section if not skip_section: cleaned_output.append(line) # Return the cleaned output with blank lines filtered out return "\n".join(line for line in cleaned_output if line) def split_into_chunks(content, chunk_size=1000): """Splits large content into smaller chunks.""" return [content[i:i + chunk_size] for i in range(0, len(content), chunk_size)] def generate_detailed_documentation(file_contents, functionality_description): """ Generates detailed documentation using Gemini directly. Args: file_contents (dict): A dictionary with file paths as keys and their content as values. functionality_description (str): A description of the functionality to document. Returns: str: The generated documentation. """ 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: ' Project Summary: ' 2. Refine the user-defined functionality: ' Functionality Summary: ' 3. Describe the functionality flow: ' Functionality Flow: ' 4. Generate detailed documentation for each function in the codebase: ' Function Documentation: For each function: - '': - Summary: - Inputs:
- Outputs:
- Dependencies: - Data structures:
- Algorithmic Details: - Error Handling: - Assumptions: - Example Usage: ' Please return only the required documentation in the specified format. Code files: """ for file_path, content in file_contents.items(): prompt += f"\nFile: {os.path.basename(file_path)}\n{content}\n" response = model.generate_content(prompt) 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 '. - 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 = [] for line in lines: stripped_line = line.strip() if stripped_line.startswith("*"): # Only replace if the '*' is at the start or followed by whitespace 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}") if st.sidebar.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"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.") # Prompt user for functionality description functionality = st.text_area( "Describe the functionality", placeholder="e.g., Explain the function of the file `main.py`", ) # Button to start analyzing functionality if st.button("Analyze"): 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 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) # Save the documentation in session state for exporting and viewing st.session_state.generated_documentation = documentation # 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: Functionality Flow: """ # 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: st.error("Project folder not found. Ensure the GitHub repository was cloned successfully.") 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 # Define paths for PDF and Markdown files user_folder = os.path.join("user_projects", st.session_state.username) os.makedirs(user_folder, exist_ok=True) # Ensure the user folder exists 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") # Show appropriate buttons based on file generation status pdf_generated = os.path.exists(pdf_path) markdown_generated = os.path.exists(markdown_path) # Generate and replace buttons dynamically if not pdf_generated: if st.button("Generate PDF"): try: generate_pdf(documentation, pdf_path) st.session_state.pdf_generated = True # Set a flag for PDF generation st.rerun() # Trigger a rerun to show the download button 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.session_state.markdown_generated = True # Set a flag for Markdown generation st.rerun() # Trigger a rerun to show the download button except Exception as e: st.error(f"Failed to generate Markdown file: {e}") else: 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.write("Click 'Generate Markdown File' to be able to download and store the documentation as a markdown file") # Helper function to generate PDF def generate_pdf(documentation, pdf_path): pdf = FPDF() pdf.set_auto_page_break(auto=True, margin=15) pdf.add_page() pdf.set_font("Courier", size=12) # Maintain the IDE-like font style # Add headers and content for line in documentation.splitlines(): try: if line.startswith("- '"): pdf.set_font("Courier", style="B", size=12) # Bold for specific lines pdf.multi_cell(0, 10, line.encode('latin-1', 'replace').decode('latin-1')) elif 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) # Bold larger headers pdf.multi_cell(0, 10, line.encode('latin-1', 'replace').decode('latin-1')) else: pdf.set_font("Courier", size=12) # Regular for other lines pdf.multi_cell(0, 10, line.encode('latin-1', 'replace').decode('latin-1')) except UnicodeEncodeError as e: pdf.multi_cell(0, 10, line.encode('latin-1', 'replace').decode('latin-1')) # Save the PDF pdf.output(pdf_path) # Helper function to generate Markdown file 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. """ formatted_lines = [] for line in documentation.splitlines(): # Add bold formatting to specific lines (e.g., function names) if line.startswith("- '") and line.endswith("':"): formatted_lines.append(f"**{line.strip()}**") # Add section headers with markdown syntax elif line.startswith("Project Summary:") or line.startswith("Functionality Summary:") or \ line.startswith("Functionality Flow:") or line.startswith("Function Documentation:"): formatted_lines.append(f"# {line.strip()}") else: # Regular text as-is formatted_lines.append(line) # Write the formatted lines to a markdown file 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"): 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", ) def project_view_page(): # Sidebar with logout and return 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() # Main content for project page st.subheader(f"Project: {st.session_state.current_project}") st.write("Manage your project and explore its files.") # Buttons for documentation functionality if st.button("Generate Documentation"): st.session_state.page = "generate_documentation" st.rerun() if st.button("Saved Documentation"): st.session_state.page = "saved_documentation" st.rerun() # Toggle file structure display (if required) if "show_file_structure" not in st.session_state: st.session_state.show_file_structure = False if st.button("Show File Structure"): st.session_state.show_file_structure = not st.session_state.show_file_structure if st.session_state.show_file_structure: user_folder = os.path.join("user_projects", st.session_state.username) project_folder = os.path.join(user_folder, st.session_state.current_project) st.write("File structure:") for root, dirs, files in os.walk(project_folder): level = root.replace(project_folder, "").count(os.sep) indent = " " * 4 * level if level == 0: st.write(f"📂 {os.path.basename(root)}") else: with st.expander(f"{indent}📂 {os.path.basename(root)}"): sub_indent = " " * 4 * (level + 1) for file in files: st.write(f"{sub_indent}📄 {file}") if __name__ == "__main__": main()