server.py

from fastapi import FastAPI
from pydantic import BaseModel
import subprocess
import tempfile
import os
import uuid
import shutil

app = FastAPI()

class ExecuteRequest(BaseModel):
    solution: str
    test: str
    language: str = "python3"

@app.post("/execute")
def execute(req: ExecuteRequest):
    solution = req.solution
    test = req.test
    language = req.language
    if language != "python3":
        return {"status": "failed", "error": f"Language {language} not supported."}

    tmp_dir = os.path.join(tempfile.gettempdir(), f"surprisal-py-{uuid.uuid4().hex}")
    os.makedirs(tmp_dir, exist_ok=True)

    try:
        sol_path = os.path.join(tmp_dir, "solution.py")
        spec_path = os.path.join(tmp_dir, "evaluation_spec.py")

        with open(sol_path, "w") as f:
            f.write(solution)
        with open(spec_path, "w") as f:
            f.write(test)

        # Run with 10s timeout
        # Using -m unittest to match original behavior
        result = subprocess.run(
            ["python3", "-m", "unittest", "evaluation_spec"],
            cwd=tmp_dir,
            capture_output=True,
            text=True,
            timeout=10,
            env={**os.environ, "PYTHONPATH": tmp_dir}
        )

        return {
            "status": "accepted" if result.returncode == 0 else "rejected",
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {"status": "rejected", "stderr": "Execution timed out."}
    except Exception as e:
        return {"status": "failed", "error": str(e)}
    finally:
        shutil.rmtree(tmp_dir, ignore_errors=True)

@app.get("/health")
def health():
    return "OK"

if __name__ == "__main__":
    import uvicorn
    port = int(os.getenv("PORT", 3003))
    uvicorn.run(app, host="0.0.0.0", port=port)