Nitesh321 Claude Opus 4.6 commited on
Commit
016c7fd
Β·
1 Parent(s): d8571a4

fix: implement abstract methods in SQLAnalyticsEnv and use sync() wrapper

Browse files

- Implement _step_payload, _parse_result, _parse_state in client.py
- Use .sync() wrapper since EnvClient is async-first
- Access StepResult.observation for env-specific fields (db_schema, error, result)
- Access StepResult.reward and StepResult.done for step-level fields

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Files changed (2) hide show
  1. client.py +20 -4
  2. inference.py +10 -10
client.py CHANGED
@@ -36,7 +36,9 @@ Example usage
36
  env.close()
37
  """
38
 
39
- from openenv.core.env_client import EnvClient
 
 
40
 
41
  try:
42
  from .models import SQLAction, SQLObservation, SQLState
@@ -49,8 +51,8 @@ class SQLAnalyticsEnv(EnvClient[SQLAction, SQLObservation, SQLState]):
49
  Typed client for the SQL Analytics OpenEnv environment.
50
 
51
  Inherits the full EnvClient interface:
52
- reset(**kwargs) -> SQLObservation
53
- step(action: SQLAction) -> SQLObservation
54
  state() -> SQLState
55
  close()
56
 
@@ -58,4 +60,18 @@ class SQLAnalyticsEnv(EnvClient[SQLAction, SQLObservation, SQLState]):
58
  SQLAnalyticsEnv.from_env(hf_repo_id)
59
  SQLAnalyticsEnv.from_docker_image(image_tag)
60
  """
61
- pass
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
36
  env.close()
37
  """
38
 
39
+ from typing import Any, Dict
40
+
41
+ from openenv.core.env_client import EnvClient, StepResult
42
 
43
  try:
44
  from .models import SQLAction, SQLObservation, SQLState
 
51
  Typed client for the SQL Analytics OpenEnv environment.
52
 
53
  Inherits the full EnvClient interface:
54
+ reset(**kwargs) -> StepResult[SQLObservation]
55
+ step(action: SQLAction) -> StepResult[SQLObservation]
56
  state() -> SQLState
57
  close()
58
 
 
60
  SQLAnalyticsEnv.from_env(hf_repo_id)
61
  SQLAnalyticsEnv.from_docker_image(image_tag)
62
  """
63
+
64
+ def _step_payload(self, action: SQLAction) -> Dict[str, Any]:
65
+ return action.model_dump()
66
+
67
+ def _parse_result(self, payload: Dict[str, Any]) -> StepResult[SQLObservation]:
68
+ obs_dict = payload.get("observation", {})
69
+ obs = SQLObservation(**obs_dict)
70
+ return StepResult(
71
+ observation=obs,
72
+ reward=payload.get("reward"),
73
+ done=payload.get("done", False),
74
+ )
75
+
76
+ def _parse_state(self, payload: Dict[str, Any]) -> SQLState:
77
+ return SQLState(**payload)
inference.py CHANGED
@@ -141,11 +141,11 @@ def _extract_sql(text: str) -> str:
141
 
142
  # ── Agent loop ────────────────────────────────────────────────────────────────
143
 
144
- def run_task(client: OpenAI, env: SQLAnalyticsEnv, task_id: str) -> float:
145
  """Run one episode. Returns the final score in [0, 1]."""
146
- obs = env.reset(task_id=task_id)
147
- schema = obs.db_schema or ""
148
- task_description = obs.task_description or ""
149
 
150
  log_start(task=task_id, env=BENCHMARK, model=MODEL_NAME)
151
 
@@ -186,12 +186,12 @@ def run_task(client: OpenAI, env: SQLAnalyticsEnv, task_id: str) -> float:
186
  print(f"[DEBUG] No SQL at step={step}, stopping.", file=sys.stderr, flush=True)
187
  break
188
 
189
- obs = env.step(SQLAction(sql=sql))
190
 
191
- reward = obs.reward or 0.0
192
- done = obs.done
193
- error = obs.error
194
- result_rows = len(obs.result) if obs.result else 0
195
  steps_taken = step
196
 
197
  rewards.append(reward)
@@ -223,7 +223,7 @@ def run_task(client: OpenAI, env: SQLAnalyticsEnv, task_id: str) -> float:
223
 
224
  def main() -> None:
225
  client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
226
- env = SQLAnalyticsEnv(base_url=ENV_BASE_URL)
227
 
228
  scores = {}
229
  try:
 
141
 
142
  # ── Agent loop ────────────────────────────────────────────────────────────────
143
 
144
+ def run_task(client: OpenAI, env, task_id: str) -> float:
145
  """Run one episode. Returns the final score in [0, 1]."""
146
+ result = env.reset(task_id=task_id)
147
+ schema = result.observation.db_schema or ""
148
+ task_description = result.observation.task_description or ""
149
 
150
  log_start(task=task_id, env=BENCHMARK, model=MODEL_NAME)
151
 
 
186
  print(f"[DEBUG] No SQL at step={step}, stopping.", file=sys.stderr, flush=True)
187
  break
188
 
189
+ result = env.step(SQLAction(sql=sql))
190
 
191
+ reward = result.reward or 0.0
192
+ done = result.done
193
+ error = result.observation.error
194
+ result_rows = len(result.observation.result) if result.observation.result else 0
195
  steps_taken = step
196
 
197
  rewards.append(reward)
 
223
 
224
  def main() -> None:
225
  client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
226
+ env = SQLAnalyticsEnv(base_url=ENV_BASE_URL).sync()
227
 
228
  scores = {}
229
  try: