Update app.py

app.py

@@ -5,7 +5,7 @@ import os
 import json
 import io
 import subprocess  # To call process_hf_dataset.py
-from database import init_chromadb, store_program, query_programs, load_chromadb_from_hf, DB_NAME
+from database import init_chromadb, store_program, query_programs, load_chromadb_from_hf, DB_NAME, create_collection
 import logging
 
 # Set up logging
@@ -46,6 +46,10 @@ def index():
                 vectors = [part['vector'] for part in parts]
                 store_program(client, code_input, sequence, vectors, DB_NAME)
                 logger.info(f"Stored code: {filename}")
+                # Verify storage
+                collection = create_collection(client, DB_NAME)
+                count = collection.count()
+                logger.info(f"ChromaDB now contains {count} entries")
             except Exception as e:
                 logger.error(f"Error storing code {filename}: {e}")
                 return f"Error storing code: {e}", 500
@@ -60,6 +64,10 @@ def index():
                 vectors = [part['vector'] for part in parts]
                 store_program(client, code_input, sequence, vectors, DB_NAME)
                 logger.info(f"Stored code: {filename}")
+                # Verify storage
+                collection = create_collection(client, DB_NAME)
+                count = collection.count()
+                logger.info(f"ChromaDB now contains {count} entries")
             except Exception as e:
                 logger.error(f"Error storing code {filename}: {e}")
                 return f"Error storing code: {e}", 500
@@ -70,6 +78,8 @@ def index():
                 client = load_chromadb_from_hf()
                 query_results = query_programs(client, operations, DB_NAME)
                 logger.info(f"Queried operations: {operations}")
+                # Verify query results
+                logger.info(f"Found {len(query_results)} matching programs in ChromaDB")
                 return render_template(
                     'results_partial.html',
                     parts=None,
@@ -88,6 +98,8 @@ def index():
                 client = load_chromadb_from_hf()
                 query_results = query_programs(client, None, DB_NAME, semantic_query=semantic_query)
                 logger.info(f"Queried semantically: {semantic_query}")
+                # Verify query results
+                logger.info(f"Found {len(query_results)} matching programs in ChromaDB")
                 return render_template(
                     'results_partial.html',
                     parts=None,
@@ -111,12 +123,20 @@ 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
+                if collection is None or not hasattr(collection, 'add'):
+                    raise ValueError("ChromaDB collection creation failed")
+                logger.info("Verified ChromaDB collection is valid")
                 
                 # Process dataset
                 result = subprocess.run(['python', 'process_hf_dataset.py'], check=True, capture_output=True, text=True, cwd=os.path.dirname(__file__))
                 logger.info(f"Process Hugging Face dataset output: {result.stdout}")
                 if result.stderr:
                     logger.error(f"Process Hugging Face dataset errors: {result.stderr}")
+                # Verify database population
+                collection = create_collection(client, DB_NAME)
+                count = collection.count()
+                logger.info(f"ChromaDB now contains {count} entries after processing")
                 return render_template(
                     'results_partial.html',
                     parts=None,
@@ -139,12 +159,20 @@ def index():
                 client = init_chromadb()
                 collection = client.get_or_create_collection(DB_NAME)
                 logger.info(f"Using existing or new ChromaDB collection: {DB_NAME}")
+                # Verify collection
+                if collection is None or not hasattr(collection, 'add'):
+                    raise ValueError("ChromaDB collection access failed")
+                logger.info("Verified ChromaDB collection is valid")
                 
                 # Process dataset
                 result = subprocess.run(['python', 'process_hf_dataset.py'], check=True, capture_output=True, text=True, cwd=os.path.dirname(__file__))
                 logger.info(f"Load Hugging Face dataset output: {result.stdout}")
                 if result.stderr:
                     logger.error(f"Load Hugging Face dataset errors: {result.stderr}")
+                # Verify database population
+                collection = create_collection(client, DB_NAME)
+                count = collection.count()
+                logger.info(f"ChromaDB now contains {count} entries after loading")
                 return render_template(
                     'results_partial.html',
                     parts=None,
@@ -175,6 +203,9 @@ def index():
                 if collection is None or not hasattr(collection, 'add'):
                     raise ValueError("ChromaDB collection creation failed")
                 logger.info("Verified ChromaDB collection is valid")
+                # Verify collection is empty
+                count = collection.count()
+                logger.info(f"ChromaDB now contains {count} entries after reset (should be 0)")
                 return render_template(
                     'results_partial.html',
                     parts=None,