Upload folder using huggingface_hub

server/sql_environment.py

@@ -1,8 +1,8 @@
 import sqlite3
 import random
-from typing import Tuple
+from typing import Any, Optional, Tuple
 
-from openenv.core.env_server.types import EnvBase
+from openenv.core.env_server.interfaces import Environment
 from models import SQLAction, SQLObservation, SQLState
 from server.challenges import CHALLENGES
 
@@ -47,9 +47,22 @@ def _results_match(schema_sql: str, query_a: str, query_b: str) -> bool:
         return False
 
 
-class SQLTutorEnvironment(EnvBase[SQLAction, SQLObservation, SQLState]):
+class SQLTutorEnvironment(Environment[SQLAction, SQLObservation, SQLState]):
+    SUPPORTS_CONCURRENT_SESSIONS = True
+
+    def __init__(self):
+        super().__init__()
+        self._state = SQLState()
+
+    def reset(
+        self,
+        seed: Optional[int] = None,
+        episode_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> SQLObservation:
+        if seed is not None:
+            random.seed(seed)
 
-    def reset(self) -> Tuple[SQLObservation, SQLState]:
         challenge = random.choice(CHALLENGES)
 
         state = SQLState(
@@ -65,6 +78,8 @@ class SQLTutorEnvironment(EnvBase[SQLAction, SQLObservation, SQLState]):
             hints_used=0,
             is_resolved=False,
             cumulative_reward=0.0,
+            episode_id=episode_id,
+            step_count=0,
         )
         self._state = state
 
@@ -81,12 +96,20 @@ class SQLTutorEnvironment(EnvBase[SQLAction, SQLObservation, SQLState]):
             steps_taken=0,
             max_steps=state.max_steps,
             hints_used=0,
+            done=False,
+            reward=None,
         )
-        return observation, state
-
-    def step(self, action: SQLAction) -> Tuple[SQLObservation, float, bool, SQLState]:
+        return observation
+
+    def step(
+        self,
+        action: SQLAction,
+        timeout_s: Optional[float] = None,
+        **kwargs: Any,
+    ) -> SQLObservation:
         state = self._state
         state.steps_taken += 1
+        state.step_count += 1
         reward = 0.0
         done = False
         hint = None
@@ -147,9 +170,12 @@ class SQLTutorEnvironment(EnvBase[SQLAction, SQLObservation, SQLState]):
             steps_taken=state.steps_taken,
             max_steps=state.max_steps,
             hints_used=state.hints_used,
+            done=done,
+            reward=reward,
         )
 
-        return observation, reward, done, state
+        return observation
 
-    def get_state(self) -> SQLState:
+    @property
+    def state(self) -> SQLState:
         return self._state