Update app.py

app.py

@@ -9,45 +9,47 @@ import google.generativeai as genai
 # Load environment variables
 load_dotenv()
 
-# Flask app setup
+# --- 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", "insecure-key")
+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
+# --- 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)
 
-# Gemini AI setup
-model = None
-gemini_key = os.getenv("GEMINI_API_KEY")
-if gemini_key:
-    try:
-        genai.configure(api_key=gemini_key)
-        model = genai.GenerativeModel('gemini-2.0-pro')
-    except Exception as e:
-        print(f"Gemini setup failed: {e}")
-
-def get_project_files(project_path):
-    files = {}
+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):
-        for file in filenames:
-            if file.startswith('.'): continue
-            rel_path = os.path.relpath(os.path.join(root, file), project_path)
-            try:
-                with open(os.path.join(root, file), 'r', encoding='utf-8', errors='ignore') as f:
-                    files[rel_path] = f.read()
-            except Exception as e:
-                files[rel_path] = f"Error reading: {e}"
-    return files
+        # 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'] = []
@@ -55,117 +57,188 @@ def index():
 
 @app.route('/upload', methods=['POST'])
 def upload():
-    if 'project_id' not in session:
-        session['project_id'] = str(uuid.uuid4())
-        session['chat_history'] = []
+    """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 uploaded'}), 400
+        return jsonify({'error': 'No file part in the request.'}), 400
+    
     file = request.files['project_zip']
-    if not file.filename.endswith('.zip'):
-        return jsonify({'error': 'Only .zip supported'}), 400
+    if file.filename == '' or not file.filename.endswith('.zip'):
+        return jsonify({'error': 'Invalid file. Please upload a .zip file.'}), 400
 
-    # Extract project
-    project_id = session['project_id']
+    # 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):
-        for root, dirs, files in os.walk(project_path, topdown=False):
-            for f in files: os.remove(os.path.join(root, f))
-            for d in dirs: os.rmdir(os.path.join(root, d))
-    os.makedirs(project_path, exist_ok=True)
+        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)
-    with zipfile.ZipFile(zip_path, 'r') as z:
-        z.extractall(project_path)
-    os.remove(zip_path)
-
-    # Read and prepare context
-    files = get_project_files(project_path)
-    prompt = "**Analyze this project and give suggestions for improvements, refactoring, or bugs.**\n\n"
-    for path, content in files.items():
-        prompt += f"File: `{path}`\n```\n{content[:1500]}\n```\n\n"
+    
+    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
 
-    # Ask Gemini
+    # --- New Logic: Upload file to Gemini ---
     try:
-        ai_response = "Gemini model is not configured."
-        if model:
-            reply = model.generate_content(prompt)
-            ai_response = reply.text
-    except Exception as e:
-        ai_response = f"Error talking to Gemini: {e}"
-
-    # Save chat history
-    session['chat_history'] = [
-        {"role": "user", "parts": [prompt]},
-        {"role": "model", "parts": [ai_response]}
-    ]
-    session.modified = True
-
-    return jsonify({
-        "message": "Project uploaded and analyzed by Gemini.",
-        "file_tree": list(files.keys()),
-        "chat_history": [
-            {"role": "user", "content": "Project uploaded."},
+        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():
-    if not model:
-        return jsonify({"error": "Gemini API not configured"}), 500
-    msg = request.json.get("message")
-    if not msg:
-        return jsonify({"error": "Empty message"}), 400
+    """Handles follow-up chat messages, maintaining file context."""
+    if not gemini_key:
+        return jsonify({"error": "Gemini API key not configured."}), 500
 
-    if 'project_id' not in session:
-        session['project_id'] = str(uuid.uuid4())
-        session['chat_history'] = []
+    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
 
-    session['chat_history'].append({"role": "user", "parts": [msg]})
     try:
-        chat = model.start_chat(history=session['chat_history'])
-        reply = chat.send_message(msg).text
-        session['chat_history'].append({"role": "model", "parts": [reply]})
-        return jsonify({"reply": reply})
+        # 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:
-        return jsonify({"error": str(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": []})
-    path = os.path.join(app.config['UPLOAD_FOLDER'], project_id)
-    return jsonify({"file_tree": list(get_project_files(path).keys())})
+    
+    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')
-    path = request.args.get('path')
-    if not project_id or not path:
-        return jsonify({"error": "Missing file path"}), 400
-    full = os.path.join(app.config['UPLOAD_FOLDER'], project_id, path)
+    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, 'r', encoding='utf-8', errors='ignore') as f:
+        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)})
-
+        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 uploaded", 404
-    path = os.path.join(app.config['UPLOAD_FOLDER'], project_id)
-    zip_path = os.path.join(app.config['UPLOAD_FOLDER'], f"{project_id}.zip")
-    with zipfile.ZipFile(zip_path, 'w') as zipf:
-        for root, _, files in os.walk(path):
+        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:
-                fpath = os.path.join(root, file)
-                zipf.write(fpath, os.path.relpath(fpath, path))
-    return send_from_directory(app.config['UPLOAD_FOLDER'], f"{project_id}.zip", as_attachment=True)
+                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)
+    app.run(host='0.0.0.0', port=7860)