Update server/app.py

server/app.py

@@ -1,14 +1,36 @@
-import os
-import sys
-import uvicorn
+from fastapi import FastAPI
+from env import TrafficEnv
 
-# Add the parent directory to sys.path so 'inference.py' can be imported and env modules
-sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+CONFIG = {
+    "max_steps": 20,
+    "max_queue": 20,
+    "arrival_rate": (0, 2),
+    "discharge_rate": (3, 5),
+    "emergency_prob": 0.02,
+    "switch_penalty": 0.2,
+    "starvation_threshold": 10,
+    "burst_prob": 0.1,
+    "burst_multiplier": 1.2,
+}
 
-from inference import app
+app = FastAPI()
+env = TrafficEnv(CONFIG)
 
-def main():
-    uvicorn.run("server.app:app", host="0.0.0.0", port=7860)
+@app.get("/")
+def root():
+    return {"message": "Traffic Env Running"}
 
-if __name__ == "__main__":
-    main()
+@app.post("/reset")
+def reset():
+    return env.reset()
+
+@app.post("/step")
+def step(action: dict):
+    action_value = action.get("action", 0)
+    next_state, reward, done, info = env.step(action_value)
+    return {
+        "next_state": next_state,
+        "reward": reward,
+        "done": done,
+        "info": info,
+    }