Update database.py

database.py

@@ -24,7 +24,7 @@ PERSIST_DIR = "./chroma_data"  # Directory for persistent storage (optional)
 USE_GPU = False  # Default to CPU, set to True for GPU if available
 
 def init_chromadb(persist_dir=PERSIST_DIR):
-    """Initialize ChromaDB client, optionally with persistent storage, with error handling."""
+    """Initialize ChromaDB client, optionally with persistent storage, with error handling and logging."""
     try:
         # Use persistent storage if directory exists, otherwise in-memory
         if os.path.exists(persist_dir):
@@ -39,17 +39,19 @@ def init_chromadb(persist_dir=PERSIST_DIR):
         raise
 
 def create_collection(client, collection_name=DB_NAME):
-    """Create or get a ChromaDB collection for Python programs, with error handling."""
+    """Create or get a ChromaDB collection for Python programs, with error handling and logging."""
     try:
         collection = client.get_or_create_collection(name=collection_name)
-        logger.info(f"Using ChromaDB collection: {collection_name}")
+        logger.info(f"Using ChromaDB collection: {collection_name}, contains {collection.count()} entries")
+        if collection is None or not hasattr(collection, 'add'):
+            raise ValueError("ChromaDB collection creation or access failed")
         return collection
     except Exception as e:
         logger.error(f"Error creating or getting collection {collection_name}: {e}")
         raise
 
 def store_program(client, code, sequence, vectors, collection_name=DB_NAME):
-    """Store a program in ChromaDB with its code, sequence, and vectors, with error handling."""
+    """Store a program in ChromaDB with its code, sequence, and vectors, with error handling and logging."""
     try:
         collection = create_collection(client, collection_name)
         
@@ -65,7 +67,7 @@ def store_program(client, code, sequence, vectors, collection_name=DB_NAME):
             ids=[program_id],
             embeddings=[flattened_vectors]  # Pass as 6D vector
         )
-        logger.info(f"Stored program in ChromaDB: {program_id}")
+        logger.info(f"Stored program in ChromaDB: {program_id}, total entries: {collection.count()}")
         return program_id
     except Exception as e:
         logger.error(f"Error storing program in ChromaDB: {e}")
@@ -93,13 +95,14 @@ def populate_sample_db(client):
             parts, sequence = parse_python_code(code)
             vectors = [part['vector'] for part in parts]
             store_program(client, code, sequence, vectors)
-        logger.info("Populated ChromaDB with sample programs")
+        collection = create_collection(client, DB_NAME)
+        logger.info(f"Populated ChromaDB with sample programs, total entries: {collection.count()}")
     except Exception as e:
         logger.error(f"Error populating sample database: {e}")
         raise
 
 def query_programs(client, operations, collection_name=DB_NAME, top_k=5, semantic_query=None):
-    """Query ChromaDB for programs matching the operations sequence or semantic description, with error handling."""
+    """Query ChromaDB for programs matching the operations sequence or semantic description, with error handling and logging."""
     try:
         collection = create_collection(client, collection_name)
         
@@ -139,7 +142,7 @@ def query_programs(client, operations, collection_name=DB_NAME, top_k=5, semanti
                 similarity = cosine_similarity([query_vector], [semantic_vector])[0][0] if semantic_vector and query_vector else 0
                 matching_programs.append({'id': meta['id'], 'code': doc, 'similarity': similarity, 'description': meta.get('description_tokens', ''), 'program_vectors': meta.get('program_vectors', '[]')})
         
-        logger.info(f"Queried {len(matching_programs)} programs from ChromaDB")
+        logger.info(f"Queried {len(matching_programs)} programs from ChromaDB, total entries: {collection.count()}")
         return sorted(matching_programs, key=lambda x: x['similarity'], reverse=True)
     except Exception as e:
         logger.error(f"Error querying programs from ChromaDB: {e}")
@@ -238,7 +241,7 @@ def generate_semantic_vector(description, total_lines=100, use_gpu=False):
     return vector
 
 def save_chromadb_to_hf(dataset_name=HF_DATASET_NAME, token=os.getenv("HF_KEY")):
-    """Save ChromaDB data to Hugging Face Dataset, with error handling."""
+    """Save ChromaDB data to Hugging Face Dataset, with error handling and logging."""
     try:
         client = init_chromadb()
         collection = client.get_collection(DB_NAME)
@@ -255,16 +258,19 @@ def save_chromadb_to_hf(dataset_name=HF_DATASET_NAME, token=os.getenv("HF_KEY"))
         
         # Create a Hugging Face Dataset
         dataset = Dataset.from_dict(data)
+        logger.info(f"Created Hugging Face Dataset with {len(data['code'])} entries")
         
         # Push to Hugging Face Hub
         dataset.push_to_hub(dataset_name, token=token)
         logger.info(f"Dataset pushed to Hugging Face Hub as {dataset_name}")
+        # Verify push (optional, could check dataset on Hub)
+        logger.info(f"Verified Hugging Face dataset push with {len(dataset)} entries")
     except Exception as e:
         logger.error(f"Error pushing dataset to Hugging Face Hub: {e}")
         raise
 
 def load_chromadb_from_hf(dataset_name=HF_DATASET_NAME, token=os.getenv("HF_KEY")):
-    """Load ChromaDB data from Hugging Face Dataset, handle empty dataset, with error handling."""
+    """Load ChromaDB data from Hugging Face Dataset, handle empty dataset, with error handling and logging."""
     try:
         dataset = load_dataset(dataset_name, split="train", token=token)
         client = init_chromadb()
@@ -272,15 +278,17 @@ def load_chromadb_from_hf(dataset_name=HF_DATASET_NAME, token=os.getenv("HF_KEY"
         
         for item in dataset:
             store_program(client, item["code"], item["sequence"].split(','), item["program_vectors"])
-        logger.info(f"Loaded {len(dataset)} entries from Hugging Face Hub into ChromaDB")
+        collection = create_collection(client, DB_NAME)
+        logger.info(f"Loaded {len(dataset)} entries from Hugging Face Hub into ChromaDB, total entries: {collection.count()}")
         return client
     except Exception as e:
         logger.error(f"Error loading dataset from Hugging Face: {e}")
         # Fallback: Create empty collection
         client = init_chromadb()
-        create_collection(client)
+        collection = create_collection(client)
+        logger.info(f"Created empty ChromaDB collection: {DB_NAME}, contains {collection.count()} entries")
         return client
 
 if __name__ == '__main__':
     client = load_chromadb_from_hf()
-    logger.info("Database initialized or loaded from Hugging Face Hub")
+    logger.info("Database initialized or loaded from Hugging Face Hub, contains {client.get_collection(DB_NAME).count()} entries")