Update app.py

app.py

@@ -1,244 +1,146 @@
 import os
 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
 from dotenv import load_dotenv
+import google.generativeai as genai
 
-# Load environment variables
+# Load env
 load_dotenv()
 
-# Initialize Flask app
 app = Flask(__name__)
-app.config["SECRET_KEY"] = os.environ.get("FLASK_SECRET_KEY", 'default-insecure-secret-key')
-
-# Configure server-side sessions
-app.config["SESSION_PERMANENT"] = False
+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_FILE_DIR"] = './flask_session'
+app.config["SESSION_PERMANENT"] = False
 Session(app)
 
-# Configure Gemini AI
-try:
-    gemini_api_key = os.environ.get("GEMINI_API_KEY")
-    if not gemini_api_key:
-        print("GEMINI_API_KEY environment variable not found.")
-        model = None
-    else:
-        genai.configure(api_key=gemini_api_key)
-        model = genai.GenerativeModel('gemini-2.0-flash')
-except Exception as e:
-    print(f"Error configuring Gemini API: {e}")
-    model = None
-
-# Configure upload folder
 UPLOAD_FOLDER = 'uploads'
 app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
 os.makedirs(UPLOAD_FOLDER, exist_ok=True)
-os.makedirs('./flask_session', exist_ok=True)
+os.makedirs(app.config['SESSION_FILE_DIR'], exist_ok=True)
+
+# Configure Gemini
+gemini_api_key = os.getenv("GEMINI_API_KEY")
+model = None
+if gemini_api_key:
+    genai.configure(api_key=gemini_api_key)
+    model = genai.GenerativeModel('gemini-2.0-flash')
 
 def get_project_files(project_path):
-    """Get all files in a project directory with their contents."""
-    file_data = {}
-    for root, _, files in os.walk(project_path):
-        for file in files:
-            if file in ['.DS_Store', '__MACOSX']:
-                continue
-            file_path = os.path.join(root, file)
-            relative_path = os.path.relpath(file_path, project_path)
+    files = {}
+    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(file_path, 'r', encoding='utf-8', errors='ignore') as f:
-                    content = f.read()
-                file_data[relative_path] = content
+                with open(os.path.join(root, file), 'r', encoding='utf-8', errors='ignore') as f:
+                    files[rel_path] = f.read()
             except Exception as e:
-                file_data[relative_path] = f"Error reading file: {e}"
-    return file_data
+                files[rel_path] = f"Error reading: {e}"
+    return files
 
 @app.route('/')
 def index():
-    """Render the main application page."""
     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_project():
-    """Handle project zip file upload."""
-    if 'project_id' not in session:
-        session['project_id'] = str(uuid.uuid4())
-        session['chat_history'] = []
-
+def upload():
     if 'project_zip' not in request.files:
-        return jsonify({"error": "No file part"}), 400
-    
+        return jsonify({'error': 'No file uploaded'}), 400
     file = request.files['project_zip']
-    if file.filename == '':
-        return jsonify({"error": "No selected file"}), 400
-
-    if file and file.filename.endswith('.zip'):
-        project_id = session['project_id']
-        project_path = os.path.join(app.config['UPLOAD_FOLDER'], project_id)
-        
-        # Clear existing project files
-        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)
-
-        # Prepare initial context for AI
-        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"
-
-        # Update session with chat history
-        session['chat_history'] = [
-            {"role": "user", "parts": [initial_context]},
-            {"role": "model", "parts": ["I have analyzed the project. I am ready to help. How can I assist you with your code?"]}
+    if not file.filename.endswith('.zip'):
+        return jsonify({'error': 'Only .zip supported'}), 400
+
+    project_id = session['project_id']
+    project_path = os.path.join(app.config['UPLOAD_FOLDER'], project_id)
+    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)
+
+    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)
+
+    files = get_project_files(project_path)
+
+    context = "User uploaded a project:\n"
+    for path, content in files.items():
+        context += f"**File:** `{path}`\n```\n{content[:2000]}\n```\n\n"  # clip long content
+
+    session['chat_history'] = [
+        {"role": "user", "parts": [context]},
+        {"role": "model", "parts": ["Project loaded. Ready to help."]}
+    ]
+    session.modified = True
+
+    return jsonify({
+        "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."}
         ]
-        session.modified = True
-
-        # Prepare response for frontend
-        return jsonify({
-            "message": "Project uploaded and analyzed.",
-            "file_tree": list(project_files.keys()),
-            "chat_history": [
-                {"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?"}
-            ]
-        })
-    return jsonify({"error": "Invalid file type. Please upload a .zip file."}), 400
+    })
 
 @app.route('/chat', methods=['POST'])
 def chat():
-    """Handle chat messages with the AI."""
     if not model:
-        return jsonify({"error": "Gemini API is not configured. Please check your API key."}), 500
-
-    user_message = request.json.get('message')
-    if not user_message:
+        return jsonify({"error": "Gemini API not configured"}), 500
+    msg = request.json.get("message")
+    if not msg:
         return jsonify({"error": "Empty message"}), 400
-
-    project_id = session.get('project_id')
-    if not project_id:
-        return jsonify({"error": "Session expired. Please upload your project again."}), 400
-
-    project_path = os.path.join(app.config['UPLOAD_FOLDER'], project_id)
-    session['chat_history'].append({"role": "user", "parts": [user_message]})
-    session.modified = True
-
+    session['chat_history'].append({"role": "user", "parts": [msg]})
     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()
-            if file_path_line.startswith('`') and file_path_line.endswith('`'):
-                file_path = file_path_line.strip('`')
-                try:
-                    content_start = ai_response.find('```')
-                    code_block = ai_response[content_start:]
-                    code_content = '\n'.join(code_block.split('\n')[1:-1])
-                    full_path = os.path.join(project_path, file_path)
-                    os.makedirs(os.path.dirname(full_path), exist_ok=True)
-                    with open(full_path, 'w', encoding='utf-8') as f:
-                        f.write(code_content)
-                except Exception:
-                    pass
-        elif ai_response.strip().startswith("DELETE:"):
-            file_to_delete = ai_response.strip().split(":", 1)[1].strip()
-            full_path = os.path.join(project_path, file_to_delete)
-            if os.path.exists(full_path):
-                try:
-                    os.remove(full_path)
-                except OSError as e:
-                    print(f"Error deleting file {full_path}: {e}")
-
-        return jsonify({"reply": ai_response})
+        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})
     except Exception as e:
         return jsonify({"error": str(e)}), 500
 
 @app.route('/file_tree')
-def get_file_tree():
-    """Get the current project's file tree."""
-    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": []})
-    
-    try:
-        project_files = get_project_files(project_path)
-        return jsonify({"file_tree": list(project_files.keys())})
-    except Exception as e:
+def file_tree():
+    pid = session.get('project_id')
+    if not pid:
         return jsonify({"file_tree": []})
+    path = os.path.join(app.config['UPLOAD_FOLDER'], pid)
+    return jsonify({"file_tree": list(get_project_files(path).keys())})
 
 @app.route('/file_content')
-def get_file_content():
-    """Get 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": "No file selected"}), 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
-    
+def file_content():
+    pid = session.get('project_id')
+    filepath = request.args.get('path')
+    full = os.path.join(app.config['UPLOAD_FOLDER'], pid, filepath)
     try:
-        with open(full_path, 'r', encoding='utf-8', errors='ignore') as f:
-            content = f.read()
-        return jsonify({"content": content})
+        with open(full, 'r', encoding='utf-8', errors='ignore') as f:
+            return jsonify({"content": f.read()})
     except Exception as e:
-        return jsonify({"error": str(e)}), 500
+        return jsonify({"error": str(e)})
 
 @app.route('/download')
-def download_project():
-    """Download the current project as a zip file."""
-    project_id = session.get('project_id')
-    if not project_id:
-        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)
-    
-    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):
+def download():
+    pid = session.get('project_id')
+    if not pid:
+        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:
+        for root, _, files in os.walk(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)
+                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)
+
+if __name__ == "__main__":
+    app.run(host='0.0.0.0', port=7860)