Update app.py

app.py

@@ -4,10 +4,9 @@ from parser import parse_python_code
 import os
 import json
 import io
-from database import init_chromadb, populate_sample_db, store_program, query_programs, load_chromadb_from_hf, HF_DATASET_NAME, HF_TOKEN
+from database import init_chromadb, populate_sample_db, store_program, query_programs, load_chromadb_from_hf, HF_DATASET_NAME, HF_TOKEN, DB_NAME
 
 # User-configurable variables
-DB_NAME = "python_programs"  # ChromaDB collection name
 UPLOAD_DIR = "./uploads"  # Directory for uploads
 
 app = Flask(__name__)
@@ -50,7 +49,7 @@ def index():
             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
+            # Handle query for operations (category sequence)
             operations = [op.strip() for op in request.form['query_ops'].split(',')]
             client = load_chromadb_from_hf(HF_DATASET_NAME, HF_TOKEN)  # Load from Hugging Face
             query_results = query_programs(client, operations, DB_NAME)
@@ -62,6 +61,19 @@ def index():
                 code_input=None,
                 query_results=query_results
             )
+        elif 'semantic_query' in request.form and request.form['semantic_query'].strip():
+            # Handle semantic query (natural language description)
+            semantic_query = request.form['semantic_query']
+            client = load_chromadb_from_hf(HF_DATASET_NAME, HF_TOKEN)  # Load from Hugging Face
+            query_results = query_programs(client, None, DB_NAME, semantic_query=semantic_query)
+            return render_template(
+                'results_partial.html',
+                parts=None,
+                filename=filename,
+                reconstructed_code=None,
+                code_input=None,
+                query_results=query_results
+            )
 
         if parts:
             indexed_parts = [{'index': i + 1, **part} for i, part in enumerate(parts)]
@@ -89,7 +101,8 @@ def index():
 @app.route('/export_json', methods=['POST'])
 def export_json():
     parts = request.json.get('parts', [])
-    export_data = [{'vector': part['vector'], 'source': part['source']} for part in parts]
+    export_data = [{'vector': part['vector'], 'source': part['source'], 'description': generate_description_tokens([part['category']], [part['vector']])}
+                   for part in parts]
     json_str = json.dumps(export_data, indent=2)
     buffer = io.BytesIO(json_str.encode('utf-8'))
     buffer.seek(0)
@@ -100,6 +113,31 @@ def export_json():
         mimetype='application/json'
     )
 
+def generate_description_tokens(sequence, vectors):
+    """Generate semantic description tokens for a program based on its sequence and vectors."""
+    tokens = []
+    category_descriptions = {
+        'import': 'imports module',
+        'function': 'defines function',
+        'assigned_variable': 'assigns variable',
+        'input_variable': 'input parameter',
+        'returned_variable': 'returns value',
+        'if': 'conditional statement',
+        'return': 'returns result',
+        'try': 'try block',
+        'except': 'exception handler',
+        'expression': 'expression statement',
+        'spacer': 'empty line or comment'
+    }
+    
+    for cat, vec in zip(sequence, vectors):
+        if cat in category_descriptions:
+            tokens.append(f"{category_descriptions[cat]}:{cat}")
+            # Add vector-derived features (e.g., level, span) as tokens
+            tokens.append(f"level:{vec[1]}")
+            tokens.append(f"span:{vec[3]:.2f}")
+    return " ".join(tokens)
+
 if __name__ == '__main__':
     if not os.path.exists(UPLOAD_DIR):
         os.makedirs(UPLOAD_DIR)