shivamtech9395/maths_reasoning_env

1"""2server/app.py — FastAPI application for MathReasoningEnv.3Exposes /reset, /step, /state, /health endpoints.4"""5 6from fastapi import FastAPI, HTTPException7from fastapi.middleware.cors import CORSMiddleware8from pydantic import BaseModel9from typing import Optional10import sys11import os12 13# Allow imports from parent package when running standalone14sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))15 16from models import MathAction, MathObservation, MathState17from server.math_environment import MathReasoningEnvironment18 19app = FastAPI(20    title="MathReasoningEnv",21    description="OpenEnv-compatible Math Reasoning Environment with 5 graded task types",22    version="0.1.0",23)24 25app.add_middleware(26    CORSMiddleware,27    allow_origins=["*"],28    allow_methods=["*"],29    allow_headers=["*"],30)31 32# Global environment instance (single session for simplicity)33env = MathReasoningEnvironment()34 35 36class ResetRequest(BaseModel):37    task_type: Optional[str] = None38 39 40# ──────────────────────────────────────────────41# Endpoints42# ──────────────────────────────────────────────43 44@app.get("/health")45def health():46    return {"status": "ok", "env": "math_reasoning_env"}47 48 49@app.post("/reset", response_model=MathObservation)50def reset(req: ResetRequest = ResetRequest()):51    """Reset the environment and return the first problem."""52    obs = env.reset(task_type=req.task_type)53    return obs54 55 56@app.post("/step", response_model=MathObservation)57def step(action: MathAction):58    """Submit an answer and receive feedback + next problem."""59    if env.state.episode_id == "":60        raise HTTPException(status_code=400, detail="Call /reset before /step")61    obs = env.step(action)62    return obs63 64 65@app.get("/state", response_model=MathState)66def state():67    """Return the current internal state."""68    return env.state69 70 71@app.get("/tasks")72def list_tasks():73    """List all supported task types with grader info."""74    return {75        "tasks": [76            {"name": "arithmetic", "description": "Basic arithmetic (+, -, *, /)", "grader": True},77            {"name": "algebra", "description": "Solve linear equations", "grader": True},78            {"name": "word_problems", "description": "Multi-step word problems", "grader": True},79            {"name": "number_theory", "description": "GCD, LCM problems", "grader": True},80            {"name": "geometry", "description": "Area, perimeter calculations", "grader": True},81        ]82    }83