# app.py
from flask import Flask, request, render_template, jsonify, send_file
from parser import parse_python_code
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, create_collection, save_chromadb_to_hf, HF_DATASET_NAME
from datasets import Dataset
import logging
# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
# User-configurable variables
UPLOAD_DIR = "./uploads" # Directory for uploads
app = Flask(__name__)
def reconstruct_code(parts):
"""Reconstruct the original code from parsed parts."""
sorted_parts = sorted(parts, key=lambda p: p['location'][0])
return ''.join(part['source'] for part in sorted_parts)
@app.route('/', methods=['GET', 'POST'])
def index():
if request.method == 'POST':
parts = None
filename = 'unnamed.py'
code_input = None
query_results = None
# Handle file upload or pasted code (parsing)
if 'file' in request.files and request.files['file'].filename:
file = request.files['file']
if not file.filename.endswith('.py'):
return 'Invalid file type. Please upload a Python file.', 400
filename = file.filename
file_path = os.path.join(UPLOAD_DIR, filename)
file.save(file_path)
with open(file_path, 'r') as f:
code_input = f.read()
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}")
# 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
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'
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}")
# 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
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(',')]
try:
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,
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']
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}")
# Verify query results
logger.info(f"Found {len(query_results)} matching programs in ChromaDB")
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:
# Reset ChromaDB collection
client = init_chromadb()
try:
client.delete_collection(DB_NAME)
logger.info(f"Deleted ChromaDB collection: {DB_NAME}")
except Exception as e:
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")
# Verify collection is empty
count = collection.count()
logger.info(f"ChromaDB now contains {count} entries after reset (should be 0)")
# Reset Hugging Face dataset (replace with empty dataset)
try:
empty_data = {
"code": [],
"sequence": [],
"vectors": [],
"description_tokens": [],
"program_vectors": []
}
empty_dataset = Dataset.from_dict(empty_data)
empty_dataset.push_to_hub(HF_DATASET_NAME, token=os.getenv("HF_KEY"))
logger.info(f"Replaced Hugging Face dataset {HF_DATASET_NAME} with empty dataset")
except Exception as e:
logger.error(f"Error replacing Hugging Face dataset: {e}")
raise
# 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,
filename="Hugging Face Dataset Processed",
reconstructed_code=None,
code_input=None,
query_results=None,
message="Hugging Face dataset processed and stored successfully with fresh database and empty dataset."
)
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:
# Check if collection exists, get or create if needed
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")
# Verify collection state
count = collection.count()
logger.info(f"ChromaDB contains {count} entries before loading")
# 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,
filename="Hugging Face Dataset Loaded",
reconstructed_code=None,
code_input=None,
query_results=None,
message="Hugging Face dataset loaded and stored successfully."
)
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 and Hugging Face dataset (no repopulation with samples)
try:
client = init_chromadb()
try:
client.delete_collection(DB_NAME)
logger.info(f"Deleted ChromaDB collection: {DB_NAME}")
except Exception as e:
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
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)")
# Reset Hugging Face dataset (replace with empty dataset)
try:
empty_data = {
"code": [],
"sequence": [],
"vectors": [],
"description_tokens": [],
"program_vectors": []
}
empty_dataset = Dataset.from_dict(empty_data)
empty_dataset.push_to_hub(HF_DATASET_NAME, token=os.getenv("HF_KEY"))
logger.info(f"Replaced Hugging Face dataset {HF_DATASET_NAME} with empty dataset")
except Exception as e:
logger.error(f"Error replacing Hugging Face dataset: {e}")
raise
return render_template(
'results_partial.html',
parts=None,
filename="Database Reset",
reconstructed_code=None,
code_input=None,
query_results=None,
message="Database and Hugging Face dataset reset successfully."
)
except Exception as e:
logger.error(f"Error resetting database: {e}")
return f"Error resetting database: {e}", 500
if parts:
indexed_parts = [{'index': i + 1, **part} for i, part in enumerate(parts)]
reconstructed_code = reconstruct_code(indexed_parts)
return render_template(
'results_partial.html',
parts=indexed_parts,
filename=filename,
reconstructed_code=reconstructed_code,
code_input=code_input,
query_results=None
)
return 'No file, code, or query provided', 400
# Initial page load (start empty, no default population)
logger.info("Application started, database empty until triggered by buttons")
return render_template('index.html', parts=None, filename=None, reconstructed_code=None, code_input=None, query_results=None)
@app.route('/export_json', methods=['POST'])
def export_json():
parts = request.json.get('parts', [])
export_data = [{'vector': part['vector'], 'source': part['source'], 'description': generate_description_tokens([part['category']], [part['vector']])} for part in parts]
json_str = json.dumps(export_data, indent=2)
buffer = io.BytesIO(json_str.encode('utf-8'))
buffer.seek(0)
return send_file(
buffer,
as_attachment=True,
download_name='code_vectors.json',
mimetype='application/json'
)
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 " ".join(tokens)
if __name__ == '__main__':
if not os.path.exists(UPLOAD_DIR):
os.makedirs(UPLOAD_DIR)
app.run(host="0.0.0.0", port=7860) # Bind to all interfaces for Hugging Face Spaces