albert-einstein-09/codedark
3
1"""2CodeDark FastAPI Server3 4OpenEnv-compatible HTTP server for CodeDark environment.5Provides /reset, /step, /state, and /health endpoints.6"""7 8import os9from contextlib import asynccontextmanager10from typing import Optional11 12import uvicorn13from fastapi import FastAPI, HTTPException14from fastapi.middleware.cors import CORSMiddleware15from fastapi.responses import HTMLResponse16 17from ..models import (18 CodeDarkAction,19 CodeDarkObservation,20 CodeDarkState,21 ResetRequest,22 StepRequest,23 HealthResponse,24)25from .environment import CodeDarkEnvironment26 27 28# Global environment instance29_env: Optional[CodeDarkEnvironment] = None30 31 32def get_env() -> CodeDarkEnvironment:33 """Get or create environment instance."""34 global _env35 if _env is None:36 _env = CodeDarkEnvironment(37 data_dir=os.environ.get("CODEDARK_DATA_DIR"),38 tasks_path=os.environ.get("CODEDARK_TASKS_PATH"),39 max_turns=int(os.environ.get("CODEDARK_MAX_TURNS", "10")),40 )41 return _env42 43 44@asynccontextmanager45async def lifespan(app: FastAPI):46 """Lifespan context manager for startup/shutdown."""47 # Startup: initialize environment48 get_env()49 yield50 # Shutdown: cleanup if needed51 global _env52 _env = None53 54 55# Create FastAPI app56app = FastAPI(57 title="CodeDark Environment",58 description="Multi-turn data analytics environment for RL agent training",59 version="0.1.0",60 lifespan=lifespan,61)62 63# Add CORS middleware64app.add_middleware(65 CORSMiddleware,66 allow_origins=["*"],67 allow_credentials=True,68 allow_methods=["*"],69 allow_headers=["*"],70)71 72 73@app.get("/", response_class=HTMLResponse)74async def root():75 """Landing page with API documentation."""76 env = get_env()77 return f"""78 <!DOCTYPE html>79 <html>80 <head>81 <title>CodeDark Environment</title>82 <style>83 body {{ font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;84 max-width: 800px; margin: 50px auto; padding: 20px; background: #0d1117; color: #c9d1d9; }}85 h1 {{ color: #f0883e; }}86 h2 {{ color: #8b949e; border-bottom: 1px solid #30363d; padding-bottom: 8px; }}87 code {{ background: #161b22; padding: 2px 6px; border-radius: 4px; color: #7ee787; }}88 pre {{ background: #161b22; padding: 16px; border-radius: 8px; overflow-x: auto; }}89 .endpoint {{ background: #21262d; padding: 12px; margin: 8px 0; border-radius: 6px; }}90 .method {{ color: #7ee787; font-weight: bold; }}91 .path {{ color: #79c0ff; }}92 a {{ color: #58a6ff; }}93 .stats {{ display: flex; gap: 20px; margin: 20px 0; }}94 .stat {{ background: #21262d; padding: 16px; border-radius: 8px; text-align: center; }}95 .stat-value {{ font-size: 24px; color: #f0883e; font-weight: bold; }}96 .stat-label {{ color: #8b949e; font-size: 12px; }}97 </style>98 </head>99 <body>100 <h1>CodeDark Environment Server</h1>101 <p>OpenEnv-compatible multi-turn data analytics environment for RL agent training.</p>102 103 <div class="stats">104 <div class="stat"><div class="stat-value">{len(env.tasks)}</div><div class="stat-label">Tasks</div></div>105 <div class="stat"><div class="stat-value">750K</div><div class="stat-label">Bank Rows</div></div>106 <div class="stat"><div class="stat-value">500K</div><div class="stat-label">Road Rows</div></div>107 <div class="stat"><div class="stat-value">5</div><div class="stat-label">Tools</div></div>108 </div>109 110 <h2>API Endpoints</h2>111 <div class="endpoint"><span class="method">GET</span> <span class="path">/health</span> - Health check</div>112 <div class="endpoint"><span class="method">POST</span> <span class="path">/reset</span> - Reset for new episode</div>113 <div class="endpoint"><span class="method">POST</span> <span class="path">/step</span> - Execute action</div>114 <div class="endpoint"><span class="method">GET</span> <span class="path">/state</span> - Current state</div>115 <div class="endpoint"><span class="method">GET</span> <span class="path">/metadata</span> - Environment info</div>116 <div class="endpoint"><span class="method">GET</span> <span class="path">/schema</span> - Type schemas</div>117 <div class="endpoint"><span class="method">GET</span> <span class="path">/docs</span> - Interactive API docs</div>118 119 <h2>Quick Start</h2>120 <pre>121import requests122 123BASE = "https://albert-einstein-09-codedark.hf.space"124 125# Reset for new task126obs = requests.post(f"{{BASE}}/reset").json()127print(f"Task: {{obs['question']}}")128 129# Run Python code130obs = requests.post(f"{{BASE}}/step", json={{131 "tool": "run_python",132 "args": "<code>result = df.shape</code>"133}}).json()134print(f"Result: {{obs['stdout']}}")135 136# Submit answer137obs = requests.post(f"{{BASE}}/step", json={{138 "tool": "submit_answer",139 "args": "<answer>42.5</answer>"140}}).json()141print(f"Reward: {{obs['reward']}}")142 </pre>143 144 <h2>Tools</h2>145 <ul>146 <li><code>run_python</code> - Execute Python/pandas code</li>147 <li><code>read_notes</code> - Read saved notes</li>148 <li><code>save_note</code> - Save note for later</li>149 <li><code>clarify</code> - Ask clarifying question</li>150 <li><code>submit_answer</code> - Submit final answer</li>151 </ul>152 153 <h2>Links</h2>154 <p>155 <a href="/docs">Interactive API Docs</a> |156 <a href="https://github.com/vj-09/codeblue-env">GitHub</a> |157 <a href="https://www.analytics-rl.com">Leaderboard</a>158 </p>159 </body>160 </html>161 """162 163 164@app.get("/health", response_model=HealthResponse)165async def health():166 """Health check endpoint."""167 return HealthResponse(168 status="healthy",169 environment="codedark",170 version="0.1.0",171 )172 173 174@app.post("/reset", response_model=CodeDarkObservation)175async def reset(request: ResetRequest = None):176 """Reset environment for a new episode.177 178 Args:179 request: Optional reset request with task_id and seed180 181 Returns:182 Initial observation183 """184 env = get_env()185 186 if request is None:187 request = ResetRequest()188 189 obs = env.reset(task_id=request.task_id, seed=request.seed)190 return obs191 192 193@app.post("/step", response_model=CodeDarkObservation)194async def step(request: StepRequest):195 """Execute an action and return observation.196 197 Args:198 request: Step request with tool and args199 200 Returns:201 Observation after action execution202 """203 env = get_env()204 205 # Validate tool206 valid_tools = ["run_python", "read_notes", "save_note", "clarify", "submit_answer"]207 if request.tool not in valid_tools:208 raise HTTPException(209 status_code=400,210 detail=f"Invalid tool: {request.tool}. Valid tools: {valid_tools}",211 )212 213 action = CodeDarkAction(tool=request.tool, args=request.args)214 obs = env.step(action)215 return obs216 217 218@app.get("/state", response_model=CodeDarkState)219async def state():220 """Get current environment state.221 222 Returns:223 Current CodeDarkState224 """225 env = get_env()226 return env.state227 228 229@app.get("/metadata")230async def metadata():231 """Get environment metadata.232 233 Returns:234 Environment metadata dict235 """236 env = get_env()237 return {238 "name": "codedark",239 "version": "0.1.0",240 "description": "Multi-turn data analytics environment for RL agent training",241 "max_turns": env.max_turns,242 "max_clarifications": env.max_clarifications,243 "num_tasks": len(env.tasks),244 "tools": [245 {"name": "run_python", "description": "Execute Python/pandas code"},246 {"name": "read_notes", "description": "Read all saved notes"},247 {"name": "save_note", "description": "Save a note for later recall"},248 {"name": "clarify", "description": "Ask clarifying question (max 2)"},249 {"name": "submit_answer", "description": "Submit final answer"},250 ],251 "reward_structure": {252 "max_reward": 1.0,253 "components": [254 {"name": "correctness", "weight": 0.80},255 {"name": "efficiency", "weight": 0.10},256 {"name": "token_cost", "weight": 0.10},257 ],258 },259 }260 261 262@app.get("/schema")263async def schema():264 """Get environment schema for Action, Observation, State.265 266 Returns:267 JSON schemas for all types268 """269 return {270 "action": CodeDarkAction.model_json_schema(),271 "observation": CodeDarkObservation.model_json_schema(),272 "state": CodeDarkState.model_json_schema(),273 }274 275 276def main():277 """Run the server."""278 uvicorn.run(279 "codedark.server.app:app",280 host=os.environ.get("HOST", "0.0.0.0"),281 port=int(os.environ.get("PORT", "8000")),282 reload=os.environ.get("RELOAD", "false").lower() == "true",283 )284 285 286if __name__ == "__main__":287 main()288