import os import uuid import zipfile from flask import Flask, request, jsonify, render_template, session, send_from_directory from flask_session import Session from dotenv import load_dotenv import google.generativeai as genai # Load environment variables load_dotenv() # --- Gemini AI Setup --- # NOTE: It's crucial to use a model that supports file uploads, like Gemini 1.5 Pro. # The original model 'gemini-2.0-pro' is not a valid model name. # Ensure your GEMINI_API_KEY is set in your .env file. gemini_key = os.getenv("GEMINI_API_KEY") if gemini_key: genai.configure(api_key=gemini_key) # --- Flask App Setup --- app = Flask(__name__) app.config['SECRET_KEY'] = os.getenv("FLASK_SECRET_KEY", "super-secret-key-change-me") app.config['SESSION_TYPE'] = 'filesystem' app.config['SESSION_PERMANENT'] = False app.config['SESSION_FILE_DIR'] = './flask_session' Session(app) # --- Folder Setup --- UPLOAD_FOLDER = 'uploads' app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER os.makedirs(UPLOAD_FOLDER, exist_ok=True) os.makedirs(app.config['SESSION_FILE_DIR'], exist_ok=True) def get_project_file_tree(project_path): """Generates a list of file paths within a project directory.""" file_list = [] for root, _, filenames in os.walk(project_path): # Ignore hidden files and directories if any(part.startswith('.') for part in root.split(os.sep)): continue for filename in filenames: if filename.startswith('.'): continue # Create a relative path from the project_path base rel_path = os.path.relpath(os.path.join(root, filename), project_path) file_list.append(rel_path) return file_list @app.route('/') def index(): """Serves the main HTML page.""" # Initialize session if it's new if 'project_id' not in session: session['project_id'] = str(uuid.uuid4()) session['chat_history'] = [] return render_template('index.html') @app.route('/upload', methods=['POST']) def upload(): """Handles the upload and initial analysis of a .zip project.""" if not gemini_key: return jsonify({'error': 'Gemini API key is not configured on the server.'}), 500 if 'project_zip' not in request.files: return jsonify({'error': 'No file part in the request.'}), 400 file = request.files['project_zip'] if file.filename == '' or not file.filename.endswith('.zip'): return jsonify({'error': 'Invalid file. Please upload a .zip file.'}), 400 # Ensure a unique project path for the session project_id = session.get('project_id', str(uuid.uuid4())) session['project_id'] = project_id project_path = os.path.join(app.config['UPLOAD_FOLDER'], project_id) # Clean up old files if any and recreate directory if os.path.exists(project_path): import shutil shutil.rmtree(project_path) os.makedirs(project_path) # Save and extract the zip file to get the file tree for the UI zip_path = os.path.join(project_path, file.filename) file.save(zip_path) try: with zipfile.ZipFile(zip_path, 'r') as z: z.extractall(project_path) except zipfile.BadZipFile: return jsonify({'error': 'The uploaded file is not a valid zip archive.'}), 400 # --- New Logic: Upload file to Gemini --- try: print("Uploading file to Gemini API...") # We upload the saved zip file directly project_file = genai.upload_file(path=zip_path, display_name=file.filename) print(f"File uploaded successfully. Name: {project_file.name}") # Store the Gemini file reference name in the session session['gemini_file_name'] = project_file.name # --- Initial Analysis --- model = genai.GenerativeModel(model_name="gemini-1.5-pro-latest") prompt = [ "You are an expert AI code assistant.", "Analyze the entire codebase provided in this zip file. Give me a high-level summary of the project, identify potential bugs, suggest improvements or refactoring, and point out any security vulnerabilities.", "Structure your response clearly using Markdown.", project_file # Pass the file reference here ] print("Generating initial analysis...") response = model.generate_content(prompt) ai_response = response.text # Store chat history for the UI session['chat_history'] = [ {"role": "user", "content": "Project uploaded for analysis."}, {"role": "assistant", "content": ai_response} ] session.modified = True # Get file tree for the UI sidebar file_tree = get_project_file_tree(project_path) return jsonify({ "message": "Project uploaded and analyzed.", "file_tree": file_tree, "chat_history": session['chat_history'] }) except Exception as e: print(f"Error during Gemini API call: {e}") return jsonify({"error": f"An error occurred while communicating with the AI model: {e}"}), 500 @app.route('/chat', methods=['POST']) def chat(): """Handles follow-up chat messages, maintaining file context.""" if not gemini_key: return jsonify({"error": "Gemini API key not configured."}), 500 message = request.json.get("message") if not message: return jsonify({"error": "Empty message received."}), 400 gemini_file_name = session.get('gemini_file_name') if not gemini_file_name: return jsonify({"error": "No project context found. Please upload a project first."}), 400 try: # Re-instantiate the file object from the stored name project_file = genai.get_file(name=gemini_file_name) # --- Use the same model and construct the prompt with history --- model = genai.GenerativeModel(model_name="gemini-2.0-flash") # We don't need to use start_chat. We can pass the file context with each call. prompt = [ message, project_file # Pass the file reference with every message ] response = model.generate_content(prompt) # Update and save chat history if 'chat_history' not in session: session['chat_history'] = [] session['chat_history'].append({"role": "user", "content": message}) session['chat_history'].append({"role": "assistant", "content": response.text}) session.modified = True return jsonify({"reply": response.text}) except Exception as e: print(f"Error during Gemini chat: {e}") return jsonify({"error": str(e)}), 500 @app.route('/file_tree') def file_tree(): """Returns the file structure of the uploaded project.""" project_id = session.get('project_id') if not project_id: return jsonify({"file_tree": []}) project_path = os.path.join(app.config['UPLOAD_FOLDER'], project_id) if not os.path.exists(project_path): return jsonify({"file_tree": []}) return jsonify({"file_tree": get_project_file_tree(project_path)}) @app.route('/file_content') def file_content(): """Returns the content of a specific file.""" project_id = session.get('project_id') file_path = request.args.get('path') if not project_id or not file_path: return jsonify({"error": "Missing project or file path."}), 400 # Security: Prevent directory traversal attacks base_path = os.path.abspath(os.path.join(app.config['UPLOAD_FOLDER'], project_id)) full_path = os.path.abspath(os.path.join(base_path, file_path)) if not full_path.startswith(base_path): return jsonify({"error": "Access denied."}), 403 try: with open(full_path, 'r', encoding='utf-8', errors='ignore') as f: return jsonify({"content": f.read()}) except FileNotFoundError: return jsonify({"error": "File not found."}), 404 except Exception as e: return jsonify({"error": str(e)}), 500 @app.route('/download') def download(): """Allows the user to download the current state of the project as a zip.""" project_id = session.get('project_id') if not project_id: return "No project found in session.", 404 project_path = os.path.join(app.config['UPLOAD_FOLDER'], project_id) zip_filename = f"{project_id}.zip" zip_path = os.path.join(app.config['UPLOAD_FOLDER'], zip_filename) with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf: for root, _, files in os.walk(project_path): # Do not include the original uploaded zip file in the new zip if os.path.basename(root) == project_id and any(f.endswith('.zip') for f in files): files = [f for f in files if not f.endswith('.zip')] for file in files: file_full_path = os.path.join(root, file) # Write file with a relative path inside the zip zipf.write(file_full_path, os.path.relpath(file_full_path, project_path)) return send_from_directory( app.config['UPLOAD_FOLDER'], zip_filename, as_attachment=True, download_name=f"project_{project_id}.zip" # Set a more user-friendly download name ) if __name__ == "__main__": app.run(host='0.0.0.0', port=7860)