Update server.py

server.py

@@ -1,13 +1,14 @@
+import os
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
-import subprocess
-import tempfile
-import os
 
 app = FastAPI()
 
+# Fix: Set a writable cache directory
+os.environ["TRANSFORMERS_CACHE"] = "/app/cache"
+
 # Load DeepSeek-Coder-V2-Base Model
 model_name = "deepseek-ai/DeepSeek-Coder-V2-Base"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -16,9 +17,6 @@ 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"""
@@ -30,37 +28,6 @@ def generate_code(request: CodeRequest):
 
     return {"generated_code": generated_code}
 
-@app.post("/test-code")
-def test_code(request: TestRequest):
-    """Runs automated testing on the 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"""
-    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)