DarkyCodez/precision-ag-env

1from fastapi import FastAPI
2from pydantic import BaseModel
3from typing import Dict, Any
4import uvicorn
5
6from core_sim import GreenhouseEnv
7
8# Pydantic Models
9class Action(BaseModel):
10    """Action model for the environment."""
11    action: int
12
13
14class Observation(BaseModel):
15    """Observation model - sensor readings from the environment."""
16    soil_moisture: float
17    nitrogen_level: float
18    crop_health: float
19
20
21class State(BaseModel):
22    """State model - internal variables of the environment."""
23    soil_moisture: float
24    nitrogen_level: float
25    crop_health: float
26    turn_count: int
27
28
29# Create FastAPI app
30app = FastAPI(title="GreenhouseEnv Server")
31
32# Global environment instance
33env = GreenhouseEnv()
34
35
36# Endpoints
37@app.post("/reset")
38def reset() -> Observation:
39    """
40    Reset the environment to initial state.
41    
42    Returns:
43        Observation: Initial observation of the environment
44    """
45    env.reset()
46    state = env.get_state()
47    return Observation(
48        soil_moisture=state["soil_moisture"],
49        nitrogen_level=state["nitrogen_level"],
50        crop_health=state["crop_health"]
51    )
52
53
54@app.post("/step")
55def step(action: Action) -> Dict[str, Any]:
56    """
57    Execute one step with the given action.
58    
59    Args:
60        action: Action model containing the action to take (0-3)
61    
62    Returns:
63        Dictionary containing:
64        - observation: current sensor readings
65        - reward: float between 0.0 and 1.0
66        - done: boolean (episode terminated)
67        - info: additional information dict
68    """
69    # Execute step
70    reward = env.step(action.action)
71    state = env.get_state()
72    
73    # Determine if episode is done (e.g., crop health depleted or harvest action taken)
74    done = (state["crop_health"] <= 0.0) or (action.action == 3)
75    
76    # Create observation from state
77    observation = {
78        "soil_moisture": state["soil_moisture"],
79        "nitrogen_level": state["nitrogen_level"],
80        "crop_health": state["crop_health"]
81    }
82    
83    # Additional info
84    info = {
85        "turn_count": state["turn_count"],
86        "action_taken": action.action
87    }
88    
89    return {
90        "observation": observation,
91        "reward": reward,
92        "done": done,
93        "info": info
94    }
95
96
97@app.get("/state")
98def get_state() -> State:
99    """
100    Get the current state of the environment.
101    
102    Returns:
103        State: Current internal variables
104    """
105    state = env.get_state()
106    return State(
107        soil_moisture=state["soil_moisture"],
108        nitrogen_level=state["nitrogen_level"],
109        crop_health=state["crop_health"],
110        turn_count=state["turn_count"]
111    )
112
113
114if __name__ == "__main__":
115    uvicorn.run(app, host="0.0.0.0", port=7860)
116