Spaces:
Sleeping
Sleeping
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>
- client.py +20 -4
- inference.py +10 -10
client.py
CHANGED
|
@@ -36,7 +36,9 @@ Example usage
|
|
| 36 |
env.close()
|
| 37 |
"""
|
| 38 |
|
| 39 |
-
from
|
|
|
|
|
|
|
| 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 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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
|
| 145 |
"""Run one episode. Returns the final score in [0, 1]."""
|
| 146 |
-
|
| 147 |
-
schema =
|
| 148 |
-
task_description =
|
| 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 |
-
|
| 190 |
|
| 191 |
-
reward =
|
| 192 |
-
done =
|
| 193 |
-
error =
|
| 194 |
-
result_rows = len(
|
| 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:
|