Update app.py

app.py

@@ -5,9 +5,9 @@ import os
 import json
 import io
 import subprocess  # To call process_hf_dataset.py
+from database import init_chromadb, populate_sample_db, store_program, query_programs, load_chromadb_from_hf, DB_NAME
 
 # User-configurable variables
-DB_NAME = "python_programs"  # ChromaDB collection name
 UPLOAD_DIR = "./uploads"  # Directory for uploads
 
 app = Flask(__name__)
@@ -37,7 +37,6 @@ def index():
                 code_input = f.read()
             parts, sequence = parse_python_code(code_input)
             # Store in ChromaDB
-            from database import init_chromadb, store_program
             client = init_chromadb()
             vectors = [part['vector'] for part in parts]
             store_program(client, code_input, sequence, vectors, DB_NAME)
@@ -48,13 +47,11 @@ def index():
                 filename += '.py'
             parts, sequence = parse_python_code(code_input)
             vectors = [part['vector'] for part in parts]
-            from database import init_chromadb, store_program
             client = init_chromadb()
             store_program(client, code_input, sequence, vectors, DB_NAME)
         elif 'query_ops' in request.form and request.form['query_ops'].strip():
             # Handle query for operations (category sequence)
             operations = [op.strip() for op in request.form['query_ops'].split(',')]
-            from database import load_chromadb_from_hf, query_programs
             client = load_chromadb_from_hf()
             query_results = query_programs(client, operations, DB_NAME)
             return render_template(
@@ -68,7 +65,6 @@ def index():
         elif 'semantic_query' in request.form and request.form['semantic_query'].strip():
             # Handle semantic query (natural language description)
             semantic_query = request.form['semantic_query']
-            from database import load_chromadb_from_hf, query_programs
             client = load_chromadb_from_hf()
             query_results = query_programs(client, None, DB_NAME, semantic_query=semantic_query)
             return render_template(
@@ -80,8 +76,17 @@ def index():
                 query_results=query_results
             )
         elif 'process_hf' in request.form:
-            # Trigger processing of Hugging Face dataset
+            # Trigger processing of Hugging Face dataset with fresh database
             try:
+                # Reset ChromaDB collection
+                client = init_chromadb()
+                try:
+                    client.delete_collection(DB_NAME)
+                except:
+                    pass  # Collection may not exist
+                client.create_collection(DB_NAME)
+                
+                # Process dataset
                 subprocess.run(['python', 'process_hf_dataset.py'], check=True)
                 return render_template(
                     'results_partial.html',
@@ -90,10 +95,31 @@ def index():
                     reconstructed_code=None,
                     code_input=None,
                     query_results=None,
-                    message="Hugging Face dataset processed and stored successfully."
+                    message="Hugging Face dataset processed and stored successfully with fresh database."
                 )
             except subprocess.CalledProcessError as e:
                 return f"Error processing Hugging Face dataset: {e}", 500
+        elif 'reset_db' in request.form:
+            # Reset ChromaDB collection
+            try:
+                client = init_chromadb()
+                try:
+                    client.delete_collection(DB_NAME)
+                except:
+                    pass  # Collection may not exist
+                client.create_collection(DB_NAME)
+                populate_sample_db(client)  # Optionally repopulate with samples
+                return render_template(
+                    'results_partial.html',
+                    parts=None,
+                    filename="Database Reset",
+                    reconstructed_code=None,
+                    code_input=None,
+                    query_results=None,
+                    message="Database reset and repopulated with samples successfully."
+                )
+            except Exception as e:
+                return f"Error resetting database: {e}", 500
 
         if parts:
             indexed_parts = [{'index': i + 1, **part} for i, part in enumerate(parts)]
@@ -109,15 +135,12 @@ def index():
         return 'No file, code, or query provided', 400
 
     # Initial page load
-    from database import load_chromadb_from_hf
     client = load_chromadb_from_hf()
     # If no dataset exists locally, populate with samples
     try:
         if not client.list_collections()[0].name == DB_NAME:
-            from database import populate_sample_db
             populate_sample_db(client)
     except:
-        from database import populate_sample_db
         populate_sample_db(client)
     return render_template('index.html', parts=None, filename=None, reconstructed_code=None, code_input=None, query_results=None)