# database.py import chromadb from parser import parse_python_code import os from sklearn.metrics.pairwise import cosine_similarity import numpy as np # User-configurable variables DB_NAME = "python_programs" # ChromaDB collection name 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): """Initialize ChromaDB client, optionally with persistent storage.""" try: # Use persistent storage if directory exists, otherwise in-memory if os.path.exists(persist_dir): client = chromadb.PersistentClient(path=persist_dir) else: client = chromadb.Client() return client except Exception as e: print(f"Error initializing ChromaDB: {e}") return chromadb.Client() # Fallback to in-memory def create_collection(client, collection_name=DB_NAME): """Create or get a ChromaDB collection for Python programs.""" try: collection = client.get_collection(name=collection_name) except: collection = client.create_collection(name=collection_name) return collection def store_program(client, code, sequence, vectors, collection_name=DB_NAME): """Store a program in ChromaDB with its code, sequence, and vectors.""" collection = create_collection(client, collection_name) # Flatten vectors to ensure they are a list of numbers flattened_vectors = [item for sublist in vectors for item in sublist] # Store program data (ID, code, sequence, vectors) program_id = str(hash(code)) # Use hash of code as ID for uniqueness collection.add( documents=[code], metadatas=[{"sequence": ",".join(sequence)}], ids=[program_id], embeddings=[flattened_vectors] # Pass as flat list ) return program_id def populate_sample_db(client): """Populate ChromaDB with sample Python programs.""" samples = [ """ import os def add_one(x): y = x + 1 return y """, """ def multiply(a, b): c = a * b if c > 0: return c """ ] for code in samples: parts, sequence = parse_python_code(code) 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.""" 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"] ) # 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) 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}) return sorted(matching_programs, key=lambda x: x['similarity'], reverse=True) def create_vector(category, level, location, total_lines, parent_path): """Helper to create a vector for query (matches parser's create_vector).""" category_map = { 'import': 1, 'function': 2, 'async_function': 3, 'class': 4, 'if': 5, 'while': 6, 'for': 7, 'try': 8, 'expression': 9, 'spacer': 10, 'other': 11, 'elif': 12, 'else': 13, 'except': 14, 'finally': 15, 'return': 16, 'assigned_variable': 17, 'input_variable': 18, 'returned_variable': 19 } category_id = category_map.get(category, 0) start_line, end_line = location span = (end_line - start_line + 1) / total_lines center_pos = ((start_line + end_line) / 2) / total_lines parent_depth = len(parent_path) parent_weight = sum(category_map.get(parent.split('[')[0].lower(), 0) * (1 / (i + 1)) for i, parent in enumerate(parent_path)) / max(1, len(category_map)) return [category_id, level, center_pos, span, parent_depth, parent_weight] def is_subsequence(subseq, seq): """Check if subseq is a subsequence of seq.""" it = iter(seq) return all(item in it for item in subseq) 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) # Fetch all data from ChromaDB results = collection.get(include=["documents", "metadatas", "embeddings"]) 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 } # Create a Hugging Face Dataset dataset = Dataset.from_dict(data) # Push to Hugging Face Hub dataset.push_to_hub(dataset_name, token=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 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"]}], ids=[str(hash(item["code"]))], embeddings=[item["vectors"]] ) return client if __name__ == '__main__': client = init_chromadb() populate_sample_db(client) # Uncomment to save to Hugging Face # save_chromadb_to_hf()