Spaces:
Sleeping
Sleeping
It was added the milvus control status
Browse files
app.py
CHANGED
@@ -7,6 +7,8 @@ from chat_with_project import query_project
|
|
7 |
from get_prompts import get_prompt_for_mode
|
8 |
from dotenv import load_dotenv, set_key
|
9 |
from milvus import initialize_milvus, DEFAULT_MILVUS_HOST, DEFAULT_MILVUS_PORT, DEFAULT_COLLECTION_NAME, DEFAULT_DIMENSION, DEFAULT_MAX_RETRIES, DEFAULT_RETRY_DELAY
|
|
|
|
|
10 |
|
11 |
# --- Configuration and Setup ---
|
12 |
|
@@ -14,6 +16,10 @@ from milvus import initialize_milvus, DEFAULT_MILVUS_HOST, DEFAULT_MILVUS_PORT,
|
|
14 |
WORKSPACE_DIR = "workspace"
|
15 |
EXTRACTION_DIR = "extraction"
|
16 |
|
|
|
|
|
|
|
|
|
17 |
def clear_directories():
|
18 |
"""Clears the workspace and extraction directories."""
|
19 |
for directory in [WORKSPACE_DIR, EXTRACTION_DIR]:
|
@@ -21,23 +27,27 @@ def clear_directories():
|
|
21 |
shutil.rmtree(directory)
|
22 |
os.makedirs(directory, exist_ok=True)
|
23 |
|
|
|
24 |
# Clear directories at startup
|
25 |
clear_directories()
|
26 |
|
27 |
# --- API Key Management ---
|
28 |
|
|
|
29 |
def ensure_env_file_exists():
|
30 |
"""Ensures that a .env file exists in the project root."""
|
31 |
if not os.path.exists(".env"):
|
32 |
with open(".env", "w") as f:
|
33 |
f.write("") # Create an empty .env file
|
34 |
|
|
|
35 |
def load_api_key():
|
36 |
"""Loads the API key from the .env file or the environment."""
|
37 |
ensure_env_file_exists()
|
38 |
load_dotenv()
|
39 |
return os.environ.get("OPENAI_API_KEY")
|
40 |
|
|
|
41 |
def update_api_key(api_key):
|
42 |
"""Updates the API key in the .env file."""
|
43 |
if api_key:
|
@@ -47,18 +57,20 @@ def update_api_key(api_key):
|
|
47 |
else:
|
48 |
return "API key cannot be empty."
|
49 |
|
|
|
50 |
def is_api_key_set():
|
51 |
"""Checks if the API key is set."""
|
52 |
return bool(load_api_key())
|
53 |
|
54 |
# --- Core Functionalities ---
|
55 |
|
|
|
56 |
def process_zip(zip_file_path):
|
57 |
"""Extracts a zip file, analyzes content, and stores information."""
|
58 |
try:
|
59 |
# Clear existing workspace and extraction directories before processing
|
60 |
clear_directories()
|
61 |
-
|
62 |
# Extract the zip file
|
63 |
with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
|
64 |
zip_ref.extractall(WORKSPACE_DIR)
|
@@ -71,8 +83,12 @@ def process_zip(zip_file_path):
|
|
71 |
except Exception as e:
|
72 |
return f"An error occurred: {e}"
|
73 |
|
74 |
-
|
|
|
|
|
|
|
75 |
"""Initializes or loads the Milvus vector database."""
|
|
|
76 |
try:
|
77 |
# Convert string inputs to appropriate types
|
78 |
milvus_port = int(milvus_port)
|
@@ -80,14 +96,29 @@ def init_milvus(milvus_host, milvus_port, collection_name, dimension, max_retrie
|
|
80 |
max_retries = int(max_retries)
|
81 |
retry_delay = int(retry_delay)
|
82 |
|
83 |
-
|
84 |
-
|
85 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
86 |
except Exception as e:
|
|
|
87 |
return f"Error initializing Milvus: {e}"
|
88 |
|
89 |
# --- Chatbot Verification ---
|
90 |
|
|
|
91 |
def is_project_loaded():
|
92 |
"""Checks if a project has been loaded (i.e., if the extraction directory contains .pkl files)."""
|
93 |
extraction_dir = "extraction"
|
@@ -95,6 +126,8 @@ def is_project_loaded():
|
|
95 |
return bool(pkl_files)
|
96 |
|
97 |
# --- Helper Function for Developer Mode ---
|
|
|
|
|
98 |
def extract_files_from_response(response):
|
99 |
"""
|
100 |
Parses the LLM response to extract file paths and their corresponding code content.
|
@@ -114,7 +147,7 @@ def extract_files_from_response(response):
|
|
114 |
if current_file is not None:
|
115 |
# Save previous file content
|
116 |
files[current_file] = "\n".join(current_content)
|
117 |
-
|
118 |
# Start a new file
|
119 |
current_file = line.replace("--- BEGIN FILE:", "").strip()
|
120 |
current_content = []
|
@@ -137,10 +170,25 @@ def chat_ui(query, history, mode):
|
|
137 |
"""Handles the chat interaction for Analyzer, Debugger, and Developer modes."""
|
138 |
api_key = load_api_key()
|
139 |
if not api_key:
|
140 |
-
return [
|
|
|
|
|
|
|
|
|
|
|
141 |
|
142 |
if not is_project_loaded():
|
143 |
-
return [
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
144 |
|
145 |
# Initialize history if None
|
146 |
if history is None:
|
@@ -152,26 +200,35 @@ def chat_ui(query, history, mode):
|
|
152 |
|
153 |
# Pass the query and system prompt to the LLM
|
154 |
response = query_project(query, system_prompt)
|
|
|
155 |
print(f"Response from query_project: {response}")
|
156 |
|
157 |
if response is None or not response.strip():
|
158 |
response = "An error occurred during processing. Please check the logs."
|
159 |
|
|
|
|
|
|
|
|
|
160 |
if mode == "developer":
|
161 |
extracted_files = extract_files_from_response(response)
|
162 |
-
formatted_response = ""
|
163 |
|
|
|
164 |
for filepath, content in extracted_files.items():
|
165 |
-
formatted_response += f"## {filepath}\n`\n{content}\n`\n"
|
166 |
|
|
|
|
|
|
|
167 |
else:
|
168 |
# Format the output for non-developer modes
|
169 |
-
formatted_response = response
|
170 |
|
|
|
|
|
171 |
history.append((query, formatted_response))
|
172 |
|
173 |
return history, history
|
174 |
-
|
175 |
# ZIP Processing Interface
|
176 |
zip_iface = gr.Interface(
|
177 |
fn=process_zip,
|
@@ -185,12 +242,36 @@ zip_iface = gr.Interface(
|
|
185 |
milvus_iface = gr.Interface(
|
186 |
fn=init_milvus,
|
187 |
inputs=[
|
188 |
-
gr.Textbox(
|
189 |
-
|
190 |
-
|
191 |
-
|
192 |
-
|
193 |
-
gr.Textbox(
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
194 |
],
|
195 |
outputs="text",
|
196 |
title="Milvus Database Initialization",
|
@@ -203,26 +284,38 @@ chat_iface = gr.Interface(
|
|
203 |
inputs=[
|
204 |
gr.Textbox(label="Ask a question", placeholder="Type your question here"),
|
205 |
gr.State(), # Maintains chat history
|
206 |
-
gr.Radio(
|
|
|
|
|
|
|
|
|
207 |
],
|
208 |
outputs=[
|
209 |
-
gr.Chatbot(label="Chat with Project"),
|
210 |
-
"state"
|
211 |
],
|
212 |
title="Chat with your Project",
|
213 |
description="Ask questions about the data extracted from the zip file.",
|
214 |
# Example usage - Corrected to only include instruction and mode
|
215 |
examples=[
|
216 |
["What is this project about?", "analyzer"],
|
217 |
-
["Are there any potential bugs?", "debugger"],
|
218 |
["How does the data flow through the application?", "analyzer"],
|
219 |
["Explain the main components of the architecture.", "analyzer"],
|
220 |
-
["What are the dependencies of this project?",
|
221 |
-
["Are there any potential memory leaks?",
|
222 |
-
[
|
223 |
-
|
224 |
-
|
225 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
226 |
],
|
227 |
)
|
228 |
|
@@ -242,16 +335,26 @@ def get_api_key_status():
|
|
242 |
else:
|
243 |
return "API key status: Not set"
|
244 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
245 |
status_iface = gr.Interface(
|
246 |
-
fn=get_api_key_status,
|
247 |
inputs=None,
|
248 |
-
outputs="text",
|
249 |
live=True,
|
250 |
-
title="
|
251 |
)
|
252 |
|
253 |
# Add credits to the UI
|
254 |
-
credits = gr.Markdown(
|
|
|
|
|
255 |
|
256 |
# --- Main Application Launch ---
|
257 |
|
|
|
7 |
from get_prompts import get_prompt_for_mode
|
8 |
from dotenv import load_dotenv, set_key
|
9 |
from milvus import initialize_milvus, DEFAULT_MILVUS_HOST, DEFAULT_MILVUS_PORT, DEFAULT_COLLECTION_NAME, DEFAULT_DIMENSION, DEFAULT_MAX_RETRIES, DEFAULT_RETRY_DELAY
|
10 |
+
from pymilvus import connections, MilvusException, utility
|
11 |
+
import markdown # Import the markdown library
|
12 |
|
13 |
# --- Configuration and Setup ---
|
14 |
|
|
|
16 |
WORKSPACE_DIR = "workspace"
|
17 |
EXTRACTION_DIR = "extraction"
|
18 |
|
19 |
+
# Milvus connection status - Assume connected initially
|
20 |
+
milvus_connected = True
|
21 |
+
|
22 |
+
|
23 |
def clear_directories():
|
24 |
"""Clears the workspace and extraction directories."""
|
25 |
for directory in [WORKSPACE_DIR, EXTRACTION_DIR]:
|
|
|
27 |
shutil.rmtree(directory)
|
28 |
os.makedirs(directory, exist_ok=True)
|
29 |
|
30 |
+
|
31 |
# Clear directories at startup
|
32 |
clear_directories()
|
33 |
|
34 |
# --- API Key Management ---
|
35 |
|
36 |
+
|
37 |
def ensure_env_file_exists():
|
38 |
"""Ensures that a .env file exists in the project root."""
|
39 |
if not os.path.exists(".env"):
|
40 |
with open(".env", "w") as f:
|
41 |
f.write("") # Create an empty .env file
|
42 |
|
43 |
+
|
44 |
def load_api_key():
|
45 |
"""Loads the API key from the .env file or the environment."""
|
46 |
ensure_env_file_exists()
|
47 |
load_dotenv()
|
48 |
return os.environ.get("OPENAI_API_KEY")
|
49 |
|
50 |
+
|
51 |
def update_api_key(api_key):
|
52 |
"""Updates the API key in the .env file."""
|
53 |
if api_key:
|
|
|
57 |
else:
|
58 |
return "API key cannot be empty."
|
59 |
|
60 |
+
|
61 |
def is_api_key_set():
|
62 |
"""Checks if the API key is set."""
|
63 |
return bool(load_api_key())
|
64 |
|
65 |
# --- Core Functionalities ---
|
66 |
|
67 |
+
|
68 |
def process_zip(zip_file_path):
|
69 |
"""Extracts a zip file, analyzes content, and stores information."""
|
70 |
try:
|
71 |
# Clear existing workspace and extraction directories before processing
|
72 |
clear_directories()
|
73 |
+
|
74 |
# Extract the zip file
|
75 |
with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
|
76 |
zip_ref.extractall(WORKSPACE_DIR)
|
|
|
83 |
except Exception as e:
|
84 |
return f"An error occurred: {e}"
|
85 |
|
86 |
+
|
87 |
+
def init_milvus(
|
88 |
+
milvus_host, milvus_port, collection_name, dimension, max_retries, retry_delay
|
89 |
+
):
|
90 |
"""Initializes or loads the Milvus vector database."""
|
91 |
+
global milvus_connected
|
92 |
try:
|
93 |
# Convert string inputs to appropriate types
|
94 |
milvus_port = int(milvus_port)
|
|
|
96 |
max_retries = int(max_retries)
|
97 |
retry_delay = int(retry_delay)
|
98 |
|
99 |
+
# Call the modified function
|
100 |
+
success = initialize_milvus(
|
101 |
+
milvus_host,
|
102 |
+
milvus_port,
|
103 |
+
collection_name,
|
104 |
+
dimension,
|
105 |
+
max_retries,
|
106 |
+
retry_delay,
|
107 |
+
)
|
108 |
+
|
109 |
+
if success:
|
110 |
+
milvus_connected = True
|
111 |
+
return "Milvus database initialized or loaded successfully."
|
112 |
+
else:
|
113 |
+
milvus_connected = False
|
114 |
+
return "Error initializing Milvus: Unable to establish connection or initialize."
|
115 |
except Exception as e:
|
116 |
+
milvus_connected = False
|
117 |
return f"Error initializing Milvus: {e}"
|
118 |
|
119 |
# --- Chatbot Verification ---
|
120 |
|
121 |
+
|
122 |
def is_project_loaded():
|
123 |
"""Checks if a project has been loaded (i.e., if the extraction directory contains .pkl files)."""
|
124 |
extraction_dir = "extraction"
|
|
|
126 |
return bool(pkl_files)
|
127 |
|
128 |
# --- Helper Function for Developer Mode ---
|
129 |
+
|
130 |
+
|
131 |
def extract_files_from_response(response):
|
132 |
"""
|
133 |
Parses the LLM response to extract file paths and their corresponding code content.
|
|
|
147 |
if current_file is not None:
|
148 |
# Save previous file content
|
149 |
files[current_file] = "\n".join(current_content)
|
150 |
+
|
151 |
# Start a new file
|
152 |
current_file = line.replace("--- BEGIN FILE:", "").strip()
|
153 |
current_content = []
|
|
|
170 |
"""Handles the chat interaction for Analyzer, Debugger, and Developer modes."""
|
171 |
api_key = load_api_key()
|
172 |
if not api_key:
|
173 |
+
return [
|
174 |
+
(
|
175 |
+
"Error",
|
176 |
+
"OpenAI API key not set. Please set the API key in the Settings tab.",
|
177 |
+
)
|
178 |
+
], []
|
179 |
|
180 |
if not is_project_loaded():
|
181 |
+
return [
|
182 |
+
(
|
183 |
+
"Error",
|
184 |
+
"No project loaded. Please upload and process a ZIP file first.",
|
185 |
+
)
|
186 |
+
], []
|
187 |
+
|
188 |
+
if not milvus_connected:
|
189 |
+
return [
|
190 |
+
("Error", "Milvus is not connected. Please connect to Milvus first.")
|
191 |
+
], []
|
192 |
|
193 |
# Initialize history if None
|
194 |
if history is None:
|
|
|
200 |
|
201 |
# Pass the query and system prompt to the LLM
|
202 |
response = query_project(query, system_prompt)
|
203 |
+
print(f"Type response {type(response)}")
|
204 |
print(f"Response from query_project: {response}")
|
205 |
|
206 |
if response is None or not response.strip():
|
207 |
response = "An error occurred during processing. Please check the logs."
|
208 |
|
209 |
+
# Convert the response to Markdown format
|
210 |
+
md = markdown.Markdown(extensions=['fenced_code'])
|
211 |
+
formatted_response = "" # Initialize formatted_response here
|
212 |
+
|
213 |
if mode == "developer":
|
214 |
extracted_files = extract_files_from_response(response)
|
|
|
215 |
|
216 |
+
# Iterate through the extracted files and format them
|
217 |
for filepath, content in extracted_files.items():
|
218 |
+
formatted_response += f"## {filepath}\n`\n{content}\n`\n\n"
|
219 |
|
220 |
+
# Remove the last extra newline characters if they exist
|
221 |
+
if formatted_response.endswith("\n\n"):
|
222 |
+
formatted_response = formatted_response[:-2]
|
223 |
else:
|
224 |
# Format the output for non-developer modes
|
225 |
+
formatted_response = response
|
226 |
|
227 |
+
# Convert the entire response to HTML at once
|
228 |
+
formatted_response = md.convert(formatted_response)
|
229 |
history.append((query, formatted_response))
|
230 |
|
231 |
return history, history
|
|
|
232 |
# ZIP Processing Interface
|
233 |
zip_iface = gr.Interface(
|
234 |
fn=process_zip,
|
|
|
242 |
milvus_iface = gr.Interface(
|
243 |
fn=init_milvus,
|
244 |
inputs=[
|
245 |
+
gr.Textbox(
|
246 |
+
label="Milvus Host",
|
247 |
+
placeholder=DEFAULT_MILVUS_HOST,
|
248 |
+
value=DEFAULT_MILVUS_HOST,
|
249 |
+
),
|
250 |
+
gr.Textbox(
|
251 |
+
label="Milvus Port",
|
252 |
+
placeholder=DEFAULT_MILVUS_PORT,
|
253 |
+
value=DEFAULT_MILVUS_PORT,
|
254 |
+
),
|
255 |
+
gr.Textbox(
|
256 |
+
label="Collection Name",
|
257 |
+
placeholder=DEFAULT_COLLECTION_NAME,
|
258 |
+
value=DEFAULT_COLLECTION_NAME,
|
259 |
+
),
|
260 |
+
gr.Textbox(
|
261 |
+
label="Dimension",
|
262 |
+
placeholder=str(DEFAULT_DIMENSION),
|
263 |
+
value=str(DEFAULT_DIMENSION),
|
264 |
+
),
|
265 |
+
gr.Textbox(
|
266 |
+
label="Max Retries",
|
267 |
+
placeholder=str(DEFAULT_MAX_RETRIES),
|
268 |
+
value=str(DEFAULT_MAX_RETRIES),
|
269 |
+
),
|
270 |
+
gr.Textbox(
|
271 |
+
label="Retry Delay (seconds)",
|
272 |
+
placeholder=str(DEFAULT_RETRY_DELAY),
|
273 |
+
value=str(DEFAULT_RETRY_DELAY),
|
274 |
+
),
|
275 |
],
|
276 |
outputs="text",
|
277 |
title="Milvus Database Initialization",
|
|
|
284 |
inputs=[
|
285 |
gr.Textbox(label="Ask a question", placeholder="Type your question here"),
|
286 |
gr.State(), # Maintains chat history
|
287 |
+
gr.Radio(
|
288 |
+
["analyzer", "debugger", "developer"],
|
289 |
+
label="Chat Mode",
|
290 |
+
value="analyzer",
|
291 |
+
),
|
292 |
],
|
293 |
outputs=[
|
294 |
+
gr.Chatbot(label="Chat with Project", height=500),
|
295 |
+
"state", # This is to store the state,
|
296 |
],
|
297 |
title="Chat with your Project",
|
298 |
description="Ask questions about the data extracted from the zip file.",
|
299 |
# Example usage - Corrected to only include instruction and mode
|
300 |
examples=[
|
301 |
["What is this project about?", "analyzer"],
|
302 |
+
["Are there any potential bugs?", "debugger"],
|
303 |
["How does the data flow through the application?", "analyzer"],
|
304 |
["Explain the main components of the architecture.", "analyzer"],
|
305 |
+
["What are the dependencies of this project?", "analyzer"],
|
306 |
+
["Are there any potential memory leaks?", "debugger"],
|
307 |
+
[
|
308 |
+
"Identify any areas where the code could be optimized.",
|
309 |
+
"debugger",
|
310 |
+
],
|
311 |
+
[
|
312 |
+
"Please implement basic logging for the main application and save logs to a file.",
|
313 |
+
"developer",
|
314 |
+
],
|
315 |
+
[
|
316 |
+
"Can you add a try/except blocks in main functions to handle exceptions",
|
317 |
+
"developer",
|
318 |
+
],
|
319 |
],
|
320 |
)
|
321 |
|
|
|
335 |
else:
|
336 |
return "API key status: Not set"
|
337 |
|
338 |
+
|
339 |
+
def get_milvus_status():
|
340 |
+
if milvus_connected:
|
341 |
+
return "Milvus: Connected"
|
342 |
+
else:
|
343 |
+
return "Milvus: Disconnected"
|
344 |
+
|
345 |
+
|
346 |
status_iface = gr.Interface(
|
347 |
+
fn=lambda: [get_api_key_status(), get_milvus_status()],
|
348 |
inputs=None,
|
349 |
+
outputs=["text", "text"],
|
350 |
live=True,
|
351 |
+
title="Status",
|
352 |
)
|
353 |
|
354 |
# Add credits to the UI
|
355 |
+
credits = gr.Markdown(
|
356 |
+
"## Credits\n\nCreated by [Ruslan Magana Vsevolodovna](https://ruslanmv.com/)"
|
357 |
+
)
|
358 |
|
359 |
# --- Main Application Launch ---
|
360 |
|
milvus.py
CHANGED
@@ -86,12 +86,22 @@ def connect_to_milvus(host, port, max_retries, retry_delay):
|
|
86 |
return False
|
87 |
|
88 |
def initialize_milvus(host, port, collection_name, dimension, max_retries, retry_delay):
|
89 |
-
"""Initializes Milvus with parameters.
|
90 |
-
if connect_to_milvus(host, port, max_retries, retry_delay):
|
91 |
-
create_milvus_collection(host, port, collection_name, dimension)
|
92 |
-
load_data_to_milvus(host, port, collection_name)
|
93 |
-
connections.disconnect(alias='default')
|
94 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
95 |
|
96 |
if __name__ == "__main__":
|
97 |
# Use default values or environment variables if available
|
|
|
86 |
return False
|
87 |
|
88 |
def initialize_milvus(host, port, collection_name, dimension, max_retries, retry_delay):
|
89 |
+
"""Initializes Milvus with parameters.
|
|
|
|
|
|
|
|
|
90 |
|
91 |
+
Returns:
|
92 |
+
True if successfully connected and initialized, False otherwise.
|
93 |
+
"""
|
94 |
+
if connect_to_milvus(host, port, max_retries, retry_delay):
|
95 |
+
try:
|
96 |
+
create_milvus_collection(host, port, collection_name, dimension)
|
97 |
+
load_data_to_milvus(host, port, collection_name)
|
98 |
+
connections.disconnect(alias='default')
|
99 |
+
return True # Return True if everything is successful
|
100 |
+
except Exception as e:
|
101 |
+
print(f"Error during initialization: {e}")
|
102 |
+
return False # Return False if any error occurs during collection creation or data loading
|
103 |
+
else:
|
104 |
+
return False # Return False if connection failed
|
105 |
|
106 |
if __name__ == "__main__":
|
107 |
# Use default values or environment variables if available
|