Update server.py

server.py

@@ -3,13 +3,19 @@ from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
+import subprocess
+import tempfile
 
 app = FastAPI()
 
-# Fix: Set a writable cache directory
+# ✅ Fix: Set writable cache directory for Hugging Face models
 os.environ["TRANSFORMERS_CACHE"] = "/app/cache"
+os.environ["HF_HOME"] = "/app/cache"
 
-# Load DeepSeek-Coder-V2-Base Model
+# ✅ Ensure cache directory exists
+os.makedirs("/app/cache", exist_ok=True)
+
+# ✅ Load AI Model
 model_name = "deepseek-ai/DeepSeek-Coder-V2-Base"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
@@ -17,17 +23,51 @@ model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float
 class CodeRequest(BaseModel):
     user_story: str
 
+class TestRequest(BaseModel):
+    code: str
+
 @app.post("/generate-code")
 def generate_code(request: CodeRequest):
-    """Generates code based on user story"""
+    """Generates AI-powered structured code based on user story"""
     prompt = f"Generate structured code for: {request.user_story}"
-    
+
     inputs = tokenizer(prompt, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
     output = model.generate(**inputs, max_length=300)
     generated_code = tokenizer.decode(output[0], skip_special_tokens=True)
 
     return {"generated_code": generated_code}
 
+@app.post("/test-code")
+def test_code(request: TestRequest):
+    """Runs automated testing on AI-generated code"""
+    try:
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".py") as temp_file:
+            temp_file.write(request.code.encode())
+            temp_file.close()
+
+            result = subprocess.run(["pytest", temp_file.name], capture_output=True, text=True)
+            os.unlink(temp_file.name)
+
+            if result.returncode == 0:
+                return {"test_status": "All tests passed!"}
+            else:
+                return {"test_status": "Test failed!", "details": result.stderr}
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/execute-code")
+def execute_code():
+    """Executes AI-generated code and returns output"""
+    sample_code = "print('Hello from AI-generated code!')"
+    
+    try:
+        result = subprocess.run(["python3", "-c", sample_code], capture_output=True, text=True)
+        return {"status": "Execution successful!", "output": result.stdout}
+    
+    except Exception as e:
+        return {"status": "Execution failed!", "error": str(e)}
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)