Update app.py

app.py

@@ -1,19 +1,19 @@
 import os
-import zipfile
 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 env
+# Load environment variables
 load_dotenv()
 
 app = Flask(__name__)
-app.config["SECRET_KEY"] = os.getenv("FLASK_SECRET_KEY", "insecure")
-app.config["SESSION_TYPE"] = "filesystem"
-app.config["SESSION_FILE_DIR"] = './flask_session'
-app.config["SESSION_PERMANENT"] = False
+app.config['SECRET_KEY'] = os.getenv("FLASK_SECRET_KEY", "insecure-key")
+app.config['SESSION_TYPE'] = 'filesystem'
+app.config['SESSION_PERMANENT'] = False
+app.config['SESSION_FILE_DIR'] = './flask_session'
 Session(app)
 
 UPLOAD_FOLDER = 'uploads'
@@ -21,12 +21,15 @@ app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
 os.makedirs(UPLOAD_FOLDER, exist_ok=True)
 os.makedirs(app.config['SESSION_FILE_DIR'], exist_ok=True)
 
-# Configure Gemini
-gemini_api_key = os.getenv("GEMINI_API_KEY")
+# Gemini AI setup
 model = None
-if gemini_api_key:
-    genai.configure(api_key=gemini_api_key)
-    model = genai.GenerativeModel('gemini-2.0-flash')
+gemini_key = os.getenv("GEMINI_API_KEY")
+if gemini_key:
+    try:
+        genai.configure(api_key=gemini_key)
+        model = genai.GenerativeModel('gemini-2.0-flash')
+    except Exception as e:
+        print(f"Gemini setup failed: {e}")
 
 def get_project_files(project_path):
     files = {}
@@ -50,6 +53,10 @@ def index():
 
 @app.route('/upload', methods=['POST'])
 def upload():
+    if 'project_id' not in session:
+        session['project_id'] = str(uuid.uuid4())
+        session['chat_history'] = []
+
     if 'project_zip' not in request.files:
         return jsonify({'error': 'No file uploaded'}), 400
     file = request.files['project_zip']
@@ -58,6 +65,8 @@ def upload():
 
     project_id = session['project_id']
     project_path = os.path.join(app.config['UPLOAD_FOLDER'], project_id)
+
+    # Clear existing files
     if os.path.exists(project_path):
         for root, dirs, files in os.walk(project_path, topdown=False):
             for f in files:
@@ -66,6 +75,7 @@ def upload():
                 os.rmdir(os.path.join(root, d))
     os.makedirs(project_path, exist_ok=True)
 
+    # Save and extract zip
     zip_path = os.path.join(project_path, file.filename)
     file.save(zip_path)
     with zipfile.ZipFile(zip_path, 'r') as z:
@@ -73,14 +83,13 @@ def upload():
     os.remove(zip_path)
 
     files = get_project_files(project_path)
-
-    context = "User uploaded a project:\n"
+    context = "User uploaded project:\n"
     for path, content in files.items():
-        context += f"**File:** `{path}`\n```\n{content[:2000]}\n```\n\n"  # clip long content
+        context += f"**File:** `{path}`\n```\n{content[:2000]}\n```\n\n"
 
     session['chat_history'] = [
         {"role": "user", "parts": [context]},
-        {"role": "model", "parts": ["Project loaded. Ready to help."]}
+        {"role": "model", "parts": ["Project loaded. How can I help you?"]}
     ]
     session.modified = True
 
@@ -88,8 +97,8 @@ def upload():
         "message": "Project uploaded.",
         "file_tree": list(files.keys()),
         "chat_history": [
-            {"role": "user", "content": "Project context sent to AI."},
-            {"role": "assistant", "content": "Project loaded. Ready to help."}
+            {"role": "user", "content": "Project uploaded."},
+            {"role": "assistant", "content": "Project loaded. How can I help you?"}
         ]
     })
 
@@ -100,6 +109,10 @@ def chat():
     msg = request.json.get("message")
     if not msg:
         return jsonify({"error": "Empty message"}), 400
+
+    if 'project_id' not in session:
+        session['project_id'] = str(uuid.uuid4())
+        session['chat_history'] = []
     session['chat_history'].append({"role": "user", "parts": [msg]})
     try:
         chat = model.start_chat(history=session['chat_history'])
@@ -107,21 +120,23 @@ def chat():
         session['chat_history'].append({"role": "model", "parts": [reply]})
         return jsonify({"reply": reply})
     except Exception as e:
-        return jsonify({"error": str(e)}), 500
+        return jsonify({"error": str(e)})
 
 @app.route('/file_tree')
 def file_tree():
-    pid = session.get('project_id')
-    if not pid:
+    project_id = session.get('project_id')
+    if not project_id:
         return jsonify({"file_tree": []})
-    path = os.path.join(app.config['UPLOAD_FOLDER'], pid)
+    path = os.path.join(app.config['UPLOAD_FOLDER'], project_id)
     return jsonify({"file_tree": list(get_project_files(path).keys())})
 
 @app.route('/file_content')
 def file_content():
-    pid = session.get('project_id')
-    filepath = request.args.get('path')
-    full = os.path.join(app.config['UPLOAD_FOLDER'], pid, filepath)
+    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)
     try:
         with open(full, 'r', encoding='utf-8', errors='ignore') as f:
             return jsonify({"content": f.read()})
@@ -130,17 +145,17 @@ def file_content():
 
 @app.route('/download')
 def download():
-    pid = session.get('project_id')
-    if not pid:
+    project_id = session.get('project_id')
+    if not project_id:
         return "No project uploaded", 404
-    path = os.path.join(app.config['UPLOAD_FOLDER'], pid)
-    zpath = os.path.join(app.config['UPLOAD_FOLDER'], f"{pid}.zip")
-    with zipfile.ZipFile(zpath, 'w') as zipf:
+    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):
             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"{pid}.zip", as_attachment=True)
+    return send_from_directory(app.config['UPLOAD_FOLDER'], f"{project_id}.zip", as_attachment=True)
 
 if __name__ == "__main__":
     app.run(host='0.0.0.0', port=7860)