Update database.py

database.py

@@ -4,12 +4,12 @@ from parser import parse_python_code
 import os
 from sklearn.metrics.pairwise import cosine_similarity
 import numpy as np
-import dotenv
+from datasets import Dataset, load_dataset
 
 # User-configurable variables
 DB_NAME = "python_programs"  # ChromaDB collection name
-HF_DATASET_NAME = "broadfield-dev/python_program_vectors"  # Hugging Face Dataset name
-HF_TOKEN = os.getenv("HF_KEY")  # Replace with your Hugging Face API token
+HF_DATASET_NAME = "python_program_vectors"  # Hugging Face Dataset name
+HF_TOKEN = "YOUR_HUGGINGFACE_TOKEN"  # Replace with your Hugging Face API token
 PERSIST_DIR = "./chroma_data"  # Directory for persistent storage (optional)
 
 def init_chromadb(persist_dir=PERSIST_DIR):
@@ -44,7 +44,7 @@ def store_program(client, code, sequence, vectors, collection_name=DB_NAME):
     program_id = str(hash(code))  # Use hash of code as ID for uniqueness
     collection.add(
         documents=[code],
-        metadatas=[{"sequence": ",".join(sequence)}],
+        metadatas=[{"sequence": ",".join(sequence), "description_tokens": " ".join(generate_description_tokens(sequence, vectors))}],
         ids=[program_id],
         embeddings=[flattened_vectors]  # Pass as flat list
     )
@@ -72,33 +72,39 @@ def populate_sample_db(client):
         vectors = [part['vector'] for part in parts]
         store_program(client, code, sequence, vectors)
 
-def query_programs(client, operations, collection_name=DB_NAME, top_k=5):
-    """Query ChromaDB for programs matching the operations sequence."""
+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."""
     collection = create_collection(client, collection_name)
     
-    # Convert operations to a query vector (average of operation vectors)
-    query_vector = sum([create_vector(op, 0, (1, 1), 100, []) for op in operations], []) / len(operations) if operations else [0] * 6
-    
-    # Perform similarity search
-    results = collection.query(
-        query_embeddings=[query_vector],
-        n_results=top_k,
-        include=["documents", "metadatas"]
-    )
+    if semantic_query:
+        # Semantic search using description tokens
+        query_vector = generate_semantic_vector(semantic_query)
+        results = collection.query(
+            query_texts=[semantic_query],
+            n_results=top_k,
+            include=["documents", "metadatas"]
+        )
+    else:
+        # Vector-based search for operations sequence
+        query_vector = sum([create_vector(op, 0, (1, 1), 100, []) for op in operations], []) / len(operations) if operations else [0] * 6
+        results = collection.query(
+            query_embeddings=[query_vector],
+            n_results=top_k,
+            include=["documents", "metadatas"]
+        )
     
     # Process results
     matching_programs = []
     for doc, meta in zip(results['documents'][0], results['metadatas'][0]):
         sequence = meta['sequence'].split(',')
-        if is_subsequence(operations, sequence):
-            # Extract and flatten vectors from the document (assuming stored as string or list)
+        if not semantic_query or is_subsequence(operations, sequence):  # Ensure sequence match for operations
             try:
                 doc_vectors = eval(doc['vectors']) if isinstance(doc['vectors'], str) else doc['vectors']
                 program_vector = np.mean([v for v in doc_vectors if isinstance(v, (list, np.ndarray))], axis=0).tolist()
             except:
                 program_vector = [0] * 6  # Fallback for malformed vectors
             similarity = cosine_similarity([query_vector], [program_vector])[0][0] if program_vector and query_vector else 0
-            matching_programs.append({'id': meta['id'], 'code': doc, 'similarity': similarity})
+            matching_programs.append({'id': meta['id'], 'code': doc, 'similarity': similarity, 'description': meta.get('description_tokens', '')})
     
     return sorted(matching_programs, key=lambda x: x['similarity'], reverse=True)
 
@@ -124,9 +130,52 @@ def is_subsequence(subseq, seq):
     it = iter(seq)
     return all(item in it for item in subseq)
 
+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 tokens
+
+def generate_semantic_vector(description):
+    """Generate a semantic vector for a textual description (simplified for now)."""
+    # This is a placeholder—use an embedding model (e.g., CodeBERT, BERT) for real semantic search
+    tokens = description.lower().split()
+    category_weights = {
+        'import': 1, 'function': 2, 'assign': 17, 'input': 18, 'return': 19, 'if': 5, 'try': 8, 'except': 14
+    }
+    vector = [0] * 6
+    for token in tokens:
+        for cat, weight in category_weights.items():
+            if cat in token:
+                vector[0] = weight  # Use category_id as primary feature
+                vector[1] = 1  # Assume level 1 for simplicity
+                vector[2] = 0.5  # Center position (midpoint)
+                vector[3] = 0.1  # Span (small for simplicity)
+                vector[4] = 1  # Parent depth (shallow)
+                vector[5] = weight / len(category_weights)  # Parent weight
+    return vector
+
 def save_chromadb_to_hf(dataset_name=HF_DATASET_NAME, token=HF_TOKEN):
     """Save ChromaDB data to Hugging Face Dataset."""
-    from datasets import Dataset
     client = init_chromadb()
     collection = create_collection(client)
     
@@ -135,7 +184,8 @@ def save_chromadb_to_hf(dataset_name=HF_DATASET_NAME, token=HF_TOKEN):
     data = {
         "code": results["documents"],
         "sequence": [meta["sequence"] for meta in results["metadatas"]],
-        "vectors": [[item for sublist in vec for item in sublist] for vec in results["embeddings"]]  # Flatten vectors
+        "vectors": [[item for sublist in vec for item in sublist] for vec in results["embeddings"]],  # Flatten vectors
+        "description_tokens": [meta.get('description_tokens', '') for meta in results["metadatas"]]
     }
     
     # Create a Hugging Face Dataset
@@ -146,23 +196,29 @@ def save_chromadb_to_hf(dataset_name=HF_DATASET_NAME, token=HF_TOKEN):
     print(f"Dataset pushed to Hugging Face Hub as {dataset_name}")
 
 def load_chromadb_from_hf(dataset_name=HF_DATASET_NAME, token=HF_TOKEN):
-    """Load ChromaDB data from Hugging Face Dataset."""
-    from datasets import load_dataset
+    """Load ChromaDB data from Hugging Face Dataset, handle empty dataset."""
+    try:
+        dataset = load_dataset(dataset_name, split="train", token=token)
+    except Exception as e:
+        print(f"Error loading dataset from Hugging Face: {e}. Populating with samples...")
+        client = init_chromadb()
+        populate_sample_db(client)
+        save_chromadb_to_hf()  # Create and push a new dataset
+        return init_chromadb()
+    
     client = init_chromadb()
     collection = create_collection(client)
     
-    dataset = load_dataset(dataset_name, split="train", token=token)
     for item in dataset:
         collection.add(
             documents=[item["code"]],
-            metadatas=[{"sequence": item["sequence"]}],
+            metadatas=[{"sequence": item["sequence"], "description_tokens": item["description_tokens"]}],
             ids=[str(hash(item["code"]))],
             embeddings=[item["vectors"]]
         )
     return client
 
 if __name__ == '__main__':
-    client = init_chromadb()
-    populate_sample_db(client)
+    client = load_chromadb_from_hf()
     # Uncomment to save to Hugging Face
     # save_chromadb_to_hf()