Athspi commited on
Commit
d9e48e0
·
verified ·
1 Parent(s): a672c5c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +81 -75
app.py CHANGED
@@ -3,19 +3,22 @@ import zipfile
3
  import uuid
4
  import google.generativeai as genai
5
  from flask import Flask, request, jsonify, render_template, session, send_from_directory
6
-
7
- # --- Gemini API Configuration ---
8
- # Load the API key from the .env file
9
  from dotenv import load_dotenv
10
- load_dotenv()
11
 
12
- # Configure the Gemini API client
13
- genai.configure(api_key=os.environ.get("GEMINI_API_KEY"))
14
- model = genai.GenerativeModel('gemini-1.5-flash') # Or another suitable model
15
 
16
  # Initialize the Flask app
17
  app = Flask(__name__)
18
- app.secret_key = 'your_very_secret_key' # Change this to a random secret key
 
 
 
 
 
 
 
 
19
 
20
  # Configure the upload folder
21
  UPLOAD_FOLDER = 'uploads'
@@ -24,13 +27,16 @@ app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
24
  if not os.path.exists(UPLOAD_FOLDER):
25
  os.makedirs(UPLOAD_FOLDER)
26
 
 
 
27
  def get_project_files(project_path):
28
- """
29
- Gets the file tree and content of all files in the project directory.
30
- """
31
  file_data = {}
32
  for root, _, files in os.walk(project_path):
33
  for file in files:
 
 
 
34
  file_path = os.path.join(root, file)
35
  relative_path = os.path.relpath(file_path, project_path)
36
  try:
@@ -41,26 +47,31 @@ def get_project_files(project_path):
41
  file_data[relative_path] = f"Error reading file: {e}"
42
  return file_data
43
 
 
 
44
  @app.route('/')
45
  def index():
46
- """
47
- Renders the main page.
48
- """
49
  if 'project_id' not in session:
50
  session['project_id'] = str(uuid.uuid4())
51
  session['chat_history'] = []
52
  return render_template('index.html')
53
 
 
54
  @app.route('/upload', methods=['POST'])
55
  def upload_project():
56
- """
57
- Handles the project zip file upload.
58
- """
 
 
 
59
  if 'project_zip' not in request.files:
60
  return jsonify({"error": "No file part"}), 400
61
  file = request.files['project_zip']
62
  if file.filename == '':
63
  return jsonify({"error": "No selected file"}), 400
 
64
  if file and file.filename.endswith('.zip'):
65
  project_id = session['project_id']
66
  project_path = os.path.join(app.config['UPLOAD_FOLDER'], project_id)
@@ -77,102 +88,100 @@ def upload_project():
77
  project_files = get_project_files(project_path)
78
  session['project_files'] = project_files
79
 
80
- # Initial context for the AI
81
- initial_context = "The user has uploaded a new project. Here are the files and their content:\n\n"
82
  for path, content in project_files.items():
83
  initial_context += f"**File:** `{path}`\n```\n{content}\n```\n\n"
84
 
85
- # Initialize chat history for Gemini
86
  session['chat_history'] = [
87
  {"role": "user", "parts": [initial_context]},
88
- {"role": "model", "parts": ["I have understood the project structure and content. How can I help you improve it?"]}
89
  ]
90
 
91
- # Prepare response for the frontend (needs 'role' and 'content' keys)
92
  frontend_chat_history = [
93
  {"role": "user", "content": initial_context},
94
- {"role": "assistant", "content": "I have understood the project structure and content. How can I help you improve it?"}
95
  ]
96
 
97
  return jsonify({
98
- "message": "Project uploaded and understood.",
99
  "file_tree": list(project_files.keys()),
100
  "chat_history": frontend_chat_history
101
  })
102
- return jsonify({"error": "Invalid file type, please upload a .zip file"}), 400
 
 
103
 
104
  @app.route('/chat', methods=['POST'])
105
  def chat():
106
- """
107
- Handles the chat interaction with the AI.
108
- """
 
109
  user_message = request.json.get('message')
110
  if not user_message:
111
  return jsonify({"error": "Empty message"}), 400
112
 
113
  project_id = session.get('project_id')
114
- if not project_id:
115
- return jsonify({"error": "No project uploaded"}), 400
116
-
117
  project_path = os.path.join(app.config['UPLOAD_FOLDER'], project_id)
118
 
119
- # Add user message to chat history
120
  session['chat_history'].append({"role": "user", "parts": [user_message]})
121
  session.modified = True
122
 
123
- # Start a chat session with the model using the history
124
  chat_session = model.start_chat(history=session['chat_history'])
125
-
126
  try:
127
- # System instruction for Gemini
128
- system_instruction = """You are an expert programmer AI assistant. When the user asks you to make changes to the code, you must respond with the file path and the complete, updated code for that file within ```python ... ```, ```html ... ```, or other appropriate markdown code blocks. For creating a new file, respond with the new file path and the content. For deleting a file, respond with 'DELETE: [file_path]'. For suggesting improvements with diffs, provide the original and suggested code in separate blocks.
129
- Example for modifying a file:
130
- `app/main.py`
131
- ```python
132
- # updated python code here
133
- ```
134
- Example for creating a file:
135
- `static/css/new_style.css`
136
- ```css
137
- /* new css code */
138
- ```
139
- """
140
- # We combine the system instruction with the user message for better context
141
- full_prompt = f"{system_instruction}\n\nUser Question: {user_message}"
142
-
143
- response = chat_session.send_message(full_prompt)
144
  ai_response = response.text
145
 
146
  # Add AI response to chat history
147
  session['chat_history'].append({"role": "model", "parts": [ai_response]})
148
  session.modified = True
149
 
150
-
151
- # --- AI File Operation Logic (remains the same) ---
152
  if "```" in ai_response:
153
  lines = ai_response.split('\n')
154
  file_path_line = lines[0].strip()
155
  if file_path_line.startswith('`') and file_path_line.endswith('`'):
156
  file_path = file_path_line.strip('`')
157
  try:
158
- content_start = ai_response.find('```') + 3
159
- content_end = ai_response.rfind('```')
160
- # Skip the language hint (e.g., python, html)
161
- code_content = ai_response[content_start:].split('\n', 1)[1]
162
- code_content = code_content[:content_end - content_start - len(lines[1]) - 1]
163
 
164
  full_path = os.path.join(project_path, file_path)
165
  os.makedirs(os.path.dirname(full_path), exist_ok=True)
166
  with open(full_path, 'w', encoding='utf-8') as f:
167
  f.write(code_content)
168
  except IndexError:
169
- pass # Or handle error where code block is malformed
170
 
171
- elif ai_response.startswith("DELETE:"):
172
- file_to_delete = ai_response.split(":", 1)[1].strip()
173
  full_path = os.path.join(project_path, file_to_delete)
174
  if os.path.exists(full_path):
175
- os.remove(full_path)
 
 
 
176
 
177
  return jsonify({"reply": ai_response})
178
 
@@ -180,23 +189,22 @@ def chat():
180
  return jsonify({"error": str(e)}), 500
181
 
182
 
183
- # --- Other routes (/file_tree, /file_content, /download) remain unchanged ---
184
- # (You can copy them directly from the previous response)
185
-
186
  @app.route('/file_tree')
187
  def get_file_tree():
 
188
  project_id = session.get('project_id')
189
  if not project_id:
190
- return jsonify({"error": "No project uploaded"}), 400
191
  project_path = os.path.join(app.config['UPLOAD_FOLDER'], project_id)
192
  if not os.path.exists(project_path):
193
- return jsonify({"error": "Project not found"}), 404
194
  file_tree = list(get_project_files(project_path).keys())
195
  return jsonify({"file_tree": file_tree})
196
 
197
 
198
  @app.route('/file_content')
199
  def get_file_content():
 
200
  project_id = session.get('project_id')
201
  file_path = request.args.get('path')
202
  if not project_id or not file_path:
@@ -205,7 +213,6 @@ def get_file_content():
205
  full_path = os.path.join(app.config['UPLOAD_FOLDER'], project_id, file_path)
206
  if not os.path.exists(full_path):
207
  return jsonify({"error": "File not found"}), 404
208
-
209
  try:
210
  with open(full_path, 'r', encoding='utf-8') as f:
211
  content = f.read()
@@ -216,15 +223,14 @@ def get_file_content():
216
 
217
  @app.route('/download')
218
  def download_project():
219
- """
220
- Zips and sends the current project state for download.
221
- """
222
  project_id = session.get('project_id')
223
  if not project_id:
224
- return "No project to download", 404
225
 
226
  project_path = os.path.join(app.config['UPLOAD_FOLDER'], project_id)
227
- zip_path = os.path.join(app.config['UPLOAD_FOLDER'], f"{project_id}.zip")
 
228
 
229
  with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
230
  for root, _, files in os.walk(project_path):
@@ -233,8 +239,8 @@ def download_project():
233
  arcname = os.path.relpath(file_path, project_path)
234
  zipf.write(file_path, arcname)
235
 
236
- return send_from_directory(app.config['UPLOAD_FOLDER'], f"{project_id}.zip", as_attachment=True)
237
 
238
 
239
  if __name__ == '__main__':
240
- app.run(host="0.0.0.0", port=7860)
 
3
  import uuid
4
  import google.generativeai as genai
5
  from flask import Flask, request, jsonify, render_template, session, send_from_directory
 
 
 
6
  from dotenv import load_dotenv
 
7
 
8
+ # --- API and App Configuration ---
9
+ load_dotenv() # Load environment variables from .env file
 
10
 
11
  # Initialize the Flask app
12
  app = Flask(__name__)
13
+ app.secret_key = os.environ.get("FLASK_SECRET_KEY", 'your_default_secret_key_change_me')
14
+
15
+ # Configure the Gemini API client
16
+ try:
17
+ genai.configure(api_key=os.environ.get("GEMINI_API_KEY"))
18
+ model = genai.GenerativeModel('gemini-1.5-flash') # Or 'gemini-pro'
19
+ except Exception as e:
20
+ print(f"Error configuring Gemini API: {e}")
21
+ model = None
22
 
23
  # Configure the upload folder
24
  UPLOAD_FOLDER = 'uploads'
 
27
  if not os.path.exists(UPLOAD_FOLDER):
28
  os.makedirs(UPLOAD_FOLDER)
29
 
30
+
31
+ # --- Helper Functions ---
32
  def get_project_files(project_path):
33
+ """Gets the file tree and content of all files in the project directory."""
 
 
34
  file_data = {}
35
  for root, _, files in os.walk(project_path):
36
  for file in files:
37
+ # Skip .DS_Store and other unwanted files
38
+ if file == '.DS_Store':
39
+ continue
40
  file_path = os.path.join(root, file)
41
  relative_path = os.path.relpath(file_path, project_path)
42
  try:
 
47
  file_data[relative_path] = f"Error reading file: {e}"
48
  return file_data
49
 
50
+
51
+ # --- Flask Routes ---
52
  @app.route('/')
53
  def index():
54
+ """Renders the main page and initializes the session."""
 
 
55
  if 'project_id' not in session:
56
  session['project_id'] = str(uuid.uuid4())
57
  session['chat_history'] = []
58
  return render_template('index.html')
59
 
60
+
61
  @app.route('/upload', methods=['POST'])
62
  def upload_project():
63
+ """Handles the project zip file upload."""
64
+ # FIX: Ensure session is initialized to prevent KeyError
65
+ if 'project_id' not in session:
66
+ session['project_id'] = str(uuid.uuid4())
67
+ session['chat_history'] = []
68
+
69
  if 'project_zip' not in request.files:
70
  return jsonify({"error": "No file part"}), 400
71
  file = request.files['project_zip']
72
  if file.filename == '':
73
  return jsonify({"error": "No selected file"}), 400
74
+
75
  if file and file.filename.endswith('.zip'):
76
  project_id = session['project_id']
77
  project_path = os.path.join(app.config['UPLOAD_FOLDER'], project_id)
 
88
  project_files = get_project_files(project_path)
89
  session['project_files'] = project_files
90
 
91
+ initial_context = "The user has uploaded a new project. Here is the file structure and content:\n\n"
 
92
  for path, content in project_files.items():
93
  initial_context += f"**File:** `{path}`\n```\n{content}\n```\n\n"
94
 
 
95
  session['chat_history'] = [
96
  {"role": "user", "parts": [initial_context]},
97
+ {"role": "model", "parts": ["I have analyzed the project. I am ready to help. How can I assist you with your code?"]}
98
  ]
99
 
 
100
  frontend_chat_history = [
101
  {"role": "user", "content": initial_context},
102
+ {"role": "assistant", "content": "I have analyzed the project. I am ready to help. How can I assist you with your code?"}
103
  ]
104
 
105
  return jsonify({
106
+ "message": "Project uploaded and analyzed.",
107
  "file_tree": list(project_files.keys()),
108
  "chat_history": frontend_chat_history
109
  })
110
+
111
+ return jsonify({"error": "Invalid file type. Please upload a .zip file."}), 400
112
+
113
 
114
  @app.route('/chat', methods=['POST'])
115
  def chat():
116
+ """Handles chat interaction with the Gemini API."""
117
+ if not model:
118
+ return jsonify({"error": "Gemini API is not configured. Please check your API key."}), 500
119
+
120
  user_message = request.json.get('message')
121
  if not user_message:
122
  return jsonify({"error": "Empty message"}), 400
123
 
124
  project_id = session.get('project_id')
 
 
 
125
  project_path = os.path.join(app.config['UPLOAD_FOLDER'], project_id)
126
 
127
+ # Add user message to history
128
  session['chat_history'].append({"role": "user", "parts": [user_message]})
129
  session.modified = True
130
 
 
131
  chat_session = model.start_chat(history=session['chat_history'])
132
+
133
  try:
134
+ system_instruction = """You are an expert programmer AI assistant. You must follow these rules:
135
+ 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:
136
+ `src/app.js`
137
+ ```javascript
138
+ // new updated javascript code
139
+ ```
140
+ 2. When asked to create a new file, respond with the new file path and its content in the same format.
141
+ 3. When asked to delete a file, respond ONLY with 'DELETE: [file_path]'.
142
+ 4. For general conversation, just respond naturally."""
143
+
144
+ # The history already provides context. We send the latest user message.
145
+ # The system instruction is part of the model's configuration or initial prompt.
146
+ # For simplicity here, we prepend it to guide this specific turn.
147
+ # A more advanced setup might use the system_instruction parameter in genai.configure
148
+ 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}"
149
+
150
+ response = chat_session.send_message(user_message) # Send only the new message
151
  ai_response = response.text
152
 
153
  # Add AI response to chat history
154
  session['chat_history'].append({"role": "model", "parts": [ai_response]})
155
  session.modified = True
156
 
157
+ # --- AI File Operation Logic ---
158
+ # A more robust solution might involve structured JSON output from the LLM
159
  if "```" in ai_response:
160
  lines = ai_response.split('\n')
161
  file_path_line = lines[0].strip()
162
  if file_path_line.startswith('`') and file_path_line.endswith('`'):
163
  file_path = file_path_line.strip('`')
164
  try:
165
+ content_start = ai_response.find('```')
166
+ code_block = ai_response[content_start:]
167
+ # Extract language and content
168
+ code_content = '\n'.join(code_block.split('\n')[1:-1])
 
169
 
170
  full_path = os.path.join(project_path, file_path)
171
  os.makedirs(os.path.dirname(full_path), exist_ok=True)
172
  with open(full_path, 'w', encoding='utf-8') as f:
173
  f.write(code_content)
174
  except IndexError:
175
+ pass # Ignore malformed code blocks for now
176
 
177
+ elif ai_response.strip().startswith("DELETE:"):
178
+ file_to_delete = ai_response.strip().split(":", 1)[1].strip()
179
  full_path = os.path.join(project_path, file_to_delete)
180
  if os.path.exists(full_path):
181
+ try:
182
+ os.remove(full_path)
183
+ except OSError as e:
184
+ print(f"Error deleting file {full_path}: {e}")
185
 
186
  return jsonify({"reply": ai_response})
187
 
 
189
  return jsonify({"error": str(e)}), 500
190
 
191
 
 
 
 
192
  @app.route('/file_tree')
193
  def get_file_tree():
194
+ """Returns the current file structure of the project."""
195
  project_id = session.get('project_id')
196
  if not project_id:
197
+ return jsonify({"error": "No project in session."}), 400
198
  project_path = os.path.join(app.config['UPLOAD_FOLDER'], project_id)
199
  if not os.path.exists(project_path):
200
+ return jsonify({"error": "Project files not found."}), 404
201
  file_tree = list(get_project_files(project_path).keys())
202
  return jsonify({"file_tree": file_tree})
203
 
204
 
205
  @app.route('/file_content')
206
  def get_file_content():
207
+ """Returns the content of a specific file."""
208
  project_id = session.get('project_id')
209
  file_path = request.args.get('path')
210
  if not project_id or not file_path:
 
213
  full_path = os.path.join(app.config['UPLOAD_FOLDER'], project_id, file_path)
214
  if not os.path.exists(full_path):
215
  return jsonify({"error": "File not found"}), 404
 
216
  try:
217
  with open(full_path, 'r', encoding='utf-8') as f:
218
  content = f.read()
 
223
 
224
  @app.route('/download')
225
  def download_project():
226
+ """Zips and sends the current project state for download."""
 
 
227
  project_id = session.get('project_id')
228
  if not project_id:
229
+ return "No project to download.", 404
230
 
231
  project_path = os.path.join(app.config['UPLOAD_FOLDER'], project_id)
232
+ zip_filename = f"project_{project_id}.zip"
233
+ zip_path = os.path.join(app.config['UPLOAD_FOLDER'], zip_filename)
234
 
235
  with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
236
  for root, _, files in os.walk(project_path):
 
239
  arcname = os.path.relpath(file_path, project_path)
240
  zipf.write(file_path, arcname)
241
 
242
+ return send_from_directory(app.config['UPLOAD_FOLDER'], zip_filename, as_attachment=True)
243
 
244
 
245
  if __name__ == '__main__':
246
+ app.run(debug=True, port=5001)