Update app.py

app.py

@@ -4,7 +4,11 @@ from parser import parse_python_code
 import os
 import json
 import io
-from database import init_chromadb, populate_sample_db, store_program, query_programs
+from database import init_chromadb, populate_sample_db, store_program, query_programs, load_chromadb_from_hf, HF_DATASET_NAME, HF_TOKEN
+
+# User-configurable variables
+DB_NAME = "python_programs"  # ChromaDB collection name
+UPLOAD_DIR = "./uploads"  # Directory for uploads
 
 app = Flask(__name__)
 
@@ -27,15 +31,15 @@ def index():
             if not file.filename.endswith('.py'):
                 return 'Invalid file type. Please upload a Python file.', 400
             filename = file.filename
-            file_path = os.path.join('uploads', filename)
+            file_path = os.path.join(UPLOAD_DIR, filename)
             file.save(file_path)
             with open(file_path, 'r') as f:
                 code_input = f.read()
             parts, sequence = parse_python_code(code_input)
             # Store in ChromaDB
-            vectors = [part['vector'] for part in parts]
             client = init_chromadb()
-            store_program(client, code_input, sequence, vectors)
+            vectors = [part['vector'] for part in parts]
+            store_program(client, code_input, sequence, vectors, DB_NAME)
         elif 'code' in request.form and request.form['code'].strip():
             code_input = request.form['code']
             filename = request.form.get('filename', 'unnamed.py') or 'unnamed.py'
@@ -44,12 +48,12 @@ def index():
             parts, sequence = parse_python_code(code_input)
             vectors = [part['vector'] for part in parts]
             client = init_chromadb()
-            store_program(client, code_input, sequence, vectors)
+            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
             operations = [op.strip() for op in request.form['query_ops'].split(',')]
-            client = init_chromadb()
-            query_results = query_programs(client, operations)
+            client = load_chromadb_from_hf(HF_DATASET_NAME, HF_TOKEN)  # Load from Hugging Face
+            query_results = query_programs(client, operations, DB_NAME)
             return render_template(
                 'results_partial.html',
                 parts=None,
@@ -73,8 +77,13 @@ def index():
         return 'No file, code, or query provided', 400
 
     # Initial page load
-    client = init_chromadb()  # Ensure ChromaDB is initialized
-    populate_sample_db(client)  # Populate with sample data
+    client = load_chromadb_from_hf(HF_DATASET_NAME, HF_TOKEN)  # Load from Hugging Face on startup
+    # If no dataset exists locally, populate with samples
+    try:
+        if not client.list_collections()[0].name == DB_NAME:
+            populate_sample_db(client)
+    except:
+        populate_sample_db(client)
     return render_template('index.html', parts=None, filename=None, reconstructed_code=None, code_input=None, query_results=None)
 
 @app.route('/export_json', methods=['POST'])
@@ -92,6 +101,6 @@ def export_json():
     )
 
 if __name__ == '__main__':
-    if not os.path.exists('uploads'):
-        os.makedirs('uploads')
+    if not os.path.exists(UPLOAD_DIR):
+        os.makedirs(UPLOAD_DIR)
     app.run(host="0.0.0.0", port=7860)