Spaces:
Sleeping
Sleeping
File size: 2,301 Bytes
798461f 016c7fd 798461f 016c7fd 798461f 016c7fd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | """
client.py
---------
Client for the SQL Analytics Environment.
SQLAnalyticsEnv wraps EnvClient with the three typed models specific to
this environment (SQLAction, SQLObservation, SQLState). All connection
and session management is inherited from EnvClient.
Example usage
-------------
# Against a local server:
with SQLAnalyticsEnv(base_url="http://localhost:8000") as env:
obs = env.reset(task_id="task_003")
print(obs.task_description)
result_obs = env.step(
SQLAction(sql="SELECT name, SUM(total_amount) FROM ...")
)
print(result_obs.result, result_obs.reward)
# Against a HuggingFace Space:
env = SQLAnalyticsEnv.from_env("your-hf-username/sql-analytics-env")
try:
obs = env.reset(task_id="task_001")
...
finally:
env.close()
# From a local Docker image:
env = SQLAnalyticsEnv.from_docker_image("sql-analytics-env:latest")
try:
obs = env.reset(task_id="task_002")
...
finally:
env.close()
"""
from typing import Any, Dict
from openenv.core.env_client import EnvClient, StepResult
try:
from .models import SQLAction, SQLObservation, SQLState
except ImportError:
from models import SQLAction, SQLObservation, SQLState
class SQLAnalyticsEnv(EnvClient[SQLAction, SQLObservation, SQLState]):
"""
Typed client for the SQL Analytics OpenEnv environment.
Inherits the full EnvClient interface:
reset(**kwargs) -> StepResult[SQLObservation]
step(action: SQLAction) -> StepResult[SQLObservation]
state() -> SQLState
close()
Convenience class methods (from EnvClient):
SQLAnalyticsEnv.from_env(hf_repo_id)
SQLAnalyticsEnv.from_docker_image(image_tag)
"""
def _step_payload(self, action: SQLAction) -> Dict[str, Any]:
return action.model_dump()
def _parse_result(self, payload: Dict[str, Any]) -> StepResult[SQLObservation]:
obs_dict = payload.get("observation", {})
obs = SQLObservation(**obs_dict)
return StepResult(
observation=obs,
reward=payload.get("reward"),
done=payload.get("done", False),
)
def _parse_state(self, payload: Dict[str, Any]) -> SQLState:
return SQLState(**payload)
|