hadokenvskikoken/codegen_api

1from fastapi import FastAPI, HTTPException2from pydantic import BaseModel3import torch4from transformers import AutoModelForCausalLM, AutoTokenizer5import autopep86import subprocess7import time8import re9import os10from pathlib import Path11from fastapi.middleware.cors import CORSMiddleware12import tempfile13 14app = FastAPI(title="Code Evaluation & Optimization API")15 16# CORS Configuration17app.add_middleware(18    CORSMiddleware,19    allow_origins=["*"],20    allow_credentials=True,21    allow_methods=["*"],22    allow_headers=["*"],23)24 25# Environment Setup - Modified for Hugging Face Spaces26CACHE_DIR = Path(os.getenv("HF_HOME", "/tmp/huggingface"))27CACHE_DIR.mkdir(parents=True, exist_ok=True)28os.environ["TRANSFORMERS_CACHE"] = str(CACHE_DIR)29os.environ["HF_HOME"] = str(CACHE_DIR)30 31# Load AI Model - Using smaller model for Spaces compatibility32MODEL_NAME = "codellama/CodeLlama-7b-instruct-hf"  # More suitable for API use33 34try:35    tokenizer = AutoTokenizer.from_pretrained(36        MODEL_NAME,37        cache_dir=str(CACHE_DIR))38    model = AutoModelForCausalLM.from_pretrained(39        MODEL_NAME,40        device_map="auto",41        torch_dtype=torch.float16,42        cache_dir=str(CACHE_DIR))43except Exception as e:44    print(f"Model loading warning: {str(e)}")45    model = None46    tokenizer = None47 48# Request Model49class CodeRequest(BaseModel):50    code: str51    language: str = "python"52 53def create_temp_file(code: str, extension: str) -> str:54    """Create temporary file in writable directory with proper permissions"""55    temp_dir = "/tmp/code_files"56    os.makedirs(temp_dir, exist_ok=True)57    fd, path = tempfile.mkstemp(suffix=f".{extension}", dir=temp_dir)58    with os.fdopen(fd, 'w') as tmp:59        tmp.write(code)60    os.chmod(path, 0o777)  # Ensure executable permissions61    return path62 63def cleanup_temp_files():64    """Clean up temporary files"""65    temp_dir = "/tmp/code_files"66    if os.path.exists(temp_dir):67        for filename in os.listdir(temp_dir):68            file_path = os.path.join(temp_dir, filename)69            try:70                if os.path.isfile(file_path):71                    os.unlink(file_path)72            except Exception as e:73                print(f"Error deleting {file_path}: {e}")74 75# Helper Functions76def evaluate_code(user_code: str, lang: str) -> dict:77    """Evaluate code for correctness, performance, and security"""78    start_time = time.time()79    file_ext = {"python": "py", "java": "java", "cpp": "cpp", "javascript": "js"}.get(lang, "txt")80    81    try:82        filename = create_temp_file(user_code, file_ext)83        84        commands = {85            "python": ["python3", filename],86            "java": ["javac", filename, "&&", "java", filename.replace(".java", "")],87            "cpp": ["g++", filename, "-o", f"{filename}.out", "&&", f"./{filename}.out"],88            "javascript": ["node", filename]89        }90 91        if lang not in commands:92            return {"status": "error", "message": "Unsupported language", "score": 0}93 94        result = subprocess.run(95            " ".join(commands[lang]),96            capture_output=True,97            text=True,98            timeout=5,99            shell=True100        )101        exec_time = time.time() - start_time102        correctness = 1 if result.returncode == 0 else 0103        error_message = None if correctness else result.stderr.strip()104 105        # Scoring logic106        readability_score = 20 if len(user_code) < 200 else 10107        efficiency_score = 30 if exec_time < 1 else 10108        security_score = 20 if "eval(" not in user_code and "exec(" not in user_code else 0109        total_score = (correctness * 50) + readability_score + efficiency_score + security_score110 111        feedback = []112        if correctness == 0:113            feedback.append("❌ Error in Code Execution! Check syntax or logic errors.")114            feedback.append(f"📌 Error Details: {error_message}")115        else:116            feedback.append("✅ Code executed successfully!")117 118        if efficiency_score < 30:119            feedback.append("⚡ Performance Issue: Code took longer to execute. Optimize loops or calculations.")120        if readability_score < 20:121            feedback.append("📖 Readability Issue: Code is lengthy. Break into smaller functions.")122        if security_score == 0:123            feedback.append("🔒 Security Risk: Avoid using eval() or exec().")124 125        return {126            "status": "success" if correctness else "error",127            "execution_time": round(exec_time, 3) if correctness else None,128            "score": max(0, min(100, total_score)),129            "feedback": "\n".join(feedback),130            "error_details": error_message if not correctness else None131        }132 133    except subprocess.TimeoutExpired:134        return {"status": "error", "message": "Execution timed out", "score": 0}135    except Exception as e:136        return {"status": "error", "message": str(e), "score": 0}137    finally:138        cleanup_temp_files()139 140def optimize_code_ai(user_code: str, lang: str) -> str:141    """Generate optimized code using AI"""142    if model is None or tokenizer is None:143        raise HTTPException(status_code=503, detail="AI service temporarily unavailable")144    145    try:146        if lang == "python":147            user_code = autopep8.fix_code(user_code)148            user_code = re.sub(r"eval\((.*)\)", r"int(\1)  # Removed eval for security", user_code)149            user_code = re.sub(r"/ 0", "/ 1  # Fixed division by zero", user_code)150        151        prompt = f"""Optimize this {lang} code for better performance and readability:152```{lang}153{user_code}154        """155        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)156        157        with torch.no_grad():158            outputs = model.generate(**inputs, max_length=1024, temperature=0.7)159        160        optimized_code = tokenizer.decode(outputs[0], skip_special_tokens=True)161        162        # Extract code between the last code block markers163        code_blocks = re.findall(r'```(?:python)?\n(.*?)\n```', optimized_code, re.DOTALL)164        if code_blocks:165            optimized_code = code_blocks[-1]  # Get the last code block166        167        return optimized_code.strip() if optimized_code else user_code168    except Exception as e:169        raise HTTPException(status_code=500, detail=f"AI optimization failed: {str(e)}")170 171# API Endpoints172@app.post("/evaluate")173async def evaluate_endpoint(request: CodeRequest):174    try:175        result = evaluate_code(request.code, request.language)176        return {"status": "success", "result": result}177    except Exception as e:178        raise HTTPException(status_code=400, detail=str(e))179 180@app.post("/optimize")181async def optimize_endpoint(request: CodeRequest):182    try:183        optimized = optimize_code_ai(request.code, request.language)184        return {"status": "success", "optimized_code": optimized}185    except Exception as e:186        raise HTTPException(status_code=400, detail=str(e))187 188@app.get("/")189def health_check():190    return {191        "status": "API is running",192        "model": MODEL_NAME if model else "Not loaded",193        "endpoints": {194            "evaluate": "POST /evaluate",195            "optimize": "POST /optimize"196        }197    }198 199@app.on_event("shutdown")200def shutdown_event():201    cleanup_temp_files()202 203if __name__ == "__main__":204    import uvicorn205    uvicorn.run("main:app", host="0.0.0.0", port=7860)