Spaces:
Sleeping
Sleeping
File size: 3,822 Bytes
a8a3c90 3867c62 a8a3c90 | 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 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | """QueryForge Environment Client."""
from typing import Any, Dict, List, Optional
import httpx
from openenv.core import EnvClient
from openenv.core.client_types import StepResult
from openenv.core.env_server.types import State
try:
from .models import SQLAction, SQLObservation, TaskSpec
except ImportError:
from models import SQLAction, SQLObservation, TaskSpec
class QueryforgeEnv(EnvClient[SQLAction, SQLObservation, State]):
"""
Client for the QueryForge SQL Debugger & Optimiser environment.
Maintains a persistent WebSocket connection to the environment server.
Each client instance has its own dedicated session (isolated task state).
Example:
>>> with QueryforgeEnv(base_url="http://localhost:8000") as env:
... obs = env.reset()
... print(obs.task_title)
... print(obs.task_description)
...
... result = env.step(SQLAction(sql="SELECT name, age FROM users WHERE age > 30"))
... print(result.reward, result.observation.feedback)
Example with Docker:
>>> env = QueryforgeEnv.from_docker_image("queryforge-env:latest")
>>> try:
... obs = env.reset()
... result = env.step(SQLAction(sql="SELECT ..."))
... finally:
... env.close()
"""
def _step_payload(self, action: SQLAction) -> Dict:
return {"sql": action.sql}
def _parse_result(self, payload: Dict) -> StepResult[SQLObservation]:
obs_data = payload.get("observation", {})
observation = SQLObservation(
task_id=obs_data.get("task_id", ""),
task_level=obs_data.get("task_level", ""),
task_title=obs_data.get("task_title", ""),
task_description=obs_data.get("task_description", ""),
syntax_valid=obs_data.get("syntax_valid", False),
execution_success=obs_data.get("execution_success", False),
execution_error=obs_data.get("execution_error"),
rows_returned=obs_data.get("rows_returned", 0),
feedback=obs_data.get("feedback", ""),
hint=obs_data.get("hint", ""),
attempt=obs_data.get("attempt", 0),
best_score=obs_data.get("best_score", 0.0),
done=payload.get("done", False),
reward=payload.get("reward", 0.0),
metadata=obs_data.get("metadata", {}),
)
return StepResult(
observation=observation,
reward=payload.get("reward", 0.0),
done=payload.get("done", False),
)
def _parse_state(self, payload: Dict) -> State:
return State(
episode_id=payload.get("episode_id"),
step_count=payload.get("step_count", 0),
)
# ── Task Registry helpers ─────────────────────────────────────────────────
def register_task(self, spec: TaskSpec) -> Dict[str, Any]:
"""Register a custom task on the server. Returns the server response dict."""
resp = httpx.post(
f"{self.base_url}/tasks",
json=spec.model_dump(),
timeout=10,
)
resp.raise_for_status()
return resp.json()
def list_tasks(self) -> List[Dict[str, Any]]:
"""Return all registered tasks (built-in + custom) as a list of dicts."""
resp = httpx.get(f"{self.base_url}/tasks", timeout=10)
resp.raise_for_status()
return resp.json()
def delete_task(self, task_id: str) -> Dict[str, Any]:
"""Remove a custom task by ID. Raises httpx.HTTPStatusError on 403/404."""
resp = httpx.delete(f"{self.base_url}/tasks/{task_id}", timeout=10)
resp.raise_for_status()
return resp.json()
|