Update server.py

server.py

@@ -1,20 +1,65 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 import subprocess
+import tempfile
+import os
 
 app = FastAPI()
 
+# Load DeepSeek-Coder-V2-Base 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")
+
 class CodeRequest(BaseModel):
     user_story: str
 
+class TestRequest(BaseModel):
+    code: str
+
 @app.post("/generate-code")
 def generate_code(request: CodeRequest):
-    result = subprocess.run(
-        ["ollama", "run", "deepseek-coder-v2", request.user_story],
-        capture_output=True,
-        text=True
-    )
-    return {"generated_code": result.stdout.strip()}
+    """Generates 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 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