Update app.py

app.py

@@ -40,50 +40,65 @@ def index():
             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
-            client = init_chromadb()
-            vectors = [part['vector'] for part in parts]
-            store_program(client, code_input, sequence, vectors, DB_NAME)
-            logger.info(f"Stored code: {filename}")
+            try:
+                parts, sequence = parse_python_code(code_input)
+                client = init_chromadb()
+                vectors = [part['vector'] for part in parts]
+                store_program(client, code_input, sequence, vectors, DB_NAME)
+                logger.info(f"Stored code: {filename}")
+            except Exception as e:
+                logger.error(f"Error storing code {filename}: {e}")
+                return f"Error storing code: {e}", 500
         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'
             if not filename.endswith('.py'):
                 filename += '.py'
-            parts, sequence = parse_python_code(code_input)
-            vectors = [part['vector'] for part in parts]
-            client = init_chromadb()
-            store_program(client, code_input, sequence, vectors, DB_NAME)
-            logger.info(f"Stored code: {filename}")
+            try:
+                parts, sequence = parse_python_code(code_input)
+                client = init_chromadb()
+                vectors = [part['vector'] for part in parts]
+                store_program(client, code_input, sequence, vectors, DB_NAME)
+                logger.info(f"Stored code: {filename}")
+            except Exception as e:
+                logger.error(f"Error storing code {filename}: {e}")
+                return f"Error storing code: {e}", 500
         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(',')]
-            client = load_chromadb_from_hf()
-            query_results = query_programs(client, operations, DB_NAME)
-            logger.info(f"Queried operations: {operations}")
-            return render_template(
-                'results_partial.html',
-                parts=None,
-                filename=filename,
-                reconstructed_code=None,
-                code_input=None,
-                query_results=query_results
-            )
+            try:
+                client = load_chromadb_from_hf()
+                query_results = query_programs(client, operations, DB_NAME)
+                logger.info(f"Queried operations: {operations}")
+                return render_template(
+                    'results_partial.html',
+                    parts=None,
+                    filename=filename,
+                    reconstructed_code=None,
+                    code_input=None,
+                    query_results=query_results
+                )
+            except Exception as e:
+                logger.error(f"Error querying operations: {e}")
+                return f"Error querying operations: {e}", 500
         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()
-            query_results = query_programs(client, None, DB_NAME, semantic_query=semantic_query)
-            logger.info(f"Queried semantically: {semantic_query}")
-            return render_template(
-                'results_partial.html',
-                parts=None,
-                filename=filename,
-                reconstructed_code=None,
-                code_input=None,
-                query_results=query_results
-            )
+            try:
+                client = load_chromadb_from_hf()
+                query_results = query_programs(client, None, DB_NAME, semantic_query=semantic_query)
+                logger.info(f"Queried semantically: {semantic_query}")
+                return render_template(
+                    'results_partial.html',
+                    parts=None,
+                    filename=filename,
+                    reconstructed_code=None,
+                    code_input=None,
+                    query_results=query_results
+                )
+            except Exception as e:
+                logger.error(f"Error querying semantically: {e}")
+                return f"Error querying semantically: {e}", 500
         elif 'process_hf' in request.form:
             # Trigger processing of Hugging Face dataset with fresh database
             try:
@@ -114,6 +129,9 @@ def index():
             except subprocess.CalledProcessError as e:
                 logger.error(f"Error processing Hugging Face dataset: {e.stderr}")
                 return f"Error processing Hugging Face dataset: {e.stderr}", 500
+            except Exception as e:
+                logger.error(f"Unexpected error processing Hugging Face dataset: {e}")
+                return f"Unexpected error processing Hugging Face dataset: {e}", 500
         elif 'load_dataset' in request.form:
             # Trigger loading of Hugging Face dataset without resetting
             try:
@@ -139,6 +157,9 @@ def index():
             except subprocess.CalledProcessError as e:
                 logger.error(f"Error loading Hugging Face dataset: {e.stderr}")
                 return f"Error loading Hugging Face dataset: {e.stderr}", 500
+            except Exception as e:
+                logger.error(f"Unexpected error loading Hugging Face dataset: {e}")
+                return f"Unexpected error loading Hugging Face dataset: {e}", 500
         elif 'reset_db' in request.form:
             # Reset ChromaDB collection (no repopulation with samples)
             try:
@@ -150,9 +171,10 @@ def index():
                     logger.warning(f"Failed to delete collection {DB_NAME}: {e}")
                 collection = client.create_collection(DB_NAME)
                 logger.info(f"Created fresh ChromaDB collection: {DB_NAME}")
-                # Verify collection creation by checking if it's iterable (fix for 'NoneType' error)
+                # Verify collection creation
                 if collection is None or not hasattr(collection, 'add'):
                     raise ValueError("ChromaDB collection creation failed")
+                logger.info("Verified ChromaDB collection is valid")
                 return render_template(
                     'results_partial.html',
                     parts=None,