import streamlit as st import sqlite3 import hashlib import os import google.generativeai as genai import zipfile from git import Repo # 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 == "view_documentation": view_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): with open(file_path, 'r') as file: file_contents[file_path] = file.read() 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 += "If the functionality does not exist in the project, return 'Functionality does not exist.', otherwise, for each relevant function, provide:\n" prompt += "1. The function name.\n" prompt += "2. Dependencies on other functions or modules.\n" prompt += "Also return a more clear version of the user entered functionality at the top of your response. Provide no commentary, just what has been asked please." 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 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 identify relevant functions.") # 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) # Ensure compatibility with GitHub repositories if os.path.exists(project_folder): try: # Call the function to identify relevant functions result = identify_required_functions(project_folder, functionality) # Check if Gemini returned an error message if "No relevant functions found" in result: st.error("No relevant functions were found for the specified functionality. Please refine your input.") else: st.success("Relevant functions identified successfully!") st.text_area("Relevant Functions", result, height=400) 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.") # Button to navigate back to the project page if st.button("Back to Project"): st.session_state.page = "project_view" st.rerun() #------------------------------------------------------------------------------------------------------------------------------------------------------------------------ def view_documentation_page(): st.subheader(f"View Documentation for {st.session_state.current_project}") st.write("This page will display the generated documentation for the selected project.") if st.button("Back to Project"): st.session_state.page = "project_view" st.rerun() 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("View Documentation"): st.session_state.page = "view_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()