chore: types

Author: rycerzes

src/core/env_server/__init__.py

@@ -15,7 +15,7 @@
     deserialize_action_with_preprocessing,
     serialize_observation,
 )
-from .types import Action, Observation, State, SchemaResponse
+from .types import Action, Observation, State, SchemaResponse, HealthResponse
 from .web_interface import create_web_interface_app, WebInterfaceManager
 
 __all__ = [
@@ -29,6 +29,7 @@
     "Observation",
     "State",
     "SchemaResponse",
+    "HealthResponse",
     # Base transforms
     "CompositeTransform",
     "NullTransform",

src/core/env_server/http_server.py

@@ -17,7 +17,7 @@
 import inspect
 import os
 from concurrent.futures import ThreadPoolExecutor
-from typing import Any, Dict, Optional, Type
+from typing import Optional, Type
 
 from fastapi import Body, FastAPI, HTTPException, status
 from pydantic import ValidationError
@@ -38,6 +38,7 @@
     StepResponse,
     EnvironmentMetadata,
     SchemaResponse,
+    HealthResponse,
 )
 
 
@@ -109,17 +110,14 @@ def _get_valid_kwargs(self, sig, kwargs, skip_params=None):
 
         return valid_kwargs
 
-    def register_routes(self, app: Any) -> None:
+    def register_routes(self, app: FastAPI) -> None:
         """
         Register HTTP routes on a FastAPI application.
 
         Args:
             app: FastAPI application instance
         """
 
-        if not isinstance(app, FastAPI):
-            raise TypeError("app must be a FastAPI instance")
-
         # Helper function to handle reset endpoint
         async def reset_handler(
             request: ResetRequest = Body(default_factory=ResetRequest),
@@ -283,8 +281,8 @@ async def step(request: StepRequest) -> StepResponse:
             ),
             GetEndpointConfig(
                 path="/health",
-                handler=lambda: {"status": "healthy"},
-                response_model=Dict[str, str],
+                handler=lambda: HealthResponse(status="healthy"),
+                response_model=HealthResponse,
                 tag="Health",
                 summary="Health check",
                 description="Check if the environment server is running and healthy.",
@@ -347,7 +345,7 @@ def create_app(
     action_cls: Type[Action],
     observation_cls: Type[Observation],
     env_name: Optional[str] = None,
-) -> Any:
+) -> FastAPI:
     """
     Create a FastAPI application with or without web interface.
 
@@ -385,7 +383,7 @@ def create_fastapi_app(
     env: Environment,
     action_cls: Type[Action],
     observation_cls: Type[Observation],
-) -> Any:
+) -> FastAPI:
     """Create a FastAPI application with comprehensive documentation."""
     try:
         from fastapi import FastAPI

src/core/env_server/route_config.py

@@ -12,22 +12,19 @@
 """
 
 from dataclasses import dataclass
-from typing import Callable, List, Type, TypeVar
+from typing import Callable, List, Type
 
 from fastapi import FastAPI
 from pydantic import BaseModel
 
-# TypeVar for generic response types
-T = TypeVar("T", bound=BaseModel)
-
 
 @dataclass
 class GetEndpointConfig:
     """Configuration for a simple GET endpoint."""
 
     path: str
     handler: Callable[[], BaseModel | dict]
-    response_model: Type[BaseModel] | Type[dict]
+    response_model: Type[BaseModel] | type[dict]
     tag: str
     summary: str
     description: str

src/core/env_server/types.py

@@ -44,7 +44,7 @@ class Observation(BaseModel):
     )
 
     done: bool = Field(default=False, description="Whether the episode has terminated")
-    reward: Union[bool, int, float, None] = Field(
+    reward: bool | int | float | None = Field(
         default=None, description="Reward signal from the last action"
     )
     metadata: Dict[str, Any] = Field(
@@ -201,3 +201,14 @@ class SchemaResponse(BaseModel):
     state: Dict[str, Any] = Field(
         description="JSON schema for environment state objects"
     )
+
+
+class HealthResponse(BaseModel):
+    """Response model for health check endpoint."""
+
+    model_config = ConfigDict(
+        extra="forbid",
+        validate_assignment=True,
+    )
+
+    status: str = Field(description="Health status of the environment server")