Spaces:
Sleeping
Sleeping
| """FastAPI and WebSocket entrypoint for SpreadsheetBench OpenEnv.""" | |
| from __future__ import annotations | |
| from collections.abc import Callable | |
| import os | |
| from fastapi import FastAPI | |
| from openenv.core.env_server.http_server import create_app | |
| from openenv.core.env_server.interfaces import Environment | |
| from starlette.types import ASGIApp, Receive, Scope, Send | |
| from starlette.websockets import WebSocketDisconnect | |
| from spreadsheetbench_env.models import ( | |
| SpreadsheetBenchAction, | |
| SpreadsheetBenchObservation, | |
| native_tool_specs, | |
| ) | |
| from spreadsheetbench_env.server.environment import SpreadsheetBenchEnvironment | |
| _SEND_AFTER_CLOSE = 'Cannot call "send" once a close message has been sent.' | |
| class DisconnectedWebSocketMiddleware: | |
| """Suppress only OpenEnv errors caused by an already-disconnected client.""" | |
| def __init__(self, app: ASGIApp) -> None: | |
| self.app = app | |
| async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: | |
| try: | |
| await self.app(scope, receive, send) | |
| except WebSocketDisconnect: | |
| if scope["type"] != "websocket": | |
| raise | |
| except RuntimeError as error: | |
| if scope["type"] != "websocket" or str(error) != _SEND_AFTER_CLOSE: | |
| raise | |
| def _configured_capacity() -> int: | |
| raw = os.getenv("MAX_CONCURRENT_ENVS", "1") | |
| try: | |
| capacity = int(raw) | |
| except ValueError as error: | |
| raise ValueError(f"MAX_CONCURRENT_ENVS must be an integer, got {raw!r}") from error | |
| if capacity < 1: | |
| raise ValueError("MAX_CONCURRENT_ENVS must be at least 1") | |
| return capacity | |
| def _configured_default_task_id() -> int | None: | |
| value = os.getenv("DEFAULT_TASK_ID") | |
| if value is None or not value.strip(): | |
| return None | |
| try: | |
| return int(value) | |
| except ValueError as error: | |
| raise ValueError(f"DEFAULT_TASK_ID must be an integer, got {value!r}") from error | |
| def _configured_environment_class( | |
| default_task_id: int | None, | |
| ) -> type[SpreadsheetBenchEnvironment]: | |
| if default_task_id is None: | |
| return SpreadsheetBenchEnvironment | |
| class ConfiguredSpreadsheetBenchEnvironment(SpreadsheetBenchEnvironment): | |
| def __init__(self) -> None: | |
| super().__init__(default_task_id=default_task_id) | |
| ConfiguredSpreadsheetBenchEnvironment.__name__ = "SpreadsheetBenchEnvironment" | |
| return ConfiguredSpreadsheetBenchEnvironment | |
| def create_spreadsheetbench_app( | |
| max_concurrent_envs: int | None = None, | |
| environment_factory: Callable[[], Environment] | None = None, | |
| ) -> FastAPI: | |
| """Create the pinned OpenEnv app without constructing runtime data eagerly.""" | |
| capacity = _configured_capacity() if max_concurrent_envs is None else max_concurrent_envs | |
| if capacity < 1: | |
| raise ValueError("MAX_CONCURRENT_ENVS must be at least 1") | |
| factory = environment_factory or _configured_environment_class( | |
| _configured_default_task_id() | |
| ) | |
| app = create_app( | |
| factory, | |
| SpreadsheetBenchAction, | |
| SpreadsheetBenchObservation, | |
| env_name="spreadsheetbench_env", | |
| max_concurrent_envs=capacity, | |
| ) | |
| app.add_api_route("/tools", native_tool_specs, methods=["GET"]) | |
| app.add_middleware(DisconnectedWebSocketMiddleware) | |
| return app | |
| app = create_spreadsheetbench_app() | |
| def main(host: str = "0.0.0.0", port: int = 8000) -> None: | |
| """Run the standalone OpenEnv server.""" | |
| import uvicorn | |
| uvicorn.run(app, host=host, port=port) | |
| if __name__ == "__main__": | |
| main() | |