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 # 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): 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 += "If the functionality does not exist in the project, return 'Functionality does not exist.', otherwise, 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 += "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 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)] # Hugging Face Inference API endpoint for the model API_URL = "https://api-inference.huggingface.co/models/Qwen/Qwen2.5-Coder-32B-Instruct" qwen = os.getenv("QWEN") headers = {"Authorization": f"Bearer {qwen}"} def clean_output(output): """ Cleans the output from the Hugging Face model to ensure no unnecessary details are included. """ # Remove any remnants of prompts or file content (heuristically clean up known patterns) lines = output.splitlines() filtered_lines = [ line for line in lines if not ( line.startswith("File:") or line.startswith("User-specified functionality:") or line.startswith("Functions identified by Gemini:") or line.strip() == "" ) ] return "\n".join(filtered_lines) def validate_and_generate_documentation(api_url, headers, gemini_output, file_contents, functionality_description): """ Uses the Hugging Face Inference API to validate functions and generate documentation. Ensures that the output excludes file content and internal prompts. """ # Generate the prompt for the Qwen model prompt = f""" User-specified functionality: '{functionality_description}' Functions identified by Gemini: {gemini_output} Task: 1. Validate if the functions identified by Gemini are sufficient for implementing the functionality. 2. If not, identify all additional functions required. 3. For all relevant functions, generate detailed documentation in this format: Project Summary: Functionality Summary: Functionality Flow: Function Documentation: For each relevant function: - Summary: - Inputs:
- Outputs:
- Dependencies: - Data structures:
- Algorithmic Details: - Error Handling: - Assumptions: - Example Usage: """ # Send the prompt to the Hugging Face API payload = {"inputs": prompt} response = requests.post(api_url, headers=headers, json=payload) # Handle the API response if response.status_code == 200: api_response = response.json() if isinstance(api_response, list): # If response is a list, extract the first element output = api_response[0].get("generated_text", "No output generated.") elif isinstance(api_response, dict): # Handle dict response output = api_response.get("generated_text", "No output generated.") else: raise ValueError("Unexpected response format from Hugging Face API.") # Clean the output to exclude any file contents or prompts cleaned_output = clean_output(output) return cleaned_output else: raise ValueError(f"Error: {response.status_code}, {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) if os.path.exists(project_folder): try: # Call Gemini to identify required functions gemini_result = identify_required_functions(project_folder, functionality) # Read project files file_paths = read_project_files(project_folder) file_contents = read_files(file_paths) # Call the Hugging Face API for validation and documentation final_documentation = validate_and_generate_documentation( API_URL, headers, gemini_result, file_contents, functionality ) # Display the final documentation st.success("Documentation generated successfully!") st.text_area("Generated Documentation", final_documentation, height=600) 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()