from fastapi import FastAPI, HTTPException from pydantic import BaseModel import torch from transformers import AutoModelForCausalLM, AutoTokenizer import autopep8 import subprocess import time import re import os from pathlib import Path from fastapi.middleware.cors import CORSMiddleware import tempfile app = FastAPI(title="Code Evaluation & Optimization API") # CORS Configuration app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # Environment Setup - Modified for Hugging Face Spaces CACHE_DIR = Path(os.getenv("HF_HOME", "/tmp/huggingface")) CACHE_DIR.mkdir(parents=True, exist_ok=True) os.environ["TRANSFORMERS_CACHE"] = str(CACHE_DIR) os.environ["HF_HOME"] = str(CACHE_DIR) # Load AI Model - Using smaller model for Spaces compatibility MODEL_NAME = "codellama/CodeLlama-7b-instruct-hf" # More suitable for API use try: tokenizer = AutoTokenizer.from_pretrained( MODEL_NAME, cache_dir=str(CACHE_DIR)) model = AutoModelForCausalLM.from_pretrained( MODEL_NAME, device_map="auto", torch_dtype=torch.float16, cache_dir=str(CACHE_DIR)) except Exception as e: print(f"Model loading warning: {str(e)}") model = None tokenizer = None # Request Model class CodeRequest(BaseModel): code: str language: str = "python" def create_temp_file(code: str, extension: str) -> str: """Create temporary file in writable directory with proper permissions""" temp_dir = "/tmp/code_files" os.makedirs(temp_dir, exist_ok=True) fd, path = tempfile.mkstemp(suffix=f".{extension}", dir=temp_dir) with os.fdopen(fd, 'w') as tmp: tmp.write(code) os.chmod(path, 0o777) # Ensure executable permissions return path def cleanup_temp_files(): """Clean up temporary files""" temp_dir = "/tmp/code_files" if os.path.exists(temp_dir): for filename in os.listdir(temp_dir): file_path = os.path.join(temp_dir, filename) try: if os.path.isfile(file_path): os.unlink(file_path) except Exception as e: print(f"Error deleting {file_path}: {e}") # Helper Functions def evaluate_code(user_code: str, lang: str) -> dict: """Evaluate code for correctness, performance, and security""" start_time = time.time() file_ext = {"python": "py", "java": "java", "cpp": "cpp", "javascript": "js"}.get(lang, "txt") try: filename = create_temp_file(user_code, file_ext) commands = { "python": ["python3", filename], "java": ["javac", filename, "&&", "java", filename.replace(".java", "")], "cpp": ["g++", filename, "-o", f"{filename}.out", "&&", f"./{filename}.out"], "javascript": ["node", filename] } if lang not in commands: return {"status": "error", "message": "Unsupported language", "score": 0} result = subprocess.run( " ".join(commands[lang]), capture_output=True, text=True, timeout=5, shell=True ) exec_time = time.time() - start_time correctness = 1 if result.returncode == 0 else 0 error_message = None if correctness else result.stderr.strip() # Scoring logic readability_score = 20 if len(user_code) < 200 else 10 efficiency_score = 30 if exec_time < 1 else 10 security_score = 20 if "eval(" not in user_code and "exec(" not in user_code else 0 total_score = (correctness * 50) + readability_score + efficiency_score + security_score feedback = [] if correctness == 0: feedback.append("❌ Error in Code Execution! Check syntax or logic errors.") feedback.append(f"📌 Error Details: {error_message}") else: feedback.append("✅ Code executed successfully!") if efficiency_score < 30: feedback.append("⚡ Performance Issue: Code took longer to execute. Optimize loops or calculations.") if readability_score < 20: feedback.append("📖 Readability Issue: Code is lengthy. Break into smaller functions.") if security_score == 0: feedback.append("🔒 Security Risk: Avoid using eval() or exec().") return { "status": "success" if correctness else "error", "execution_time": round(exec_time, 3) if correctness else None, "score": max(0, min(100, total_score)), "feedback": "\n".join(feedback), "error_details": error_message if not correctness else None } except subprocess.TimeoutExpired: return {"status": "error", "message": "Execution timed out", "score": 0} except Exception as e: return {"status": "error", "message": str(e), "score": 0} finally: cleanup_temp_files() def optimize_code_ai(user_code: str, lang: str) -> str: """Generate optimized code using AI""" if model is None or tokenizer is None: raise HTTPException(status_code=503, detail="AI service temporarily unavailable") try: if lang == "python": user_code = autopep8.fix_code(user_code) user_code = re.sub(r"eval\((.*)\)", r"int(\1) # Removed eval for security", user_code) user_code = re.sub(r"/ 0", "/ 1 # Fixed division by zero", user_code) prompt = f"""Optimize this {lang} code for better performance and readability: ```{lang} {user_code} """ inputs = tokenizer(prompt, return_tensors="pt").to(model.device) with torch.no_grad(): outputs = model.generate(**inputs, max_length=1024, temperature=0.7) optimized_code = tokenizer.decode(outputs[0], skip_special_tokens=True) # Extract code between the last code block markers code_blocks = re.findall(r'```(?:python)?\n(.*?)\n```', optimized_code, re.DOTALL) if code_blocks: optimized_code = code_blocks[-1] # Get the last code block return optimized_code.strip() if optimized_code else user_code except Exception as e: raise HTTPException(status_code=500, detail=f"AI optimization failed: {str(e)}") # API Endpoints @app.post("/evaluate") async def evaluate_endpoint(request: CodeRequest): try: result = evaluate_code(request.code, request.language) return {"status": "success", "result": result} except Exception as e: raise HTTPException(status_code=400, detail=str(e)) @app.post("/optimize") async def optimize_endpoint(request: CodeRequest): try: optimized = optimize_code_ai(request.code, request.language) return {"status": "success", "optimized_code": optimized} except Exception as e: raise HTTPException(status_code=400, detail=str(e)) @app.get("/") def health_check(): return { "status": "API is running", "model": MODEL_NAME if model else "Not loaded", "endpoints": { "evaluate": "POST /evaluate", "optimize": "POST /optimize" } } @app.on_event("shutdown") def shutdown_event(): cleanup_temp_files() if __name__ == "__main__": import uvicorn uvicorn.run("main:app", host="0.0.0.0", port=7860)