File size: 9,840 Bytes
8c1a9c9
9759c17
 
 
 
 
 
d9e48e0
 
9759c17
 
61bb0ed
d9e48e0
 
 
 
 
 
 
 
 
9759c17
 
 
 
 
 
 
 
d9e48e0
 
9759c17
d9e48e0
9759c17
 
 
d9e48e0
 
 
9759c17
 
 
 
 
 
 
 
 
8c1a9c9
d9e48e0
 
8a7ffa1
8c1a9c9
d9e48e0
9759c17
 
 
8c1a9c9
 
d9e48e0
9759c17
 
d9e48e0
 
 
 
 
 
9759c17
 
 
 
 
d9e48e0
9759c17
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
d9e48e0
9759c17
 
 
 
 
d9e48e0
9759c17
 
 
 
d9e48e0
9759c17
 
 
d9e48e0
9759c17
 
 
d9e48e0
 
 
9759c17
 
8c1a9c9
d9e48e0
 
 
 
9759c17
 
 
 
 
 
 
d9e48e0
9759c17
 
8a7ffa1
9759c17
d9e48e0
9759c17
d9e48e0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9759c17
8c1a9c9
9759c17
 
 
8c1a9c9
d9e48e0
 
9759c17
 
 
 
 
 
d9e48e0
 
 
 
9759c17
 
 
 
 
 
d9e48e0
9759c17
d9e48e0
 
9759c17
 
d9e48e0
 
 
 
9759c17
 
 
 
 
 
 
 
 
d9e48e0
9759c17
 
d9e48e0
9759c17
 
d9e48e0
9759c17
 
 
 
 
 
d9e48e0
9759c17
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8c1a9c9
d9e48e0
9759c17
 
d9e48e0
9759c17
 
d9e48e0
 
9759c17
 
 
 
 
 
 
 
d9e48e0
9759c17
8c1a9c9
 
a487cc2
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
import os
import zipfile
import uuid
import google.generativeai as genai
from flask import Flask, request, jsonify, render_template, session, send_from_directory
from dotenv import load_dotenv

# --- API and App Configuration ---
load_dotenv() # Load environment variables from .env file

# Initialize the Flask app
app = Flask(__name__)
app.secret_key = os.environ.get("FLASK_SECRET_KEY", 'your_default_secret_key_change_me')

# Configure the Gemini API client
try:
    genai.configure(api_key=os.environ.get("GEMINI_API_KEY"))
    model = genai.GenerativeModel('gemini-1.5-flash') # Or 'gemini-pro'
except Exception as e:
    print(f"Error configuring Gemini API: {e}")
    model = None

# Configure the upload folder
UPLOAD_FOLDER = 'uploads'
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

if not os.path.exists(UPLOAD_FOLDER):
    os.makedirs(UPLOAD_FOLDER)


# --- Helper Functions ---
def get_project_files(project_path):
    """Gets the file tree and content of all files in the project directory."""
    file_data = {}
    for root, _, files in os.walk(project_path):
        for file in files:
            # Skip .DS_Store and other unwanted files
            if file == '.DS_Store':
                continue
            file_path = os.path.join(root, file)
            relative_path = os.path.relpath(file_path, project_path)
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    content = f.read()
                file_data[relative_path] = content
            except Exception as e:
                file_data[relative_path] = f"Error reading file: {e}"
    return file_data


# --- Flask Routes ---
@app.route('/')
def index():
    """Renders the main page and initializes the session."""
    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():
    """Handles the project zip file upload."""
    # FIX: Ensure session is initialized to prevent KeyError
    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 part"}), 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)
        if not os.path.exists(project_path):
            os.makedirs(project_path)

        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)
        session['project_files'] = project_files

        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"

        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?"]}
        ]

        frontend_chat_history = [
            {"role": "user", "content": initial_context},
            {"role": "assistant", "content": "I have analyzed the project. I am ready to help. How can I assist you with your code?"}
        ]

        return jsonify({
            "message": "Project uploaded and analyzed.",
            "file_tree": list(project_files.keys()),
            "chat_history": frontend_chat_history
        })

    return jsonify({"error": "Invalid file type. Please upload a .zip file."}), 400


@app.route('/chat', methods=['POST'])
def chat():
    """Handles chat interaction with the Gemini API."""
    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": "Empty message"}), 400

    project_id = session.get('project_id')
    project_path = os.path.join(app.config['UPLOAD_FOLDER'], project_id)

    # Add user message to history
    session['chat_history'].append({"role": "user", "parts": [user_message]})
    session.modified = True

    chat_session = model.start_chat(history=session['chat_history'])
    
    try:
        system_instruction = """You are an expert programmer AI assistant. You must follow these rules:
1.  When asked to modify a file, you MUST respond with the full file path on the first line, followed by the complete, updated code for that file inside a markdown code block. Example:
    `src/app.js`
    ```javascript
    // new updated javascript code
    ```
2.  When asked to create a new file, respond with the new file path and its content in the same format.
3.  When asked to delete a file, respond ONLY with 'DELETE: [file_path]'.
4.  For general conversation, just respond naturally."""
        
        # The history already provides context. We send the latest user message.
        # The system instruction is part of the model's configuration or initial prompt.
        # For simplicity here, we prepend it to guide this specific turn.
        # A more advanced setup might use the system_instruction parameter in genai.configure
        prompt_with_instruction = f"{system_instruction}\n\n## Project Context:\n{get_project_files(project_path)}\n\n## Conversation History:\n{session['chat_history']}\n\n## User's Latest Message:\n{user_message}"
        
        response = chat_session.send_message(user_message) # Send only the new message
        ai_response = response.text

        # Add AI response to chat history
        session['chat_history'].append({"role": "model", "parts": [ai_response]})
        session.modified = True

        # --- AI File Operation Logic ---
        # A more robust solution might involve structured JSON output from the LLM
        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:]
                    # Extract language and content
                    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 IndexError:
                    pass # Ignore malformed code blocks for now

        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})

    except Exception as e:
        return jsonify({"error": str(e)}), 500


@app.route('/file_tree')
def get_file_tree():
    """Returns the current file structure of the project."""
    project_id = session.get('project_id')
    if not project_id:
        return jsonify({"error": "No project in session."}), 400
    project_path = os.path.join(app.config['UPLOAD_FOLDER'], project_id)
    if not os.path.exists(project_path):
        return jsonify({"error": "Project files not found."}), 404
    file_tree = list(get_project_files(project_path).keys())
    return jsonify({"file_tree": file_tree})


@app.route('/file_content')
def get_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

    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') as f:
            content = f.read()
        return jsonify({"path": file_path, "content": content})
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@app.route('/download')
def download_project():
    """Zips and sends the current project state for download."""
    project_id = session.get('project_id')
    if not project_id:
        return "No project to download.", 404

    project_path = os.path.join(app.config['UPLOAD_FOLDER'], project_id)
    zip_filename = f"project_{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):
            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)