Taniieeee83 Claude Sonnet 4.6 commited on
Commit
512b667
·
1 Parent(s): 6b248b4

feat: add synchronous HTTP client (client.py)

Browse files

Thin httpx-based wrapper around the DataCleaning HTTP API exposing
reset() / step() / state() / health() with typed Pydantic return values.
Supports context manager usage for automatic connection cleanup.

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

Files changed (1) hide show
  1. client.py +114 -0
client.py ADDED
@@ -0,0 +1,114 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Synchronous HTTP client for the Data Cleaning OpenEnv environment.
3
+
4
+ Usage
5
+ -----
6
+ from client import DataCleaningEnvClient, DataCleaningAction
7
+
8
+ client = DataCleaningEnvClient(base_url="http://localhost:8000")
9
+
10
+ # Start a new episode (task_id 1/2/3 or omit for round-robin)
11
+ result = client.reset(task_id=1)
12
+ print(result.observation.task_description)
13
+
14
+ # Take a step
15
+ action = DataCleaningAction(
16
+ operation="fill_missing",
17
+ column="salary",
18
+ params={"strategy": "median"},
19
+ )
20
+ result = client.step(action)
21
+ print(result.observation.current_score, result.reward, result.done)
22
+
23
+ # Inspect state
24
+ state = client.state()
25
+ print(state.episode_id, state.errors_remaining)
26
+ """
27
+
28
+ from typing import Optional
29
+ import httpx
30
+ from pydantic import BaseModel
31
+
32
+ from models import DataCleaningAction, DataCleaningObservation, DataCleaningState
33
+
34
+
35
+ class StepResult(BaseModel):
36
+ """Returned by reset() and step()."""
37
+ observation: DataCleaningObservation
38
+ reward: float
39
+ done: bool
40
+ info: dict = {}
41
+
42
+
43
+ class DataCleaningEnvClient:
44
+ """
45
+ Thin synchronous wrapper around the DataCleaning HTTP API.
46
+
47
+ All methods raise httpx.HTTPStatusError on non-2xx responses.
48
+ """
49
+
50
+ def __init__(self, base_url: str = "http://localhost:8000", timeout: float = 30.0):
51
+ self.base_url = base_url.rstrip("/")
52
+ self._client = httpx.Client(base_url=self.base_url, timeout=timeout)
53
+
54
+ # ------------------------------------------------------------------
55
+ # Core API
56
+ # ------------------------------------------------------------------
57
+
58
+ def reset(self, task_id: Optional[int] = None) -> StepResult:
59
+ """
60
+ Start a new episode.
61
+
62
+ Parameters
63
+ ----------
64
+ task_id : int | None
65
+ 1 = Easy (fill missing values)
66
+ 2 = Medium (fix formats + duplicates)
67
+ 3 = Hard (full pipeline)
68
+ None = round-robin (1 → 2 → 3 → 1 …)
69
+ """
70
+ payload = {"task_id": task_id} if task_id is not None else {}
71
+ resp = self._client.post("/reset", json=payload)
72
+ resp.raise_for_status()
73
+ return StepResult(**resp.json())
74
+
75
+ def step(self, action: DataCleaningAction) -> StepResult:
76
+ """
77
+ Apply one cleaning operation and return the updated observation.
78
+
79
+ Parameters
80
+ ----------
81
+ action : DataCleaningAction
82
+ operation : str – one of fill_missing / drop_duplicates /
83
+ fix_format / replace_value / drop_outliers / fix_dtype
84
+ column : str – target column (optional for drop_duplicates)
85
+ params : dict – operation-specific parameters
86
+ """
87
+ resp = self._client.post("/step", json=action.model_dump())
88
+ resp.raise_for_status()
89
+ return StepResult(**resp.json())
90
+
91
+ def state(self) -> DataCleaningState:
92
+ """Return current episode metadata without modifying state."""
93
+ resp = self._client.post("/state")
94
+ resp.raise_for_status()
95
+ return DataCleaningState(**resp.json())
96
+
97
+ def health(self) -> dict:
98
+ """Ping the server. Returns {"status": "ok"} if healthy."""
99
+ resp = self._client.get("/health")
100
+ resp.raise_for_status()
101
+ return resp.json()
102
+
103
+ # ------------------------------------------------------------------
104
+ # Context manager support
105
+ # ------------------------------------------------------------------
106
+
107
+ def __enter__(self):
108
+ return self
109
+
110
+ def __exit__(self, *_):
111
+ self.close()
112
+
113
+ def close(self):
114
+ self._client.close()