Update app.py

app.py

@@ -3,7 +3,7 @@ import zipfile
 import uuid
 import google.generativeai as genai
 from flask import Flask, request, jsonify, render_template, session, send_from_directory
-from flask_session import Session # Import the extension
+from flask_session import Session
 from dotenv import load_dotenv
 
 # --- API and App Configuration ---
@@ -14,10 +14,9 @@ app.config["SECRET_KEY"] = os.environ.get("FLASK_SECRET_KEY", 'default-insecure-
 
 # --- SERVER-SIDE SESSION CONFIGURATION ---
 app.config["SESSION_PERMANENT"] = False
-app.config["SESSION_TYPE"] = "filesystem" # Store sessions on the server's filesystem
-app.config['SESSION_FILE_DIR'] = './flask_session' # Directory to store session files
-Session(app) # Initialize the session extension
-# --- END OF SESSION CONFIGURATION ---
+app.config["SESSION_TYPE"] = "filesystem"
+app.config['SESSION_FILE_DIR'] = './flask_session'
+Session(app)
 
 # Configure the Gemini API client
 try:
@@ -40,8 +39,6 @@ if not os.path.exists(UPLOAD_FOLDER):
 if not os.path.exists('./flask_session'):
     os.makedirs('./flask_session')
 
-
-# --- Helper Functions (No Changes) ---
 def get_project_files(project_path):
     file_data = {}
     for root, _, files in os.walk(project_path):
@@ -57,8 +54,6 @@ def get_project_files(project_path):
                 file_data[relative_path] = f"Error reading file: {e}"
     return file_data
 
-
-# --- Flask Routes ---
 @app.route('/')
 def index():
     if 'project_id' not in session:
@@ -66,7 +61,6 @@ def index():
         session['chat_history'] = []
     return render_template('index.html')
 
-
 @app.route('/upload', methods=['POST'])
 def upload_project():
     if 'project_id' not in session:
@@ -75,6 +69,7 @@ def upload_project():
 
     if 'project_zip' not in request.files:
         return jsonify({"error": "No file part"}), 400
+    
     file = request.files['project_zip']
     if file.filename == '':
         return jsonify({"error": "No selected file"}), 400
@@ -82,16 +77,25 @@ def upload_project():
     if file and file.filename.endswith('.zip'):
         project_id = session['project_id']
         project_path = os.path.join(app.config['UPLOAD_FOLDER'], project_id)
-        if not os.path.exists(project_path):
-            os.makedirs(project_path)
+        
+        # Clear existing project files if they exist
+        if os.path.exists(project_path):
+            for root, dirs, files in os.walk(project_path, topdown=False):
+                for name in files:
+                    os.remove(os.path.join(root, name))
+                for name in dirs:
+                    os.rmdir(os.path.join(root, name))
+        
+        os.makedirs(project_path, exist_ok=True)
         zip_path = os.path.join(project_path, file.filename)
         file.save(zip_path)
+        
         with zipfile.ZipFile(zip_path, 'r') as zip_ref:
             zip_ref.extractall(project_path)
         os.remove(zip_path)
 
         project_files = get_project_files(project_path)
-        
+
         initial_context = "The user has uploaded a new project. Here is the file structure and content:\n\n"
         for path, content in project_files.items():
             initial_context += f"**File:** `{path}`\n```\n{content}\n```\n\n"
@@ -103,7 +107,7 @@ def upload_project():
         session.modified = True
 
         frontend_chat_history = [
-            {"role": "user", "content": initial_context},
+            {"role": "user", "content": "Project context sent to AI."},
             {"role": "assistant", "content": "I have analyzed the project. I am ready to help. How can I assist you with your code?"}
         ]
 
@@ -131,15 +135,15 @@ def chat():
     session['chat_history'].append({"role": "user", "parts": [user_message]})
     session.modified = True
 
-    chat_session = model.start_chat(history=session['chat_history'])
-    
     try:
+        chat_session = model.start_chat(history=session['chat_history'])
         response = chat_session.send_message(user_message)
         ai_response = response.text
 
         session['chat_history'].append({"role": "model", "parts": [ai_response]})
         session.modified = True
 
+        # Handle file modifications from AI response
         if "```" in ai_response:
             lines = ai_response.split('\n')
             file_path_line = lines[0].strip()
@@ -168,51 +172,68 @@ def chat():
     except Exception as e:
         return jsonify({"error": str(e)}), 500
 
-
 @app.route('/file_tree')
 def get_file_tree():
     project_id = session.get('project_id')
     if not project_id:
-        return jsonify({"error": "No project in session."}), 400
+        return jsonify({"error": "No project in session. Please upload a project first."}), 400
+    
     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": list(get_project_files(project_path).keys())})
-
+    
+    try:
+        project_files = get_project_files(project_path)
+        return jsonify({"file_tree": list(project_files.keys())})
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
 
 @app.route('/file_content')
 def get_file_content():
     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
+    
+    if not project_id:
+        return jsonify({"error": "No project in session. Please upload a project first."}), 400
+    if not file_path:
+        return jsonify({"error": "Missing file path"}), 400
+    
     full_path = os.path.join(app.config['UPLOAD_FOLDER'], project_id, file_path)
     if not os.path.exists(full_path):
         return jsonify({"error": "File not found"}), 404
+    
     try:
         with open(full_path, 'r', encoding='utf-8', errors='ignore') as f:
             content = f.read()
-        return jsonify({"path": file_path, "content": content})
+        return jsonify({"content": content})
     except Exception as e:
         return jsonify({"error": str(e)}), 500
 
-
 @app.route('/download')
 def download_project():
     project_id = session.get('project_id')
     if not project_id:
-        return "No project to download.", 404
+        return "No project to download. Please upload a project first.", 404
+    
     project_path = os.path.join(app.config['UPLOAD_FOLDER'], project_id)
+    if not os.path.exists(project_path):
+        return "Project files not found.", 404
+    
     zip_filename = f"project_{project_id}.zip"
     zip_path = os.path.join(app.config['UPLOAD_FOLDER'], zip_filename)
+    
+    # Remove existing zip if it exists
+    if os.path.exists(zip_path):
+        os.remove(zip_path)
+    
     with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
         for root, _, files in os.walk(project_path):
             for file in files:
                 file_path = os.path.join(root, file)
                 arcname = os.path.relpath(file_path, project_path)
                 zipf.write(file_path, arcname)
+    
     return send_from_directory(app.config['UPLOAD_FOLDER'], zip_filename, as_attachment=True)
 
-
 if __name__ == '__main__':
     app.run(host="0.0.0.0", port=7860)