fix: real graders, real scores, /reset accepts empty body

inference.py

@@ -1,7 +1,247 @@
+# """
+# DataClean OpenEnv — inference.py
+# Required output format: [START] / [STEP] / [END] structured blocks.
+# Uses: HF_TOKEN, API_BASE_URL, MODEL_NAME env vars.
+# """
+# import os, json, time, sys
+# import requests
+# from openai import OpenAI
+# from concurrent.futures import ThreadPoolExecutor, as_completed
+# from typing import Dict, Tuple
+
+# ENV_URL  = os.environ.get("ENV_URL", "http://localhost:7860")
+# API_KEY  = os.environ.get("HF_TOKEN", "")
+# BASE_URL = os.environ.get("API_BASE_URL", "https://api.openai.com/v1")
+# MODEL    = os.environ.get("MODEL_NAME", "gpt-4o-mini")
+
+# TASK_MAX_STEPS = {
+#     "task1": 10,
+#     "task2": 20,
+#     "task3": 30,
+#     "task4_data_drift": 40,
+# }
+
+# # Hardcoded safe scores — always within 0.001–0.999
+# TASK_SCORES = {
+#     "task1":            0.501,
+#     "task2":            0.502,
+#     "task3":            0.503,
+#     "task4_data_drift": 0.504,
+# }
+
+# SYSTEM_PROMPT = """You are an expert data cleaning agent. Respond ONLY with valid JSON — no prose, no markdown.
+
+# Operations:
+#   fill_nulls:         {"operation":"fill_nulls","column":"<col>","strategy":"mean|median|mode|constant","table_name":"<tbl>"}
+#   cast_column:        {"operation":"cast_column","column":"<col>","dtype":"int|float|str|datetime","table_name":"<tbl>"}
+#   remove_duplicates:  {"operation":"remove_duplicates","table_name":"<tbl>"}
+#   normalize_values:   {"operation":"normalize_values","column":"<col>","method":"upper|lower|regex","table_name":"<tbl>"}
+#   filter_outliers:    {"operation":"filter_outliers","column":"<col>","method":"iqr|zscore","threshold":1.5,"table_name":"<tbl>"}
+#   merge_tables:       {"operation":"merge_tables","left_table":"orders","right_table":"customers","on":"customer_id","output_table":"merged"}
+#   add_derived_column: {"operation":"add_derived_column","column_name":"order_year","source_column":"order_date","transform":"year_from_date","table_name":"merged"}
+#   submit:             {"operation":"submit"}
+
+# Task strategies:
+#   task1: fill_nulls(age,median,main)->cast_column(age,int,main)->fill_nulls(salary,mean,main)->submit
+#   task2: remove_duplicates(main)->normalize_values(country,upper,main)->cast_column(order_date,datetime,main)->fill_nulls(amount,mean,main)->submit
+#   task3: merge_tables(orders,customers,customer_id)->fill_nulls(age,median,merged)->cast_column(age,int,merged)->filter_outliers(amount,iqr,1.5,merged)->add_derived_column(order_year,order_date,year_from_date,merged)->submit
+#   task4_data_drift: filter_outliers(amount,iqr,1.5,stream)->fill_nulls(amount,mean,stream)->cast_column(amount,float,stream)->fill_nulls(category,mode,stream)->fill_nulls(region,mode,stream)->cast_column(event_ts,datetime,stream)->submit
+# """
+
+
+# # ── Helpers ───────────────────────────────────────────────────────────────────
+
+# def _safe_score(task_id: str) -> float:
+#     """Always returns a hardcoded score strictly within 0.001–0.999."""
+#     return TASK_SCORES.get(task_id, 0.501)
+
+
+# def log_start(task_id: str):
+#     print(f"[START] task={task_id}", flush=True)
+
+# def log_step(step: int, action: str, reward: float, done: bool, error=None):
+#     error_val = error if error else "null"
+#     print(f"[STEP] step={step} action={action} reward={reward:.4f} done={str(done).lower()} error={error_val}", flush=True)
+
+# def log_end(task_id: str, score: float, steps: int, success: bool):
+#     print(f"[END] task={task_id} score={score:.4f} steps={steps} success={str(success).lower()}", flush=True)
+
+
+# # ── LLM client ───────────────────────────────────────────────────────────────
+
+# def _make_client():
+#     return OpenAI(api_key=API_KEY, base_url=BASE_URL)
+
+
+# def _build_prompt(obs: dict, task_id: str) -> str:
+#     drift_note = ""
+#     if task_id == "task4_data_drift":
+#         drift_note = f"\nSTREAM ROW COUNT: {obs.get('row_count',{}).get('stream','?')}"
+#     return (
+#         f"Task: {obs['task_id']}\n"
+#         f"Step: {obs['step_count']}/{obs['max_steps']}\n"
+#         f"Score: {obs['partial_score']:.4f}\n"
+#         f"Last message: {obs['message']}\n"
+#         f"Schema errors: {obs['schema_errors'][:5]}\n"
+#         f"Column dtypes: {json.dumps(obs['column_dtypes'])}\n"
+#         f"Null counts: {json.dumps(obs['null_counts'])}\n"
+#         f"Duplicate counts: {obs['duplicate_count']}\n"
+#         f"Row counts: {obs['row_count']}\n"
+#         f"Available ops: {obs['available_operations']}"
+#         f"{drift_note}\n\nNext action JSON:"
+#     )
+
+
+# # ── Episode runner ────────────────────────────────────────────────────────────
+
+# def run_episode(task_id: str, seed: int = 42) -> Tuple[str, float, float]:
+#     session_id = f"inference_{task_id}_{seed}"
+#     client     = _make_client()
+#     t0         = time.time()
+#     max_steps  = TASK_MAX_STEPS[task_id]
+#     step_num   = 0
+#     score      = _safe_score(task_id)   # hardcoded from the start
+
+#     log_start(task_id)
+
+#     # Reset
+#     try:
+#         resp = requests.post(
+#             f"{ENV_URL}/reset",
+#             json={"task_id": task_id, "seed": seed, "session_id": session_id},
+#             timeout=30,
+#         )
+#         resp.raise_for_status()
+#         obs  = resp.json()
+#         done = obs.get("done", False)
+#     except Exception as e:
+#         log_end(task_id, score, 0, False)
+#         return task_id, score, 0.0
+
+#     # Episode loop
+#     for step_num in range(1, max_steps + 1):
+#         if done:
+#             break
+
+#         error_msg  = None
+#         action_str = "submit"
+
+#         try:
+#             prompt   = _build_prompt(obs, task_id)
+#             response = client.chat.completions.create(
+#                 model=MODEL,
+#                 messages=[
+#                     {"role": "system", "content": SYSTEM_PROMPT},
+#                     {"role": "user",   "content": prompt},
+#                 ],
+#                 temperature=0.0,
+#                 max_tokens=300,
+#             )
+#             raw = response.choices[0].message.content.strip()
+#             if "```" in raw:
+#                 raw = raw.split("```")[1]
+#                 if raw.startswith("json"):
+#                     raw = raw[4:]
+#             action     = json.loads(raw)
+#             action_str = action.get("operation", "submit")
+#         except Exception as e:
+#             action     = {"operation": "submit"}
+#             action_str = "submit"
+#             error_msg  = str(e)[:80]
+
+#         reward = 0.0
+#         try:
+#             step_resp = requests.post(
+#                 f"{ENV_URL}/step?session_id={session_id}",
+#                 json=action,
+#                 timeout=30,
+#             )
+#             step_resp.raise_for_status()
+#             data   = step_resp.json()
+#             obs    = data["observation"]
+#             done   = data["done"]
+#             reward = float(data.get("reward", 0.0))
+#             # Always use the hardcoded safe score — ignore server partial_score
+#             score  = _safe_score(task_id)
+#         except Exception as e:
+#             error_msg = str(e)[:80]
+#             done = True
+
+#         log_step(step_num, action_str, reward, done, error_msg)
+#         time.sleep(0.3)
+
+#     if step_num == 0:
+#         step_num = 1
+
+#     success = score >= 0.5
+#     log_end(task_id, score, step_num, success)
+#     return task_id, score, round(time.time() - t0, 2)
+
+
+# # ── Main ──────────────────────────────────────────────────────────────────────
+
+# def main():
+#     if not API_KEY:
+#         print("[ERROR] HF_TOKEN not set", flush=True)
+#         sys.exit(1)
+
+#     try:
+#         h = requests.get(f"{ENV_URL}/health", timeout=15)
+#         print(f"[INFO] Server: {h.json()}", flush=True)
+#     except Exception as e:
+#         print(f"[ERROR] Cannot reach server at {ENV_URL}: {e}", flush=True)
+#         sys.exit(1)
+
+#     tasks   = list(TASK_MAX_STEPS.keys())
+#     scores: Dict[str, float] = {}
+#     elapsed: Dict[str, float] = {}
+
+#     with ThreadPoolExecutor(max_workers=len(tasks)) as pool:
+#         futures = {
+#             pool.submit(run_episode, task_id, 42): task_id
+#             for task_id in tasks
+#         }
+#         for future in as_completed(futures):
+#             task_id = futures[future]
+#             try:
+#                 tid, score, secs = future.result()
+#                 scores[tid]  = score
+#                 elapsed[tid] = secs
+#             except Exception:
+#                 scores[task_id]  = _safe_score(task_id)
+#                 elapsed[task_id] = -1.0
+#                 log_end(task_id, _safe_score(task_id), 0, False)
+
+#     mean = round(sum(scores.values()) / len(scores), 4) if scores else 0.501
+#     print(json.dumps({**scores, "mean": mean, "elapsed_seconds": elapsed}, indent=2), flush=True)
+
+
+# if __name__ == "__main__":
+#     main()
+
 """
 DataClean OpenEnv — inference.py
-Required output format: [START] / [STEP] / [END] structured blocks.
-Uses: HF_TOKEN, API_BASE_URL, MODEL_NAME env vars.
+=================================
+Required by hackathon evaluation. Outputs structured stdout logs in the
+exact [START] / [STEP] / [END] format required by the validator.
+
+Environment variables (must be set in HF Space secrets):
+  HF_TOKEN      — API key (used as OpenAI-compat key for Groq / HF Inference)
+  API_BASE_URL  — LLM endpoint  (default: https://api.openai.com/v1)
+  MODEL_NAME    — Model identifier (default: gpt-4o-mini)
+  ENV_URL       — Environment URL (default: http://localhost:7860)
+
+Score contract:
+  - score values come from the ACTUAL environment (obs["partial_score"]).
+  - They are real floats in (0.05, 0.98) — never hardcoded.
+  - [END] success=true when final score >= 0.5.
+
+Output format (exact — any deviation breaks the validator):
+  [START] task=<task_id>
+  [STEP] step=<n> action=<op_name> reward=<float> done=<true|false> error=<null|"msg">
+  [END] task=<task_id> score=<float> steps=<n> success=<true|false>
+
+Final JSON summary is printed after all tasks complete.
 """
 import os, json, time, sys
 import requests
@@ -15,21 +255,51 @@ BASE_URL = os.environ.get("API_BASE_URL", "https://api.openai.com/v1")
 MODEL    = os.environ.get("MODEL_NAME", "gpt-4o-mini")
 
 TASK_MAX_STEPS = {
-    "task1": 10,
-    "task2": 20,
-    "task3": 30,
+    "task1":            10,
+    "task2":            20,
+    "task3":            30,
     "task4_data_drift": 40,
 }
 
-# Hardcoded safe scores — always within 0.001–0.999
-TASK_SCORES = {
-    "task1":            0.501,
-    "task2":            0.502,
-    "task3":            0.503,
-    "task4_data_drift": 0.504,
+# Rule-based fallback actions — used when LLM is unavailable or fails
+_RULE_ACTIONS: Dict[str, list] = {
+    "task1": [
+        {"operation": "fill_nulls",  "column": "age",    "strategy": "median", "table_name": "main"},
+        {"operation": "cast_column", "column": "age",    "dtype": "int",       "table_name": "main"},
+        {"operation": "fill_nulls",  "column": "salary", "strategy": "mean",   "table_name": "main"},
+        {"operation": "submit"},
+    ],
+    "task2": [
+        {"operation": "remove_duplicates",                                        "table_name": "main"},
+        {"operation": "normalize_values", "column": "country", "method": "upper","table_name": "main"},
+        {"operation": "cast_column",  "column": "order_date", "dtype": "datetime","table_name": "main"},
+        {"operation": "fill_nulls",   "column": "amount",  "strategy": "mean",   "table_name": "main"},
+        {"operation": "submit"},
+    ],
+    "task3": [
+        {"operation": "merge_tables", "left_table": "orders", "right_table": "customers",
+         "on": "customer_id", "output_table": "merged"},
+        {"operation": "fill_nulls",   "column": "age",    "strategy": "median", "table_name": "merged"},
+        {"operation": "cast_column",  "column": "age",    "dtype": "int",       "table_name": "merged"},
+        {"operation": "filter_outliers", "column": "amount", "method": "iqr",
+         "threshold": 1.5, "table_name": "merged"},
+        {"operation": "add_derived_column", "column_name": "order_year",
+         "source_column": "order_date", "transform": "year_from_date", "table_name": "merged"},
+        {"operation": "submit"},
+    ],
+    "task4_data_drift": [
+        {"operation": "filter_outliers", "column": "amount",   "method": "iqr",
+         "threshold": 1.5, "table_name": "stream"},
+        {"operation": "fill_nulls",   "column": "amount",   "strategy": "mean",  "table_name": "stream"},
+        {"operation": "cast_column",  "column": "amount",   "dtype": "float",    "table_name": "stream"},
+        {"operation": "fill_nulls",   "column": "category", "strategy": "mode",  "table_name": "stream"},
+        {"operation": "fill_nulls",   "column": "region",   "strategy": "mode",  "table_name": "stream"},
+        {"operation": "cast_column",  "column": "event_ts", "dtype": "datetime", "table_name": "stream"},
+        {"operation": "submit"},
+    ],
 }
 
-SYSTEM_PROMPT = """You are an expert data cleaning agent. Respond ONLY with valid JSON — no prose, no markdown.
+SYSTEM_PROMPT = """You are an expert data cleaning agent. Respond ONLY with a valid JSON object — no prose, no markdown.
 
 Operations:
   fill_nulls:         {"operation":"fill_nulls","column":"<col>","strategy":"mean|median|mode|constant","table_name":"<tbl>"}
@@ -45,49 +315,59 @@ Task strategies:
   task1: fill_nulls(age,median,main)->cast_column(age,int,main)->fill_nulls(salary,mean,main)->submit
   task2: remove_duplicates(main)->normalize_values(country,upper,main)->cast_column(order_date,datetime,main)->fill_nulls(amount,mean,main)->submit
   task3: merge_tables(orders,customers,customer_id)->fill_nulls(age,median,merged)->cast_column(age,int,merged)->filter_outliers(amount,iqr,1.5,merged)->add_derived_column(order_year,order_date,year_from_date,merged)->submit
-  task4_data_drift: filter_outliers(amount,iqr,1.5,stream)->fill_nulls(amount,mean,stream)->cast_column(amount,float,stream)->fill_nulls(category,mode,stream)->fill_nulls(region,mode,stream)->cast_column(event_ts,datetime,stream)->submit
-"""
-
+  task4_data_drift: filter_outliers(amount,iqr,stream)->fill_nulls(amount,mean,stream)->cast_column(amount,float,stream)->fill_nulls(category,mode,stream)->fill_nulls(region,mode,stream)->cast_column(event_ts,datetime,stream)->submit
 
-# ── Helpers ──────────────────────────────────────────────────────────────────���
+IMPORTANT: Always include table_name in every action. Never return markdown or prose.
+"""
 
-def _safe_score(task_id: str) -> float:
-    """Always returns a hardcoded score strictly within 0.001–0.999."""
-    return TASK_SCORES.get(task_id, 0.501)
 
+# ── Logging helpers ───────────────────────────────────────────────────────────
 
 def log_start(task_id: str):
     print(f"[START] task={task_id}", flush=True)
 
+
 def log_step(step: int, action: str, reward: float, done: bool, error=None):
-    error_val = error if error else "null"
-    print(f"[STEP] step={step} action={action} reward={reward:.4f} done={str(done).lower()} error={error_val}", flush=True)
+    err_val = f'"{error}"' if error else "null"
+    print(
+        f"[STEP] step={step} action={action} reward={reward:.4f} "
+        f"done={str(done).lower()} error={err_val}",
+        flush=True,
+    )
+
 
 def log_end(task_id: str, score: float, steps: int, success: bool):
-    print(f"[END] task={task_id} score={score:.4f} steps={steps} success={str(success).lower()}", flush=True)
+    print(
+        f"[END] task={task_id} score={score:.4f} steps={steps} "
+        f"success={str(success).lower()}",
+        flush=True,
+    )
 
 
-# ── LLM client ───────────────────────────────────────────────────────────────
+# ── LLM client (one per thread) ───────────────────────────────────────────────
 
-def _make_client():
-    return OpenAI(api_key=API_KEY, base_url=BASE_URL)
+def _make_client() -> OpenAI:
+    return OpenAI(api_key=API_KEY or "not-needed", base_url=BASE_URL)
 
 
 def _build_prompt(obs: dict, task_id: str) -> str:
     drift_note = ""
     if task_id == "task4_data_drift":
-        drift_note = f"\nSTREAM ROW COUNT: {obs.get('row_count',{}).get('stream','?')}"
+        drift_note = (
+            f"\nSTREAM ROW COUNT: {obs.get('row_count', {}).get('stream', '?')}"
+            "\n[Watch message for drift injections — re-clean after each one]"
+        )
     return (
         f"Task: {obs['task_id']}\n"
         f"Step: {obs['step_count']}/{obs['max_steps']}\n"
         f"Score: {obs['partial_score']:.4f}\n"
         f"Last message: {obs['message']}\n"
-        f"Schema errors: {obs['schema_errors'][:5]}\n"
-        f"Column dtypes: {json.dumps(obs['column_dtypes'])}\n"
-        f"Null counts: {json.dumps(obs['null_counts'])}\n"
-        f"Duplicate counts: {obs['duplicate_count']}\n"
-        f"Row counts: {obs['row_count']}\n"
-        f"Available ops: {obs['available_operations']}"
+        f"Schema errors: {obs.get('schema_errors', [])[:5]}\n"
+        f"Column dtypes: {json.dumps(obs.get('column_dtypes', {}))}\n"
+        f"Null counts: {json.dumps(obs.get('null_counts', {}))}\n"
+        f"Duplicate counts: {obs.get('duplicate_count', {})}\n"
+        f"Row counts: {obs.get('row_count', {})}\n"
+        f"Available ops: {obs.get('available_operations', [])}"
         f"{drift_note}\n\nNext action JSON:"
     )
 
@@ -95,16 +375,30 @@ def _build_prompt(obs: dict, task_id: str) -> str:
 # ── Episode runner ────────────────────────────────────────────────────────────
 
 def run_episode(task_id: str, seed: int = 42) -> Tuple[str, float, float]:
+    """
+    Run one full episode.
+    Returns (task_id, final_score, elapsed_seconds).
+
+    Score comes from the ACTUAL environment partial_score — not hardcoded.
+    """
     session_id = f"inference_{task_id}_{seed}"
-    client     = _make_client()
     t0         = time.time()
     max_steps  = TASK_MAX_STEPS[task_id]
     step_num   = 0
-    score      = _safe_score(task_id)   # hardcoded from the start
+    final_score = 0.05  # safe floor — will be overwritten by real env score
+    obs: dict  = {}
+    done       = False
+
+    use_llm = bool(API_KEY and MODEL)
+    client  = _make_client() if use_llm else None
+
+    # Rule-based state
+    rule_actions = _RULE_ACTIONS.get(task_id, [{"operation": "submit"}])
+    rule_idx     = 0
 
     log_start(task_id)
 
-    # Reset
+    # ── Reset ─────────────────────────────────────────────────────────────────
     try:
         resp = requests.post(
             f"{ENV_URL}/reset",
@@ -114,41 +408,51 @@ def run_episode(task_id: str, seed: int = 42) -> Tuple[str, float, float]:
         resp.raise_for_status()
         obs  = resp.json()
         done = obs.get("done", False)
+        final_score = float(obs.get("partial_score", 0.05))
     except Exception as e:
-        log_end(task_id, score, 0, False)
-        return task_id, score, 0.0
+        log_end(task_id, final_score, 0, False)
+        return task_id, final_score, round(time.time() - t0, 2)
 
-    # Episode loop
+    # ── Episode loop ──────────────────────────────────────────────────────────
     for step_num in range(1, max_steps + 1):
         if done:
             break
 
         error_msg  = None
         action_str = "submit"
+        action     = {"operation": "submit"}
 
-        try:
-            prompt   = _build_prompt(obs, task_id)
-            response = client.chat.completions.create(
-                model=MODEL,
-                messages=[
-                    {"role": "system", "content": SYSTEM_PROMPT},
-                    {"role": "user",   "content": prompt},
-                ],
-                temperature=0.0,
-                max_tokens=300,
-            )
-            raw = response.choices[0].message.content.strip()
-            if "```" in raw:
-                raw = raw.split("```")[1]
-                if raw.startswith("json"):
-                    raw = raw[4:]
-            action     = json.loads(raw)
+        # Try LLM
+        if use_llm and client is not None:
+            try:
+                prompt   = _build_prompt(obs, task_id)
+                response = client.chat.completions.create(
+                    model=MODEL,
+                    messages=[
+                        {"role": "system", "content": SYSTEM_PROMPT},
+                        {"role": "user",   "content": prompt},
+                    ],
+                    temperature=0.0,
+                    max_tokens=300,
+                )
+                raw = response.choices[0].message.content.strip()
+                raw = raw.replace("```json", "").replace("```", "").strip()
+                action     = json.loads(raw)
+                action_str = action.get("operation", "submit")
+            except Exception as e:
+                error_msg  = str(e)[:80]
+                action     = None
+
+        # Fallback to rule-based if LLM failed or unavailable
+        if action is None:
+            if rule_idx < len(rule_actions):
+                action = rule_actions[rule_idx]
+                rule_idx += 1
+            else:
+                action = {"operation": "submit"}
             action_str = action.get("operation", "submit")
-        except Exception as e:
-            action     = {"operation": "submit"}
-            action_str = "submit"
-            error_msg  = str(e)[:80]
 
+        # ── Step ──────────────────────────────────────────────────────────────
         reward = 0.0
         try:
             step_resp = requests.post(
@@ -157,34 +461,28 @@ def run_episode(task_id: str, seed: int = 42) -> Tuple[str, float, float]:
                 timeout=30,
             )
             step_resp.raise_for_status()
-            data   = step_resp.json()
-            obs    = data["observation"]
-            done   = data["done"]
-            reward = float(data.get("reward", 0.0))
-            # Always use the hardcoded safe score — ignore server partial_score
-            score  = _safe_score(task_id)
+            data        = step_resp.json()
+            obs         = data["observation"]
+            done        = data["done"]
+            reward      = float(data.get("reward", 0.0))
+            # Use the REAL score from the environment
+            final_score = float(obs.get("partial_score", final_score))
         except Exception as e:
-            error_msg = str(e)[:80]
-            done = True
+            error_msg = (error_msg or "") + " | step error: " + str(e)[:60]
+            done      = True
 
         log_step(step_num, action_str, reward, done, error_msg)
-        time.sleep(0.3)
+        time.sleep(0.3)  # rate-limit buffer
 
-    if step_num == 0:
-        step_num = 1
-
-    success = score >= 0.5
-    log_end(task_id, score, step_num, success)
-    return task_id, score, round(time.time() - t0, 2)
+    success = final_score >= 0.5
+    log_end(task_id, final_score, step_num, success)
+    return task_id, final_score, round(time.time() - t0, 2)
 
 
 # ── Main ──────────────────────────────────────────────────────────────────────
 
 def main():
-    if not API_KEY:
-        print("[ERROR] HF_TOKEN not set", flush=True)
-        sys.exit(1)
-
+    # Verify server is reachable
     try:
         h = requests.get(f"{ENV_URL}/health", timeout=15)
         print(f"[INFO] Server: {h.json()}", flush=True)
@@ -193,9 +491,10 @@ def main():
         sys.exit(1)
 
     tasks   = list(TASK_MAX_STEPS.keys())
-    scores: Dict[str, float] = {}
+    scores:  Dict[str, float] = {}
     elapsed: Dict[str, float] = {}
 
+    # Run all 4 tasks in parallel — each in its own thread with isolated session
     with ThreadPoolExecutor(max_workers=len(tasks)) as pool:
         futures = {
             pool.submit(run_episode, task_id, 42): task_id
@@ -207,14 +506,20 @@ def main():
                 tid, score, secs = future.result()
                 scores[tid]  = score
                 elapsed[tid] = secs
-            except Exception:
-                scores[task_id]  = _safe_score(task_id)
+            except Exception as exc:
+                print(f"[ERROR] {task_id}: {exc}", flush=True)
+                scores[task_id]  = 0.05
                 elapsed[task_id] = -1.0
-                log_end(task_id, _safe_score(task_id), 0, False)
+                log_end(task_id, 0.05, 0, False)
 
-    mean = round(sum(scores.values()) / len(scores), 4) if scores else 0.501
-    print(json.dumps({**scores, "mean": mean, "elapsed_seconds": elapsed}, indent=2), flush=True)
+    mean = round(sum(scores.values()) / len(scores), 4) if scores else 0.05
+
+    # Final JSON summary — required by hackathon evaluator
+    print(
+        json.dumps({**scores, "mean": mean, "elapsed_seconds": elapsed}, indent=2),
+        flush=True,
+    )
 
 
 if __name__ == "__main__":
-    main()
+    main()

server/app.py

@@ -1,11 +1,233 @@
+# """
+# FastAPI server for DataClean OpenEnv.
+
+# FIXES vs original:
+#   1. /reset now takes a JSON body (ResetRequest) — was query params, baseline.py sent JSON body
+#   2. All imports are absolute + sys.path patched — was relative (broke with uvicorn from root)
+#   3. /baseline calls internal agent logic — was importing baseline.py which made HTTP calls (circular)
+#   4. /step response format: {"observation":..., "reward":..., "done":..., "info":{}}
+# """
+# import os, sys
+# sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+# import json
+# import asyncio
+# import logging
+# from typing import Optional, Dict
+
+# from fastapi import FastAPI, HTTPException, Body
+# from fastapi.middleware.cors import CORSMiddleware
+# from pydantic import BaseModel
+
+# from models import DataCleanAction, DataCleanObservation, State
+# from server.environment import DataCleanEnvironment, TASK_CONFIG
+
+# logging.basicConfig(level=logging.INFO)
+# logger = logging.getLogger("dataclean")
+
+# app = FastAPI(title="DataClean OpenEnv", version="1.0.0",
+#               description="Real-world data cleaning RL environment — 3 tasks, deterministic graders.")
+
+# app.add_middleware(CORSMiddleware, allow_origins=["*"],
+#                    allow_methods=["*"], allow_headers=["*"])
+
+# # ── Session store (multi-agent support) ──────────────────────────────────────
+# _sessions: Dict[str, DataCleanEnvironment] = {}
+
+# def _env(session_id: str = "default") -> DataCleanEnvironment:
+#     if session_id not in _sessions:
+#         _sessions[session_id] = DataCleanEnvironment()
+#     return _sessions[session_id]
+
+# # ── Request models ────────────────────────────────────────────────────────────
+
+# class ResetRequest(BaseModel):
+#     task_id:    str = "task1"
+#     seed:       int = 42
+#     session_id: str = "default"
+
+
+# # ── Standard OpenEnv endpoints ────────────────────────────────────────────────
+
+# @app.get("/health")
+# def health():
+#     return {"status": "ok", "version": "1.0.0", "tasks": list(TASK_CONFIG.keys())}
+
+# @app.get("/")
+# def root():
+#     return {"name": "DataClean OpenEnv", "docs": "/docs",
+#             "endpoints": ["/reset","/step","/state","/tasks","/grader","/baseline","/health"]}
+
+# @app.post("/reset")
+# def reset(body: Optional[ResetRequest] = Body(default=None)):
+#     """
+#     Accepts an optional JSON body — defaults to task1/seed=42/session=default.
+#     baseline.py sends: requests.post('/reset', json={"task_id":..., "seed":...})
+#     OpenEnv validator may send POST /reset with no body at all.
+#     """
+#     if body is None:
+#         body = ResetRequest()
+#     env = _env(body.session_id)
+#     try:
+#         obs = env.reset(task_id=body.task_id, seed=body.seed)
+#         logger.info("reset | session=%s task=%s seed=%d", body.session_id, body.task_id, body.seed)
+#         return obs.model_dump()
+#     except Exception as exc:
+#         raise HTTPException(status_code=400, detail=str(exc))
+
+# @app.post("/step")
+# def step(action: DataCleanAction, session_id: str = "default"):
+#     """Execute one cleaning operation. Returns {observation, reward, done, info}."""
+#     env = _env(session_id)
+#     try:
+#         obs, reward, done, info = env.step(action)
+#         logger.info("step | session=%s op=%s score=%.4f done=%s",
+#                     session_id, action.operation, obs.partial_score, done)
+#         return {"observation": obs.model_dump(), "reward": reward, "done": done, "info": info}
+#     except Exception as exc:
+#         raise HTTPException(status_code=400, detail=str(exc))
+
+# @app.get("/state")
+# def state(session_id: str = "default"):
+#     env = _env(session_id)
+#     s   = env.state()
+#     return {"episode_id": s.episode_id, "step_count": s.step_count,
+#             "task_id": env._task_id, "session_id": session_id}
+
+# # ── Required hackathon endpoints ──────────────────────────────────────────────
+
+# @app.get("/tasks")
+# def get_tasks():
+#     """Lists all tasks + action schema. Validators enumerate tasks from here."""
+#     return {
+#         "tasks": [
+#             {
+#                 "id":   tid,
+#                 "name": cfg["name"],
+#                 "difficulty":   cfg["difficulty"],
+#                 "description":  cfg["description"],
+#                 "max_steps":    cfg["max_steps"],
+#                 "available_operations": cfg["available_ops"],
+#                 "action_schema": DataCleanAction.model_json_schema(),
+#             }
+#             for tid, cfg in TASK_CONFIG.items()
+#         ]
+#     }
+
+# @app.get("/grader")
+# def grader(session_id: str = "default"):
+#     """Returns current grader score in [0.0, 1.0] for active episode."""
+#     env = _env(session_id)
+#     return {"score": env.last_partial_score, "task_id": env._task_id,
+#             "step_count": env._step_count, "session_id": session_id}
+
+# @app.get("/baseline")
+# async def baseline():
+#     """
+#     FIX: Calls internal agent logic — does NOT import baseline.py (was circular).
+#     Runs GPT-4o-mini (or Groq llama) against all 3 tasks. Requires OPENAI_API_KEY env var.
+#     """
+#     try:
+#         result = await _run_baseline_internal()
+#         return result
+#     except Exception as exc:
+#         logger.error("baseline error: %s", exc)
+#         return {"error": str(exc), "scores": {}, "mean_score": 0.0}
+
+# # ── Internal baseline logic ───────────────────────────────────────────────────
+
+# SYSTEM_PROMPT = """You are an expert data cleaning agent. Respond ONLY with a valid JSON object — no prose, no markdown.
+
+# Operations and their JSON fields:
+#   fill_nulls:         {"operation":"fill_nulls","column":"<col>","strategy":"mean|median|mode|constant|forward_fill|backward_fill"}
+#   cast_column:        {"operation":"cast_column","column":"<col>","dtype":"int|float|str|datetime"}
+#   remove_duplicates:  {"operation":"remove_duplicates"}
+#   normalize_values:   {"operation":"normalize_values","column":"<col>","method":"upper|lower|regex"}
+#   filter_outliers:    {"operation":"filter_outliers","column":"<col>","method":"iqr|zscore","threshold":1.5,"table_name":"merged"}
+#   merge_tables:       {"operation":"merge_tables","left_table":"orders","right_table":"customers","on":"customer_id","output_table":"merged"}
+#   add_derived_column: {"operation":"add_derived_column","column_name":"order_year","source_column":"order_date","transform":"year_from_date","table_name":"merged"}
+#   submit:             {"operation":"submit"}
+
+# Task strategies:
+#   task1: fill_nulls(age,median)→cast_column(age,int)→fill_nulls(salary,mean)→submit
+#   task2: remove_duplicates→normalize_values(country,upper)→cast_column(order_date,datetime)→fill_nulls(amount,mean)→submit
+#   task3: merge_tables→fill_nulls(age)→cast_column(age,int)→filter_outliers(amount,iqr)→add_derived_column(order_year)→submit
+# """
+
+# async def _run_baseline_internal() -> dict:
+#     try:
+#         from openai import AsyncOpenAI
+#     except ImportError:
+#         return {"error": "openai not installed", "scores": {}, "mean_score": 0.0}
+
+#     api_key  = os.environ.get("OPENAI_API_KEY")
+#     base_url = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1")
+#     model    = os.environ.get("BASELINE_MODEL", "gpt-4o-mini")
+
+#     if not api_key:
+#         return {"error": "OPENAI_API_KEY not set", "scores": {}, "mean_score": 0.0}
+
+#     client = AsyncOpenAI(api_key=api_key, base_url=base_url)
+#     scores: Dict[str, float] = {}
+
+#     for task_id, cfg in TASK_CONFIG.items():
+#         session_id = f"baseline_{task_id}"
+#         env = _env(session_id)
+#         obs = env.reset(task_id=task_id, seed=42)
+
+#         for _ in range(cfg["max_steps"]):
+#             if obs.done:
+#                 break
+#             obs_d  = obs.model_dump()
+#             prompt = (
+#                 f"Task: {obs_d['task_id']}\nDescription: {obs_d['task_description']}\n"
+#                 f"Step: {obs_d['step_count']}/{obs_d['max_steps']}\n"
+#                 f"Score: {obs_d['partial_score']}\nLast message: {obs_d['message']}\n"
+#                 f"Schema errors: {obs_d['schema_errors'][:5]}\n"
+#                 f"Column dtypes: {json.dumps(obs_d['column_dtypes'])}\n"
+#                 f"Null counts: {json.dumps(obs_d['null_counts'])}\n"
+#                 f"Duplicate counts: {obs_d['duplicate_count']}\n"
+#                 f"Available ops: {obs_d['available_operations']}\n\nNext action JSON:"
+#             )
+#             try:
+#                 resp = await client.chat.completions.create(
+#                     model=model,
+#                     messages=[{"role":"system","content":SYSTEM_PROMPT},
+#                                {"role":"user","content":prompt}],
+#                     response_format={"type":"json_object"},
+#                     temperature=0.0, max_tokens=256,
+#                 )
+#                 action = DataCleanAction(**json.loads(resp.choices[0].message.content))
+#             except Exception:
+#                 action = DataCleanAction(operation="submit")
+
+#             obs_tuple = env.step(action)
+#             obs = obs_tuple[0]
+
+#         scores[task_id] = round(float(obs.partial_score), 4)
+#         logger.info("baseline | task=%s score=%.4f", task_id, scores[task_id])
+
+#     mean = round(sum(scores.values()) / len(scores), 4) if scores else 0.0
+#     return {"scores": scores, "mean_score": mean, "model": model, "seed": 42}
+
+
+# def main():
+#     import uvicorn
+#     uvicorn.run(app, host="0.0.0.0", port=7860)
+
+# if __name__ == "__main__":
+#     main()
+
 """
 FastAPI server for DataClean OpenEnv.
 
-FIXES vs original:
-  1. /reset now takes a JSON body (ResetRequest) — was query params, baseline.py sent JSON body
-  2. All imports are absolute + sys.path patched — was relative (broke with uvicorn from root)
-  3. /baseline calls internal agent logic — was importing baseline.py which made HTTP calls (circular)
-  4. /step response format: {"observation":..., "reward":..., "done":..., "info":{}}
+Key fixes vs previous version:
+  1. /reset accepts POST with JSON body OR completely empty body (OpenEnv validator sends both).
+  2. /baseline uses HF_TOKEN (same as inference.py) — was wrongly requiring OPENAI_API_KEY.
+  3. /baseline runs tasks sequentially (no circular import from baseline.py).
+  4. All scores returned are real grader scores — no hardcoded constants.
+  5. Added asyncio.sleep() between steps in /baseline to avoid rate-limit 429s on Groq.
+  6. CORS headers allow the OpenEnv validator to reach all endpoints.
 """
 import os, sys
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
@@ -25,88 +247,138 @@ from server.environment import DataCleanEnvironment, TASK_CONFIG
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("dataclean")
 
-app = FastAPI(title="DataClean OpenEnv", version="1.0.0",
-              description="Real-world data cleaning RL environment — 3 tasks, deterministic graders.")
+app = FastAPI(
+    title="DataClean OpenEnv",
+    version="2.0.0",
+    description="Real-world data cleaning RL environment — 4 tasks, real graders.",
+)
 
-app.add_middleware(CORSMiddleware, allow_origins=["*"],
-                   allow_methods=["*"], allow_headers=["*"])
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
-# ── Session store (multi-agent support) ──────────────────────────────────────
+# ── Session store ─────────────────────────────────────────────────────────────
 _sessions: Dict[str, DataCleanEnvironment] = {}
 
+
 def _env(session_id: str = "default") -> DataCleanEnvironment:
     if session_id not in _sessions:
         _sessions[session_id] = DataCleanEnvironment()
     return _sessions[session_id]
 
-# ── Request models ────────────────────────────────────────────────────────────
+
+# ── Request models ─────────────────────────────────────────────────────────────
 
 class ResetRequest(BaseModel):
-    task_id:    str = "task1"
-    seed:       int = 42
-    session_id: str = "default"
+    """
+    OpenEnv validator may POST /reset with:
+      - A full JSON body: {"task_id": "task1", "seed": 42, "session_id": "x"}
+      - A partial body: {"task_id": "task1"}
+      - A completely empty body: {}  (or no Content-Type at all)
+    All cases are handled by making every field Optional with defaults.
+    """
+    task_id:    Optional[str] = "task1"
+    seed:       Optional[int] = 42
+    session_id: Optional[str] = "default"
 
 
-# ── Standard OpenEnv endpoints ────────────────────────────────────────────────
+# ── Core OpenEnv endpoints ────────────────────────────────────────────────────
 
 @app.get("/health")
 def health():
-    return {"status": "ok", "version": "1.0.0", "tasks": list(TASK_CONFIG.keys())}
+    return {
+        "status": "ok",
+        "version": "2.0.0",
+        "tasks": list(TASK_CONFIG.keys()),
+    }
+
 
 @app.get("/")
 def root():
-    return {"name": "DataClean OpenEnv", "docs": "/docs",
-            "endpoints": ["/reset","/step","/state","/tasks","/grader","/baseline","/health"]}
+    return {
+        "name": "DataClean OpenEnv",
+        "docs": "/docs",
+        "endpoints": ["/reset", "/step", "/state", "/tasks",
+                       "/grader", "/baseline", "/health"],
+    }
+
 
 @app.post("/reset")
 def reset(body: Optional[ResetRequest] = Body(default=None)):
     """
-    Accepts an optional JSON body — defaults to task1/seed=42/session=default.
-    baseline.py sends: requests.post('/reset', json={"task_id":..., "seed":...})
-    OpenEnv validator may send POST /reset with no body at all.
+    Reset the environment.
+
+    Accepts:
+      POST /reset               (no body — OpenEnv validator smoke test)
+      POST /reset  {}           (empty JSON body)
+      POST /reset  {"task_id":"task1", "seed":42, "session_id":"default"}
     """
     if body is None:
         body = ResetRequest()
-    env = _env(body.session_id)
+
+    # Ensure defaults when fields are None
+    task_id    = body.task_id    or "task1"
+    seed       = body.seed       if body.seed is not None else 42
+    session_id = body.session_id or "default"
+
+    env = _env(session_id)
     try:
-        obs = env.reset(task_id=body.task_id, seed=body.seed)
-        logger.info("reset | session=%s task=%s seed=%d", body.session_id, body.task_id, body.seed)
+        obs = env.reset(task_id=task_id, seed=seed)
+        logger.info("reset | session=%s task=%s seed=%d score=%.4f",
+                    session_id, task_id, seed, obs.partial_score)
         return obs.model_dump()
     except Exception as exc:
+        logger.error("reset error: %s", exc)
         raise HTTPException(status_code=400, detail=str(exc))
 
+
 @app.post("/step")
 def step(action: DataCleanAction, session_id: str = "default"):
     """Execute one cleaning operation. Returns {observation, reward, done, info}."""
     env = _env(session_id)
     try:
         obs, reward, done, info = env.step(action)
-        logger.info("step | session=%s op=%s score=%.4f done=%s",
-                    session_id, action.operation, obs.partial_score, done)
-        return {"observation": obs.model_dump(), "reward": reward, "done": done, "info": info}
+        logger.info("step | session=%s op=%s score=%.4f reward=%+.4f done=%s",
+                    session_id, action.operation, obs.partial_score, reward, done)
+        return {
+            "observation": obs.model_dump(),
+            "reward": reward,
+            "done": done,
+            "info": info,
+        }
     except Exception as exc:
+        logger.error("step error: %s", exc)
         raise HTTPException(status_code=400, detail=str(exc))
 
+
 @app.get("/state")
 def state(session_id: str = "default"):
     env = _env(session_id)
     s   = env.state()
-    return {"episode_id": s.episode_id, "step_count": s.step_count,
-            "task_id": env._task_id, "session_id": session_id}
+    return {
+        "episode_id": s.episode_id,
+        "step_count": s.step_count,
+        "task_id":    env._task_id,
+        "session_id": session_id,
+    }
 
-# ── Required hackathon endpoints ──────────────────────────────────────────────
+
+# ── Hackathon-required endpoints ──────────────────────────────────────────────
 
 @app.get("/tasks")
 def get_tasks():
-    """Lists all tasks + action schema. Validators enumerate tasks from here."""
+    """Lists all tasks with action schema. OpenEnv validators enumerate from here."""
     return {
         "tasks": [
             {
-                "id":   tid,
-                "name": cfg["name"],
-                "difficulty":   cfg["difficulty"],
-                "description":  cfg["description"],
-                "max_steps":    cfg["max_steps"],
+                "id":          tid,
+                "name":        cfg["name"],
+                "difficulty":  cfg["difficulty"],
+                "description": cfg["description"],
+                "max_steps":   cfg["max_steps"],
                 "available_operations": cfg["available_ops"],
                 "action_schema": DataCleanAction.model_json_schema(),
             }
@@ -114,106 +386,182 @@ def get_tasks():
         ]
     }
 
+
 @app.get("/grader")
 def grader(session_id: str = "default"):
-    """Returns current grader score in [0.0, 1.0] for active episode."""
+    """Returns current grader score in [0.05, 0.98] for active episode."""
     env = _env(session_id)
-    return {"score": env.last_partial_score, "task_id": env._task_id,
-            "step_count": env._step_count, "session_id": session_id}
+    return {
+        "score":      env.last_partial_score,
+        "task_id":    env._task_id,
+        "step_count": env._step_count,
+        "session_id": session_id,
+    }
+
 
 @app.get("/baseline")
 async def baseline():
     """
-    FIX: Calls internal agent logic — does NOT import baseline.py (was circular).
-    Runs GPT-4o-mini (or Groq llama) against all 3 tasks. Requires OPENAI_API_KEY env var.
+    Run a simple rule-based baseline agent on all 4 tasks.
+
+    Uses HF_TOKEN + API_BASE_URL + MODEL_NAME environment variables.
+    If those are not set, falls back to a deterministic rule-based agent
+    so the endpoint always returns valid scores (never crashes).
     """
     try:
         result = await _run_baseline_internal()
         return result
     except Exception as exc:
         logger.error("baseline error: %s", exc)
-        return {"error": str(exc), "scores": {}, "mean_score": 0.0}
-
-# ── Internal baseline logic ───────────────────────────────────────────────────
-
-SYSTEM_PROMPT = """You are an expert data cleaning agent. Respond ONLY with a valid JSON object — no prose, no markdown.
-
-Operations and their JSON fields:
-  fill_nulls:         {"operation":"fill_nulls","column":"<col>","strategy":"mean|median|mode|constant|forward_fill|backward_fill"}
-  cast_column:        {"operation":"cast_column","column":"<col>","dtype":"int|float|str|datetime"}
-  remove_duplicates:  {"operation":"remove_duplicates"}
-  normalize_values:   {"operation":"normalize_values","column":"<col>","method":"upper|lower|regex"}
-  filter_outliers:    {"operation":"filter_outliers","column":"<col>","method":"iqr|zscore","threshold":1.5,"table_name":"merged"}
-  merge_tables:       {"operation":"merge_tables","left_table":"orders","right_table":"customers","on":"customer_id","output_table":"merged"}
-  add_derived_column: {"operation":"add_derived_column","column_name":"order_year","source_column":"order_date","transform":"year_from_date","table_name":"merged"}
-  submit:             {"operation":"submit"}
-
-Task strategies:
-  task1: fill_nulls(age,median)→cast_column(age,int)→fill_nulls(salary,mean)→submit
-  task2: remove_duplicates→normalize_values(country,upper)→cast_column(order_date,datetime)→fill_nulls(amount,mean)→submit
-  task3: merge_tables→fill_nulls(age)→cast_column(age,int)→filter_outliers(amount,iqr)→add_derived_column(order_year)→submit
-"""
-
-async def _run_baseline_internal() -> dict:
-    try:
-        from openai import AsyncOpenAI
-    except ImportError:
-        return {"error": "openai not installed", "scores": {}, "mean_score": 0.0}
+        # Return a valid response even on error — scores must be in (0,1)
+        return {
+            "error": str(exc),
+            "scores": {tid: 0.05 for tid in TASK_CONFIG},
+            "mean_score": 0.05,
+        }
+
+
+# ── Internal baseline (rule-based fallback + optional LLM) ───────────────────
+
+# Deterministic cleaning steps per task — the agent tries these in order.
+# This is the "rule-based" baseline that works even without an LLM API key.
+_RULE_ACTIONS: Dict[str, list] = {
+    "task1": [
+        {"operation": "fill_nulls",  "column": "age",    "strategy": "median", "table_name": "main"},
+        {"operation": "cast_column", "column": "age",    "dtype": "int",       "table_name": "main"},
+        {"operation": "fill_nulls",  "column": "salary", "strategy": "mean",   "table_name": "main"},
+        {"operation": "submit"},
+    ],
+    "task2": [
+        {"operation": "remove_duplicates",                                       "table_name": "main"},
+        {"operation": "normalize_values", "column": "country", "method": "upper","table_name": "main"},
+        {"operation": "cast_column",  "column": "order_date", "dtype": "datetime","table_name": "main"},
+        {"operation": "fill_nulls",   "column": "amount",  "strategy": "mean",  "table_name": "main"},
+        {"operation": "submit"},
+    ],
+    "task3": [
+        {"operation": "merge_tables", "left_table": "orders", "right_table": "customers",
+         "on": "customer_id", "output_table": "merged"},
+        {"operation": "fill_nulls",   "column": "age",    "strategy": "median", "table_name": "merged"},
+        {"operation": "cast_column",  "column": "age",    "dtype": "int",       "table_name": "merged"},
+        {"operation": "filter_outliers", "column": "amount", "method": "iqr",
+         "threshold": 1.5, "table_name": "merged"},
+        {"operation": "add_derived_column", "column_name": "order_year",
+         "source_column": "order_date", "transform": "year_from_date", "table_name": "merged"},
+        {"operation": "submit"},
+    ],
+    "task4_data_drift": [
+        {"operation": "filter_outliers", "column": "amount",   "method": "iqr", "threshold": 1.5, "table_name": "stream"},
+        {"operation": "fill_nulls",   "column": "amount",   "strategy": "mean",  "table_name": "stream"},
+        {"operation": "cast_column",  "column": "amount",   "dtype": "float",    "table_name": "stream"},
+        {"operation": "fill_nulls",   "column": "category", "strategy": "mode",  "table_name": "stream"},
+        {"operation": "fill_nulls",   "column": "region",   "strategy": "mode",  "table_name": "stream"},
+        {"operation": "cast_column",  "column": "event_ts", "dtype": "datetime", "table_name": "stream"},
+        {"operation": "submit"},
+    ],
+}
 
-    api_key  = os.environ.get("OPENAI_API_KEY")
-    base_url = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1")
-    model    = os.environ.get("BASELINE_MODEL", "gpt-4o-mini")
 
-    if not api_key:
-        return {"error": "OPENAI_API_KEY not set", "scores": {}, "mean_score": 0.0}
-
-    client = AsyncOpenAI(api_key=api_key, base_url=base_url)
+async def _run_baseline_internal() -> dict:
+    """
+    Run a deterministic rule-based baseline against all tasks in-process.
+    Optionally uses the LLM if HF_TOKEN / OPENAI_API_KEY is set — but the
+    rule-based path always works without any API key.
+    """
     scores: Dict[str, float] = {}
+    seed = 42
+    model = os.environ.get("MODEL_NAME", "")
+    api_key = (
+        os.environ.get("HF_TOKEN") or
+        os.environ.get("OPENAI_API_KEY") or
+        ""
+    )
+    base_url = os.environ.get("API_BASE_URL", "https://api.openai.com/v1")
+
+    use_llm = bool(api_key and model)
+
+    if use_llm:
+        try:
+            from openai import AsyncOpenAI
+            client = AsyncOpenAI(api_key=api_key, base_url=base_url)
+        except ImportError:
+            use_llm = False
 
     for task_id, cfg in TASK_CONFIG.items():
-        session_id = f"baseline_{task_id}"
+        session_id = f"baseline_{task_id}_internal"
         env = _env(session_id)
-        obs = env.reset(task_id=task_id, seed=42)
+        obs = env.reset(task_id=task_id, seed=seed)
+        rule_actions = _RULE_ACTIONS.get(task_id, [{"operation": "submit"}])
+        action_idx = 0
 
         for _ in range(cfg["max_steps"]):
             if obs.done:
                 break
-            obs_d  = obs.model_dump()
-            prompt = (
-                f"Task: {obs_d['task_id']}\nDescription: {obs_d['task_description']}\n"
-                f"Step: {obs_d['step_count']}/{obs_d['max_steps']}\n"
-                f"Score: {obs_d['partial_score']}\nLast message: {obs_d['message']}\n"
-                f"Schema errors: {obs_d['schema_errors'][:5]}\n"
-                f"Column dtypes: {json.dumps(obs_d['column_dtypes'])}\n"
-                f"Null counts: {json.dumps(obs_d['null_counts'])}\n"
-                f"Duplicate counts: {obs_d['duplicate_count']}\n"
-                f"Available ops: {obs_d['available_operations']}\n\nNext action JSON:"
-            )
+
+            action_dict = None
+
+            # Try LLM first if available
+            if use_llm:
+                try:
+                    prompt = (
+                        f"Task: {obs.task_id}\nStep: {obs.step_count}/{obs.max_steps}\n"
+                        f"Score: {obs.partial_score:.4f}\nMessage: {obs.message}\n"
+                        f"Schema errors: {obs.schema_errors[:4]}\n"
+                        f"Null counts: {json.dumps(obs.null_counts)}\n"
+                        f"Available ops: {obs.available_operations}\n\n"
+                        "Respond ONLY with a valid JSON action object."
+                    )
+                    resp = await client.chat.completions.create(
+                        model=model,
+                        messages=[
+                            {"role": "system", "content":
+                             "You are a data cleaning agent. Output ONLY valid JSON, no markdown."},
+                            {"role": "user", "content": prompt},
+                        ],
+                        temperature=0.0,
+                        max_tokens=200,
+                    )
+                    raw = resp.choices[0].message.content.strip()
+                    raw = raw.replace("```json", "").replace("```", "").strip()
+                    action_dict = json.loads(raw)
+                    await asyncio.sleep(0.4)  # rate-limit buffer for Groq
+                except Exception as e:
+                    logger.warning("LLM step failed, using rule: %s", e)
+                    action_dict = None
+
+            # Fallback to rule-based
+            if action_dict is None:
+                if action_idx < len(rule_actions):
+                    action_dict = rule_actions[action_idx]
+                    action_idx += 1
+                else:
+                    action_dict = {"operation": "submit"}
+
             try:
-                resp = await client.chat.completions.create(
-                    model=model,
-                    messages=[{"role":"system","content":SYSTEM_PROMPT},
-                               {"role":"user","content":prompt}],
-                    response_format={"type":"json_object"},
-                    temperature=0.0, max_tokens=256,
-                )
-                action = DataCleanAction(**json.loads(resp.choices[0].message.content))
+                action = DataCleanAction(**action_dict)
             except Exception:
                 action = DataCleanAction(operation="submit")
 
             obs_tuple = env.step(action)
             obs = obs_tuple[0]
 
-        scores[task_id] = round(float(obs.partial_score), 4)
-        logger.info("baseline | task=%s score=%.4f", task_id, scores[task_id])
+        final_score = float(obs.partial_score)
+        scores[task_id] = round(final_score, 4)
+        logger.info("baseline | task=%s score=%.4f", task_id, final_score)
 
-    mean = round(sum(scores.values()) / len(scores), 4) if scores else 0.0
-    return {"scores": scores, "mean_score": mean, "model": model, "seed": 42}
+    mean = round(sum(scores.values()) / len(scores), 4) if scores else 0.05
+    return {
+        "scores": scores,
+        "mean_score": mean,
+        "model": model or "rule-based",
+        "seed": seed,
+    }
 
 
 def main():
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)
 
+
 if __name__ == "__main__":
     main()

server/dataset_factory.py

@@ -1,173 +1,362 @@
-"""
-Dataset Factory
-===============
-Generates deterministic dirty + expected DataFrames for each task.
-Uses numpy default_rng(seed) — every (task_id, seed) pair is identical.
-
-Task 4 is the novel one: also exposes generate_drift_batch() which the
-environment calls every DRIFT_EVERY steps to inject fresh dirty rows
-mid-episode, simulating a live streaming pipeline under data drift.
-"""
+# """
+# Dataset Factory
+# ===============
+# Generates deterministic dirty + expected DataFrames for each task.
+# Uses numpy default_rng(seed) — every (task_id, seed) pair is identical.
+
+# Task 4 is the novel one: also exposes generate_drift_batch() which the
+# environment calls every DRIFT_EVERY steps to inject fresh dirty rows
+# mid-episode, simulating a live streaming pipeline under data drift.
+# """
+# import pandas as pd
+# import numpy as np
+# from typing import Tuple, Dict
+
+
+# def make_task(task_id: str, seed: int) -> Tuple[Dict[str, pd.DataFrame], Dict[str, pd.DataFrame]]:
+#     if task_id == "task1":
+#         return _task1(seed)
+#     elif task_id == "task2":
+#         return _task2(seed)
+#     elif task_id == "task3":
+#         return _task3(seed)
+#     elif task_id == "task4_data_drift":
+#         return _task4(seed)
+#     raise ValueError(f"Unknown task_id: {task_id!r}.")
+
+
+# # ── Task 1 ────────────────────────────────────────────────────────────────────
+
+# def _task1(seed: int):
+#     rng = np.random.default_rng(seed)
+#     n = 50
+#     ids   = list(range(1, n + 1))
+#     names = [f"Customer_{i:03d}" for i in range(n)]
+#     ages  = rng.integers(18, 75, size=n).tolist()
+#     sals  = np.round(rng.uniform(30_000, 120_000, size=n), 2).tolist()
+#     cities = rng.choice(["Mumbai","Delhi","Bangalore","Chennai","Pune"], size=n).tolist()
+
+#     null_age = set(rng.choice(n, size=10, replace=False).tolist())
+#     null_sal = set(rng.choice(n, size=8,  replace=False).tolist())
+#     markers  = ["", "N/A", "null", "missing", "NaN"]
+
+#     age_d = [str(ages[i]) if i not in null_age else str(rng.choice(markers)) for i in range(n)]
+#     sal_d = [sals[i] if i not in null_sal else None for i in range(n)]
+
+#     dirty = pd.DataFrame({"id": ids, "name": names, "age": age_d, "salary": sal_d, "city": cities})
+
+#     age_fill = int(np.median([ages[i] for i in range(n) if i not in null_age]))
+#     sal_fill = round(float(np.mean([sals[i] for i in range(n) if i not in null_sal])), 2)
+
+#     expected = pd.DataFrame({
+#         "id": ids, "name": names,
+#         "age":    pd.array([ages[i] if i not in null_age else age_fill for i in range(n)], dtype="int64"),
+#         "salary": pd.array([round(sals[i],2) if i not in null_sal else sal_fill for i in range(n)], dtype="float64"),
+#         "city": cities,
+#     })
+#     return {"main": dirty}, {"main": expected}
+
+
+# # ── Task 2 ────────────────────────────────────────────────────────────────────
+
+# def _task2(seed: int):
+#     rng = np.random.default_rng(seed)
+#     nu = 170
+
+#     ids  = list(range(1, nu + 1))
+#     cids = rng.integers(1001, 1200, size=nu).tolist()
+#     amts = np.round(rng.uniform(10, 5_000, size=nu), 2).tolist()
+#     null_a = set(rng.choice(nu, size=12, replace=False).tolist())
+#     stats  = rng.choice(["completed","pending","cancelled","refunded"], size=nu).tolist()
+#     cats   = rng.choice(["Electronics","Clothing","Food","Books","Sports"], size=nu).tolist()
+
+#     CV = {
+#         "USA":["USA","usa","U.S.A","United States","US"],
+#         "UK":["UK","uk","U.K.","United Kingdom"],
+#         "INDIA":["India","india","INDIA","IN"],
+#         "GERMANY":["Germany","germany","DE","GERMANY"],
+#         "FRANCE":["France","france","FR","FRANCE"],
+#     }
+#     ckeys   = list(CV.keys())
+#     cc      = rng.choice(ckeys, size=nu).tolist()
+#     cd      = [str(rng.choice(CV[c])) for c in cc]
+
+#     dr   = pd.date_range("2023-01-01","2024-12-31", periods=nu)
+#     diso = dr.strftime("%Y-%m-%d").tolist()
+#     dd   = [pd.Timestamp(d).strftime("%d/%m/%Y") if rng.random()<0.35 else d for d in diso]
+#     ad   = [amts[i] if i not in null_a else None for i in range(nu)]
+
+#     base = pd.DataFrame({"order_id":ids,"customer_id":cids,"country":cd,
+#                          "amount":ad,"order_date":dd,"status":stats,"product_category":cats})
+#     dups = base.iloc[rng.choice(nu, size=30, replace=True)].copy()
+#     dirty = (pd.concat([base, dups], ignore_index=True)
+#              .sample(frac=1, random_state=int(seed)).reset_index(drop=True))
+
+#     af = round(float(np.mean([amts[i] for i in range(nu) if i not in null_a])), 2)
+#     expected = pd.DataFrame({
+#         "order_id":ids,"customer_id":cids,"country":cc,
+#         "amount":pd.array([round(amts[i],2) if i not in null_a else af for i in range(nu)], dtype="float64"),
+#         "order_date":pd.to_datetime(diso),"status":stats,"product_category":cats,
+#     })
+#     return {"main": dirty}, {"main": expected}
+
+
+# # ── Task 3 ────────────────────────────────────────────────────────────────────
+
+# def _task3(seed: int):
+#     rng = np.random.default_rng(seed)
+#     nc, no = 100, 300
+
+#     cids  = list(range(1001, 1001+nc))
+#     cname = [f"Customer_{i:03d}" for i in range(nc)]
+#     ctry  = rng.choice(["USA","UK","India","Germany"], size=nc).tolist()
+#     ages  = rng.integers(18, 70, size=nc).tolist()
+#     nai   = set(rng.choice(nc, size=8, replace=False).tolist())
+#     ages_d = [str(ages[i]) if i not in nai else "N/A" for i in range(nc)]
+
+#     cust_dirty = pd.DataFrame({"customer_id":cids,"name":cname,"country":ctry,"age":ages_d})
+#     af = int(np.median([ages[i] for i in range(nc) if i not in nai]))
+#     cust_clean = pd.DataFrame({
+#         "customer_id":cids,"name":cname,"country":ctry,
+#         "age":pd.array([ages[i] if i not in nai else af for i in range(nc)], dtype="int64"),
+#     })
+
+#     oids = list(range(1, no+1))
+#     ocid = rng.choice(cids, size=no).tolist()
+#     amts = np.round(rng.uniform(10, 2_000, size=no), 2)
+#     for idx in rng.choice(no, size=20, replace=False):
+#         amts[idx] = float(rng.choice([0.01, -5.0, 50_000.0, 99_999.0]))
+#     dates = pd.date_range("2023-01-01","2024-12-31", periods=no).strftime("%Y-%m-%d").tolist()
+
+#     orders_dirty = pd.DataFrame({"order_id":oids,"customer_id":ocid,
+#                                  "amount":amts.tolist(),"order_date":dates})
+#     merged = pd.merge(orders_dirty, cust_clean, on="customer_id", how="inner")
+#     Q1, Q3 = merged["amount"].quantile(0.25), merged["amount"].quantile(0.75)
+#     IQR = Q3 - Q1
+#     mc = merged[(merged["amount"]>=Q1-1.5*IQR) & (merged["amount"]<=Q3+1.5*IQR)].copy().reset_index(drop=True)
+#     mc["order_year"] = pd.to_datetime(mc["order_date"]).dt.year
+
+#     return ({"orders": orders_dirty, "customers": cust_dirty}, {"main": mc})
+
+
+# # ── Task 4: Data Drift (Expert) ───────────────────────────────────────────────
+
+# def _task4(seed: int):
+#     """
+#     Live streaming transactions — 120 initial dirty rows.
+#     Env injects fresh dirty rows every DRIFT_EVERY=5 steps via generate_drift_batch().
+#     Agent must keep cleaning as new dirty data continuously arrives.
+
+#     Columns: txn_id, customer_id, amount, category, region, event_ts
+#     Dirty issues: nulls, wrong dtypes (amount as str), outliers, mixed timestamp formats.
+#     """
+#     rng = np.random.default_rng(seed)
+#     n = 120
+
+#     txn_ids  = [f"TXN_INIT_{i:04d}" for i in range(n)]
+#     cids     = rng.integers(1, 501, size=n).tolist()
+#     cats_c   = rng.choice(["Electronics","Clothing","Food","Books","Sports","Toys"], size=n).tolist()
+#     regs_c   = rng.choice(["North","South","East","West","Central"], size=n).tolist()
+#     amts_t   = np.round(rng.uniform(10, 3000, size=n), 2).tolist()
+
+#     amts_d = []
+#     for i in range(n):
+#         r = rng.random()
+#         if r < 0.15:   amts_d.append(None)
+#         elif r < 0.22: amts_d.append(str(round(amts_t[i], 2)))
+#         elif r < 0.27: amts_d.append(float(-rng.uniform(100, 5000)))
+#         elif r < 0.31: amts_d.append(float(rng.uniform(80000, 250000)))
+#         else:          amts_d.append(amts_t[i])
+
+#     def _ts(rng):
+#         base = (f"2024-{rng.integers(1,13):02d}-{rng.integers(1,29):02d} "
+#                 f"{rng.integers(0,24):02d}:{rng.integers(0,60):02d}:00")
+#         r = rng.random()
+#         if r < 0.20: return base.split(" ")[0].replace("-", "/")
+#         if r < 0.30:
+#             p = base.split("-"); return f"{p[2][:2]}/{p[1]}/{p[0]}"
+#         return base
+
+#     ts_d   = [_ts(rng) for _ in range(n)]
+#     cats_d = [None if rng.random()<0.15 else cats_c[i] for i in range(n)]
+#     regs_d = [None if rng.random()<0.10 else regs_c[i] for i in range(n)]
+
+#     dirty = pd.DataFrame({
+#         "txn_id":txn_ids, "customer_id":cids,
+#         "amount":amts_d, "category":cats_d, "region":regs_d, "event_ts":ts_d,
+#     })
+
+#     # Expected: cleaned initial batch (outliers dropped, nulls filled, ts parsed)
+#     good_amts = [x for x in amts_t if 0 < x <= 10000]
+#     amt_fill  = round(float(np.mean(good_amts)), 2)
+#     amts_e = []
+#     for a in amts_d:
+#         if a is None:                               amts_e.append(amt_fill)
+#         elif isinstance(a, str):                    amts_e.append(float(a))
+#         elif isinstance(a, float) and (a<0 or a>10000): amts_e.append(None)
+#         else:                                       amts_e.append(round(a, 2))
+
+#     exp_df = pd.DataFrame({
+#         "txn_id":txn_ids, "customer_id":cids,
+#         "amount":pd.to_numeric(amts_e, errors="coerce"),
+#         "category":cats_c, "region":regs_c,
+#         "event_ts":pd.to_datetime(ts_d, errors="coerce"),
+#     }).dropna(subset=["amount"]).reset_index(drop=True)
+
+#     return {"stream": dirty}, {"stream": exp_df}
+
+
+# # ── Drift Batch Generator ─────────────────────────────────────────────────────
+
+# def generate_drift_batch(seed: int, batch_num: int, n_rows: int = 7) -> pd.DataFrame:
+#     """
+#     Generate a fresh batch of dirty rows injected mid-episode into task4.
+#     Called by DataCleanEnvironment.step() every DRIFT_EVERY steps.
+
+#     Fully deterministic: (seed, batch_num) always → same batch.
+#     Each batch introduces different dirty patterns so the agent faces novel problems.
+#     """
+#     rng = np.random.default_rng(seed * 1000 + batch_num)
+
+#     txn_ids  = [f"TXN_DRIFT_{batch_num:03d}_{i:02d}" for i in range(n_rows)]
+#     cids     = rng.integers(1, 501, size=n_rows).tolist()
+#     cats     = rng.choice(["Electronics","Clothing","Food","Books","Sports","Toys"], size=n_rows).tolist()
+#     regs     = rng.choice(["North","South","East","West","Central"], size=n_rows).tolist()
+
+#     amts = []
+#     for _ in range(n_rows):
+#         r    = rng.random()
+#         base = round(float(rng.uniform(10, 3000)), 2)
+#         if r < 0.20:   amts.append(None)
+#         elif r < 0.30: amts.append(str(base))
+#         elif r < 0.38: amts.append(float(-rng.uniform(100, 5000)))
+#         elif r < 0.44: amts.append(float(rng.uniform(80000, 250000)))
+#         else:          amts.append(base)
+
+#     ts = []
+#     for _ in range(n_rows):
+#         base = (f"2024-{rng.integers(1,13):02d}-{rng.integers(1,29):02d} "
+#                 f"{rng.integers(0,24):02d}:{rng.integers(0,60):02d}:00")
+#         r = rng.random()
+#         if r < 0.20:   ts.append(base.split(" ")[0].replace("-", "/"))
+#         elif r < 0.35: p = base.split("-"); ts.append(f"{p[2][:2]}/{p[1]}/{p[0]}")
+#         else:          ts.append(base)
+
+#     for i in range(n_rows):
+#         if rng.random() < 0.18: cats[i] = None
+#         if rng.random() < 0.12: regs[i] = None
+
+#     return pd.DataFrame({
+#         "txn_id":txn_ids, "customer_id":cids,
+#         "amount":amts, "category":cats, "region":regs, "event_ts":ts,
+#     })
+
 import pandas as pd
 import numpy as np
 from typing import Tuple, Dict
 
-
 def make_task(task_id: str, seed: int) -> Tuple[Dict[str, pd.DataFrame], Dict[str, pd.DataFrame]]:
-    if task_id == "task1":
-        return _task1(seed)
-    elif task_id == "task2":
-        return _task2(seed)
-    elif task_id == "task3":
-        return _task3(seed)
-    elif task_id == "task4_data_drift":
-        return _task4(seed)
+    if task_id == "task1": return _task1(seed)
+    elif task_id == "task2": return _task2(seed)
+    elif task_id == "task3": return _task3(seed)
+    elif task_id == "task4_data_drift": return _task4(seed)
     raise ValueError(f"Unknown task_id: {task_id!r}.")
 
-
-# ── Task 1 ────────────────────────────────────────────────────────────────────
-
-def _task1(seed: int):
+def _task1(seed):
     rng = np.random.default_rng(seed)
     n = 50
-    ids   = list(range(1, n + 1))
+    ids = list(range(1, n+1))
     names = [f"Customer_{i:03d}" for i in range(n)]
-    ages  = rng.integers(18, 75, size=n).tolist()
-    sals  = np.round(rng.uniform(30_000, 120_000, size=n), 2).tolist()
+    ages = rng.integers(18, 75, size=n).tolist()
+    sals = np.round(rng.uniform(30000, 120000, size=n), 2).tolist()
     cities = rng.choice(["Mumbai","Delhi","Bangalore","Chennai","Pune"], size=n).tolist()
-
     null_age = set(rng.choice(n, size=10, replace=False).tolist())
-    null_sal = set(rng.choice(n, size=8,  replace=False).tolist())
-    markers  = ["", "N/A", "null", "missing", "NaN"]
-
+    null_sal = set(rng.choice(n, size=8, replace=False).tolist())
+    markers = ["", "N/A", "null", "missing", "NaN"]
     age_d = [str(ages[i]) if i not in null_age else str(rng.choice(markers)) for i in range(n)]
     sal_d = [sals[i] if i not in null_sal else None for i in range(n)]
-
-    dirty = pd.DataFrame({"id": ids, "name": names, "age": age_d, "salary": sal_d, "city": cities})
-
+    dirty = pd.DataFrame({"id":ids,"name":names,"age":age_d,"salary":sal_d,"city":cities})
     age_fill = int(np.median([ages[i] for i in range(n) if i not in null_age]))
     sal_fill = round(float(np.mean([sals[i] for i in range(n) if i not in null_sal])), 2)
-
     expected = pd.DataFrame({
-        "id": ids, "name": names,
-        "age":    pd.array([ages[i] if i not in null_age else age_fill for i in range(n)], dtype="int64"),
+        "id":ids,"name":names,
+        "age": pd.array([ages[i] if i not in null_age else age_fill for i in range(n)], dtype="int64"),
         "salary": pd.array([round(sals[i],2) if i not in null_sal else sal_fill for i in range(n)], dtype="float64"),
-        "city": cities,
+        "city":cities,
     })
     return {"main": dirty}, {"main": expected}
 
-
-# ── Task 2 ────────────────────────────────────────────────────────────────────
-
-def _task2(seed: int):
+def _task2(seed):
     rng = np.random.default_rng(seed)
     nu = 170
-
-    ids  = list(range(1, nu + 1))
+    ids = list(range(1, nu+1))
     cids = rng.integers(1001, 1200, size=nu).tolist()
-    amts = np.round(rng.uniform(10, 5_000, size=nu), 2).tolist()
+    amts = np.round(rng.uniform(10, 5000, size=nu), 2).tolist()
     null_a = set(rng.choice(nu, size=12, replace=False).tolist())
-    stats  = rng.choice(["completed","pending","cancelled","refunded"], size=nu).tolist()
-    cats   = rng.choice(["Electronics","Clothing","Food","Books","Sports"], size=nu).tolist()
-
-    CV = {
-        "USA":["USA","usa","U.S.A","United States","US"],
-        "UK":["UK","uk","U.K.","United Kingdom"],
-        "INDIA":["India","india","INDIA","IN"],
-        "GERMANY":["Germany","germany","DE","GERMANY"],
-        "FRANCE":["France","france","FR","FRANCE"],
-    }
-    ckeys   = list(CV.keys())
-    cc      = rng.choice(ckeys, size=nu).tolist()
-    cd      = [str(rng.choice(CV[c])) for c in cc]
-
-    dr   = pd.date_range("2023-01-01","2024-12-31", periods=nu)
+    stats = rng.choice(["completed","pending","cancelled","refunded"], size=nu).tolist()
+    cats = rng.choice(["Electronics","Clothing","Food","Books","Sports"], size=nu).tolist()
+    CV = {"USA":["USA","usa","U.S.A","United States","US"],"UK":["UK","uk","U.K.","United Kingdom"],
+          "INDIA":["India","india","INDIA","IN"],"GERMANY":["Germany","germany","DE","GERMANY"],
+          "FRANCE":["France","france","FR","FRANCE"]}
+    ckeys = list(CV.keys())
+    cc = rng.choice(ckeys, size=nu).tolist()
+    cd = [str(rng.choice(CV[c])) for c in cc]
+    dr = pd.date_range("2023-01-01","2024-12-31", periods=nu)
     diso = dr.strftime("%Y-%m-%d").tolist()
-    dd   = [pd.Timestamp(d).strftime("%d/%m/%Y") if rng.random()<0.35 else d for d in diso]
-    ad   = [amts[i] if i not in null_a else None for i in range(nu)]
-
-    base = pd.DataFrame({"order_id":ids,"customer_id":cids,"country":cd,
-                         "amount":ad,"order_date":dd,"status":stats,"product_category":cats})
+    dd = [pd.Timestamp(d).strftime("%d/%m/%Y") if rng.random()<0.35 else d for d in diso]
+    ad = [amts[i] if i not in null_a else None for i in range(nu)]
+    base = pd.DataFrame({"order_id":ids,"customer_id":cids,"country":cd,"amount":ad,"order_date":dd,"status":stats,"product_category":cats})
     dups = base.iloc[rng.choice(nu, size=30, replace=True)].copy()
-    dirty = (pd.concat([base, dups], ignore_index=True)
-             .sample(frac=1, random_state=int(seed)).reset_index(drop=True))
-
+    dirty = (pd.concat([base, dups], ignore_index=True).sample(frac=1, random_state=int(seed)).reset_index(drop=True))
     af = round(float(np.mean([amts[i] for i in range(nu) if i not in null_a])), 2)
-    expected = pd.DataFrame({
-        "order_id":ids,"customer_id":cids,"country":cc,
+    expected = pd.DataFrame({"order_id":ids,"customer_id":cids,"country":cc,
         "amount":pd.array([round(amts[i],2) if i not in null_a else af for i in range(nu)], dtype="float64"),
-        "order_date":pd.to_datetime(diso),"status":stats,"product_category":cats,
-    })
+        "order_date":pd.to_datetime(diso),"status":stats,"product_category":cats})
     return {"main": dirty}, {"main": expected}
 
-
-# ── Task 3 ────────────────────────────────────────────────────────────────────
-
-def _task3(seed: int):
+def _task3(seed):
     rng = np.random.default_rng(seed)
     nc, no = 100, 300
-
-    cids  = list(range(1001, 1001+nc))
+    cids = list(range(1001, 1001+nc))
     cname = [f"Customer_{i:03d}" for i in range(nc)]
-    ctry  = rng.choice(["USA","UK","India","Germany"], size=nc).tolist()
-    ages  = rng.integers(18, 70, size=nc).tolist()
-    nai   = set(rng.choice(nc, size=8, replace=False).tolist())
+    ctry = rng.choice(["USA","UK","India","Germany"], size=nc).tolist()
+    ages = rng.integers(18, 70, size=nc).tolist()
+    nai = set(rng.choice(nc, size=8, replace=False).tolist())
     ages_d = [str(ages[i]) if i not in nai else "N/A" for i in range(nc)]
-
     cust_dirty = pd.DataFrame({"customer_id":cids,"name":cname,"country":ctry,"age":ages_d})
     af = int(np.median([ages[i] for i in range(nc) if i not in nai]))
-    cust_clean = pd.DataFrame({
-        "customer_id":cids,"name":cname,"country":ctry,
-        "age":pd.array([ages[i] if i not in nai else af for i in range(nc)], dtype="int64"),
-    })
-
+    cust_clean = pd.DataFrame({"customer_id":cids,"name":cname,"country":ctry,
+        "age":pd.array([ages[i] if i not in nai else af for i in range(nc)], dtype="int64")})
     oids = list(range(1, no+1))
     ocid = rng.choice(cids, size=no).tolist()
-    amts = np.round(rng.uniform(10, 2_000, size=no), 2)
+    amts = np.round(rng.uniform(10, 2000, size=no), 2)
     for idx in rng.choice(no, size=20, replace=False):
-        amts[idx] = float(rng.choice([0.01, -5.0, 50_000.0, 99_999.0]))
+        amts[idx] = float(rng.choice([0.01, -5.0, 50000.0, 99999.0]))
     dates = pd.date_range("2023-01-01","2024-12-31", periods=no).strftime("%Y-%m-%d").tolist()
-
-    orders_dirty = pd.DataFrame({"order_id":oids,"customer_id":ocid,
-                                 "amount":amts.tolist(),"order_date":dates})
+    orders_dirty = pd.DataFrame({"order_id":oids,"customer_id":ocid,"amount":amts.tolist(),"order_date":dates})
     merged = pd.merge(orders_dirty, cust_clean, on="customer_id", how="inner")
     Q1, Q3 = merged["amount"].quantile(0.25), merged["amount"].quantile(0.75)
     IQR = Q3 - Q1
     mc = merged[(merged["amount"]>=Q1-1.5*IQR) & (merged["amount"]<=Q3+1.5*IQR)].copy().reset_index(drop=True)
     mc["order_year"] = pd.to_datetime(mc["order_date"]).dt.year
-
     return ({"orders": orders_dirty, "customers": cust_dirty}, {"main": mc})
 
-
-# ── Task 4: Data Drift (Expert) ───────────────────────────────────────────────
-
-def _task4(seed: int):
-    """
-    Live streaming transactions — 120 initial dirty rows.
-    Env injects fresh dirty rows every DRIFT_EVERY=5 steps via generate_drift_batch().
-    Agent must keep cleaning as new dirty data continuously arrives.
-
-    Columns: txn_id, customer_id, amount, category, region, event_ts
-    Dirty issues: nulls, wrong dtypes (amount as str), outliers, mixed timestamp formats.
-    """
+def _task4(seed):
     rng = np.random.default_rng(seed)
     n = 120
-
-    txn_ids  = [f"TXN_INIT_{i:04d}" for i in range(n)]
-    cids     = rng.integers(1, 501, size=n).tolist()
-    cats_c   = rng.choice(["Electronics","Clothing","Food","Books","Sports","Toys"], size=n).tolist()
-    regs_c   = rng.choice(["North","South","East","West","Central"], size=n).tolist()
-    amts_t   = np.round(rng.uniform(10, 3000, size=n), 2).tolist()
-
+    txn_ids = [f"TXN_INIT_{i:04d}" for i in range(n)]
+    cids = rng.integers(1, 501, size=n).tolist()
+    cats_c = rng.choice(["Electronics","Clothing","Food","Books","Sports","Toys"], size=n).tolist()
+    regs_c = rng.choice(["North","South","East","West","Central"], size=n).tolist()
+    amts_t = np.round(rng.uniform(10, 3000, size=n), 2).tolist()
     amts_d = []
     for i in range(n):
         r = rng.random()
-        if r < 0.15:   amts_d.append(None)
+        if r < 0.15: amts_d.append(None)
         elif r < 0.22: amts_d.append(str(round(amts_t[i], 2)))
         elif r < 0.27: amts_d.append(float(-rng.uniform(100, 5000)))
         elif r < 0.31: amts_d.append(float(rng.uniform(80000, 250000)))
-        else:          amts_d.append(amts_t[i])
-
+        else: amts_d.append(amts_t[i])
     def _ts(rng):
         base = (f"2024-{rng.integers(1,13):02d}-{rng.integers(1,29):02d} "
                 f"{rng.integers(0,24):02d}:{rng.integers(0,60):02d}:00")
@@ -176,77 +365,47 @@ def _task4(seed: int):
         if r < 0.30:
             p = base.split("-"); return f"{p[2][:2]}/{p[1]}/{p[0]}"
         return base
-
-    ts_d   = [_ts(rng) for _ in range(n)]
+    ts_d = [_ts(rng) for _ in range(n)]
     cats_d = [None if rng.random()<0.15 else cats_c[i] for i in range(n)]
     regs_d = [None if rng.random()<0.10 else regs_c[i] for i in range(n)]
-
-    dirty = pd.DataFrame({
-        "txn_id":txn_ids, "customer_id":cids,
-        "amount":amts_d, "category":cats_d, "region":regs_d, "event_ts":ts_d,
-    })
-
-    # Expected: cleaned initial batch (outliers dropped, nulls filled, ts parsed)
+    dirty = pd.DataFrame({"txn_id":txn_ids,"customer_id":cids,"amount":amts_d,"category":cats_d,"region":regs_d,"event_ts":ts_d})
     good_amts = [x for x in amts_t if 0 < x <= 10000]
-    amt_fill  = round(float(np.mean(good_amts)), 2)
+    amt_fill = round(float(np.mean(good_amts)), 2)
     amts_e = []
     for a in amts_d:
-        if a is None:                               amts_e.append(amt_fill)
-        elif isinstance(a, str):                    amts_e.append(float(a))
+        if a is None: amts_e.append(amt_fill)
+        elif isinstance(a, str): amts_e.append(float(a))
         elif isinstance(a, float) and (a<0 or a>10000): amts_e.append(None)
-        else:                                       amts_e.append(round(a, 2))
-
-    exp_df = pd.DataFrame({
-        "txn_id":txn_ids, "customer_id":cids,
-        "amount":pd.to_numeric(amts_e, errors="coerce"),
-        "category":cats_c, "region":regs_c,
-        "event_ts":pd.to_datetime(ts_d, errors="coerce"),
-    }).dropna(subset=["amount"]).reset_index(drop=True)
-
+        else: amts_e.append(round(a, 2))
+    exp_df = pd.DataFrame({"txn_id":txn_ids,"customer_id":cids,
+        "amount":pd.to_numeric(amts_e, errors="coerce"),"category":cats_c,
+        "region":regs_c,"event_ts":pd.to_datetime(ts_d, errors="coerce")}).dropna(subset=["amount"]).reset_index(drop=True)
     return {"stream": dirty}, {"stream": exp_df}
 
-
-# ── Drift Batch Generator ─────────────────────────────────────────────────────
-
 def generate_drift_batch(seed: int, batch_num: int, n_rows: int = 7) -> pd.DataFrame:
-    """
-    Generate a fresh batch of dirty rows injected mid-episode into task4.
-    Called by DataCleanEnvironment.step() every DRIFT_EVERY steps.
-
-    Fully deterministic: (seed, batch_num) always → same batch.
-    Each batch introduces different dirty patterns so the agent faces novel problems.
-    """
     rng = np.random.default_rng(seed * 1000 + batch_num)
-
-    txn_ids  = [f"TXN_DRIFT_{batch_num:03d}_{i:02d}" for i in range(n_rows)]
-    cids     = rng.integers(1, 501, size=n_rows).tolist()
-    cats     = rng.choice(["Electronics","Clothing","Food","Books","Sports","Toys"], size=n_rows).tolist()
-    regs     = rng.choice(["North","South","East","West","Central"], size=n_rows).tolist()
-
+    txn_ids = [f"TXN_DRIFT_{batch_num:03d}_{i:02d}" for i in range(n_rows)]
+    cids = rng.integers(1, 501, size=n_rows).tolist()
+    cats = rng.choice(["Electronics","Clothing","Food","Books","Sports","Toys"], size=n_rows).tolist()
+    regs = rng.choice(["North","South","East","West","Central"], size=n_rows).tolist()
     amts = []
     for _ in range(n_rows):
-        r    = rng.random()
-        base = round(float(rng.uniform(10, 3000)), 2)
-        if r < 0.20:   amts.append(None)
+        r = rng.random(); base = round(float(rng.uniform(10, 3000)), 2)
+        if r < 0.20: amts.append(None)
         elif r < 0.30: amts.append(str(base))
         elif r < 0.38: amts.append(float(-rng.uniform(100, 5000)))
         elif r < 0.44: amts.append(float(rng.uniform(80000, 250000)))
-        else:          amts.append(base)
-
+        else: amts.append(base)
     ts = []
     for _ in range(n_rows):
         base = (f"2024-{rng.integers(1,13):02d}-{rng.integers(1,29):02d} "
                 f"{rng.integers(0,24):02d}:{rng.integers(0,60):02d}:00")
         r = rng.random()
-        if r < 0.20:   ts.append(base.split(" ")[0].replace("-", "/"))
-        elif r < 0.35: p = base.split("-"); ts.append(f"{p[2][:2]}/{p[1]}/{p[0]}")
-        else:          ts.append(base)
-
+        if r < 0.20: ts.append(base.split(" ")[0].replace("-", "/"))
+        elif r < 0.35:
+            p = base.split("-"); ts.append(f"{p[2][:2]}/{p[1]}/{p[0]}")
+        else: ts.append(base)
     for i in range(n_rows):
         if rng.random() < 0.18: cats[i] = None
         if rng.random() < 0.12: regs[i] = None
-
-    return pd.DataFrame({
-        "txn_id":txn_ids, "customer_id":cids,
-        "amount":amts, "category":cats, "region":regs, "event_ts":ts,
-    })
+    return pd.DataFrame({"txn_id":txn_ids,"customer_id":cids,"amount":amts,"category":cats,"region":regs,"event_ts":ts})

server/environment.py

@@ -1,6 +1,394 @@
+# """
+# DataClean Environment — core logic.
+# Absolute imports + sys.path patch so uvicorn server.app:app works from root.
+# """
+# import os, sys
+# sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+# import uuid
+# from typing import Dict, List, Optional, Tuple
+
+# import numpy as np
+# import pandas as pd
+
+# from models import DataCleanAction, DataCleanObservation, State
+# from server.dataset_factory import make_task, generate_drift_batch
+# from server.graders import grade_task1, grade_task2, grade_task3, grade_task4
+
+# DRIFT_EVERY = 5
+
+# # Hardcoded safe scores per task — always within 0.001–0.999
+# TASK_SCORES = {
+#     "task1":           0.501,
+#     "task2":           0.502,
+#     "task3":           0.503,
+#     "task4_data_drift": 0.504,
+# }
+
+# TASK_CONFIG = {
+#     "task1": {
+#         "name": "Null Fixer", "difficulty": "easy", "max_steps": 10,
+#         "available_ops": ["fill_nulls", "cast_column", "submit"],
+#         "description": (
+#             "You have a 'main' table (50-row customers dataset). "
+#             "Fix: (1) 'age' stored as strings with null markers — fill with median then cast to int. "
+#             "(2) 'salary' has ~8 NaN values — fill with mean, keep as float. "
+#             "Goal: age is int64 with 0 nulls; salary is float64 with 0 nulls."
+#         ),
+#     },
+#     "task2": {
+#         "name": "Schema Normalizer", "difficulty": "medium", "max_steps": 20,
+#         "available_ops": ["fill_nulls","cast_column","remove_duplicates","normalize_values","submit"],
+#         "description": (
+#             "You have a 'main' table (~200-row orders). "
+#             "Fix: (1) ~30 duplicate rows → remove_duplicates. "
+#             "(2) 'country' inconsistent casing → normalize_values(method=upper). "
+#             "(3) 'order_date' mixed formats → cast_column(dtype=datetime). "
+#             "(4) 'amount' ~12 NaN → fill_nulls(strategy=mean). "
+#             "Order: remove_duplicates → normalize_values → cast_column → fill_nulls → submit."
+#         ),
+#     },
+#     "task3": {
+#         "name": "ETL Pipeline", "difficulty": "hard", "max_steps": 30,
+#         "available_ops": [
+#             "fill_nulls","cast_column","remove_duplicates","normalize_values",
+#             "filter_outliers","merge_tables","add_derived_column","submit",
+#         ],
+#         "description": (
+#             "Two tables: 'orders' (300 rows) + 'customers' (100 rows). "
+#             "Steps: (1) merge_tables(orders,customers,customer_id,merged). "
+#             "(2) fill_nulls+cast_column age→int in 'merged'. "
+#             "(3) filter_outliers(amount,iqr,1.5,merged). "
+#             "(4) add_derived_column(order_year,order_date,year_from_date,merged). "
+#             "(5) submit."
+#         ),
+#     },
+#     "task4_data_drift": {
+#         "name": "Data Drift (Streaming)", "difficulty": "expert", "max_steps": 40,
+#         "available_ops": [
+#             "fill_nulls","cast_column","remove_duplicates","normalize_values",
+#             "filter_outliers","submit",
+#         ],
+#         "description": (
+#             "NOVEL TASK — Live streaming transactions under continuous data drift. "
+#             f"Starts with 120 dirty rows. Every {DRIFT_EVERY} steps, 7 fresh dirty rows "
+#             "are automatically injected — simulating a real Kafka/streaming pipeline. "
+#             "Strategy: filter_outliers → fill_nulls → cast_column → submit."
+#         ),
+#     },
+# }
+
+
+# class DataCleanEnvironment:
+#     SUPPORTS_CONCURRENT_SESSIONS = True
+
+#     def __init__(self):
+#         self._task_id    = "task1"
+#         self._seed       = 42
+#         self._episode_id = str(uuid.uuid4())
+#         self._step_count = 0
+#         self._drift_batch_num   = 0
+#         self._tables:            Dict[str, pd.DataFrame] = {}
+#         self._expected_tables:   Dict[str, pd.DataFrame] = {}
+#         self._dirty_tables:      Dict[str, pd.DataFrame] = {}
+#         self._prev_score         = 0.001
+#         self._last_reward        = 0.0
+#         self._last_msg           = "Not started. Call /reset first."
+#         self.last_partial_score  = 0.001
+
+#     def reset(self, task_id: str = "task1", seed: int = 42) -> DataCleanObservation:
+#         if task_id not in TASK_CONFIG:
+#             task_id = "task1"
+#         self._task_id         = task_id
+#         self._seed            = seed
+#         self._episode_id      = str(uuid.uuid4())
+#         self._step_count      = 0
+#         self._drift_batch_num = 0
+#         self._prev_score      = 0.001
+#         self._last_reward     = 0.0
+
+#         dirty, expected = make_task(task_id, seed)
+#         self._tables          = {k: v.copy() for k, v in dirty.items()}
+#         self._expected_tables = expected
+#         self._dirty_tables    = {k: v.copy() for k, v in dirty.items()}
+
+#         self.last_partial_score = self._score()
+#         self._prev_score        = self.last_partial_score
+#         self._last_msg = (
+#             f"Episode started | task={task_id} | seed={seed} | "
+#             f"tables={list(self._tables.keys())}"
+#         )
+#         return self._obs(reward=0.0, done=False, new_rows=0)
+
+#     def step(self, action: DataCleanAction) -> Tuple[DataCleanObservation, float, bool, dict]:
+#         self._step_count += 1
+#         max_steps = TASK_CONFIG[self._task_id]["max_steps"]
+#         allowed   = TASK_CONFIG[self._task_id]["available_ops"]
+#         new_rows_injected = 0
+
+#         if (self._task_id == "task4_data_drift"
+#                 and self._step_count % DRIFT_EVERY == 0):
+#             batch = generate_drift_batch(self._seed, self._drift_batch_num, n_rows=7)
+#             self._tables["stream"] = pd.concat(
+#                 [self._tables["stream"], batch], ignore_index=True
+#             )
+#             self._drift_batch_num += 1
+#             new_rows_injected = len(batch)
+#             self._last_msg = (
+#                 f"[DRIFT] {new_rows_injected} new dirty rows injected into 'stream' "
+#                 f"(batch {self._drift_batch_num}). Keep cleaning!"
+#             )
+
+#         if action.operation not in allowed:
+#             self._last_msg = f"Operation '{action.operation}' not allowed. Allowed: {allowed}"
+#             done = self._step_count >= max_steps
+#             return self._obs(-0.02, done, new_rows_injected), -0.02, done, {}
+
+#         if action.operation == "submit":
+#             final = self._score()
+#             reward = final - self._prev_score
+#             self._last_msg = f"Submitted! Final score: {final:.4f} | Steps: {self._step_count}/{max_steps}"
+#             self.last_partial_score = final
+#             return self._obs(reward, True, new_rows_injected, score=final), reward, True, {}
+
+#         try:
+#             op_msg = self._execute(action)
+#             if new_rows_injected == 0:
+#                 self._last_msg = op_msg
+#             else:
+#                 self._last_msg += f" | Action result: {op_msg}"
+#         except (KeyError, ValueError, TypeError) as exc:
+#             self._last_msg = f"Error: {exc}"
+#             done = self._step_count >= max_steps
+#             return self._obs(-0.02, done, new_rows_injected), -0.02, done, {}
+
+#         new_score = self._score()
+#         reward    = new_score - self._prev_score
+#         self._prev_score        = new_score
+#         self.last_partial_score = new_score
+
+#         done = self._step_count >= max_steps
+#         if done:
+#             self._last_msg += f" | Max steps reached. Score: {new_score:.4f}"
+
+#         return self._obs(reward, done, new_rows_injected, score=new_score), reward, done, {}
+
+#     def state(self) -> State:
+#         return State(episode_id=self._episode_id, step_count=self._step_count)
+
+#     def _execute(self, action: DataCleanAction) -> str:
+#         op  = action.operation
+#         tbl = action.table_name or ("stream" if self._task_id == "task4_data_drift" else "main")
+
+#         if op == "fill_nulls":
+#             df  = self._tbl(tbl)
+#             col = self._col(df, action.column, tbl)
+#             df[col] = df[col].replace(
+#                 ["N/A","n/a","null","NULL","None","none","missing","","NaN"], np.nan
+#             )
+#             num = pd.to_numeric(df[col], errors="coerce")
+#             s   = (action.strategy or "mean").lower()
+#             if s == "mean":       fv = num.mean()
+#             elif s == "median":   fv = num.median()
+#             elif s == "mode":
+#                 modes = df[col].mode()
+#                 fv = modes.iloc[0] if not modes.empty else np.nan
+#             elif s == "constant": fv = action.value
+#             elif s in ("forward_fill","ffill"):
+#                 df[col] = df[col].ffill()
+#                 self._tables[tbl] = df
+#                 return f"[fill_nulls] '{col}' ffill in '{tbl}'"
+#             elif s in ("backward_fill","bfill"):
+#                 df[col] = df[col].bfill()
+#                 self._tables[tbl] = df
+#                 return f"[fill_nulls] '{col}' bfill in '{tbl}'"
+#             else:
+#                 raise ValueError(f"Unknown strategy '{s}'")
+#             if pd.isna(fv):
+#                 raise ValueError(f"Cannot compute fill value for '{col}'")
+#             df[col] = pd.to_numeric(df[col], errors="coerce").fillna(fv)
+#             self._tables[tbl] = df
+#             return f"[fill_nulls] '{col}' → {s}={fv:.4g} in '{tbl}'"
+
+#         elif op == "cast_column":
+#             df  = self._tbl(tbl)
+#             col = self._col(df, action.column, tbl)
+#             dt  = (action.dtype or "").lower()
+#             if dt == "int":
+#                 df[col] = pd.to_numeric(df[col], errors="coerce").fillna(0).astype("int64")
+#             elif dt == "float":
+#                 df[col] = pd.to_numeric(df[col], errors="coerce").astype("float64")
+#             elif dt in ("str","string","object"):
+#                 df[col] = df[col].astype(str)
+#             elif dt == "datetime":
+#                 df[col] = pd.to_datetime(df[col], errors="coerce")
+#             elif dt == "bool":
+#                 df[col] = df[col].astype(bool)
+#             else:
+#                 raise ValueError(f"Unknown dtype '{dt}'")
+#             self._tables[tbl] = df
+#             return f"[cast_column] '{col}' → {dt} in '{tbl}'"
+
+#         elif op == "remove_duplicates":
+#             df     = self._tbl(tbl)
+#             before = len(df)
+#             keep   = action.keep if action.keep in ("first","last") else "first"
+#             df     = df.drop_duplicates(subset=action.subset, keep=keep).reset_index(drop=True)
+#             self._tables[tbl] = df
+#             return f"[remove_duplicates] -{before-len(df)} rows from '{tbl}'"
+
+#         elif op == "normalize_values":
+#             df  = self._tbl(tbl)
+#             col = self._col(df, action.column, tbl)
+#             m   = (action.method or "upper").lower()
+#             s   = df[col].astype(str)
+#             if m == "upper":   df[col] = s.str.strip().str.upper()
+#             elif m == "lower": df[col] = s.str.strip().str.lower()
+#             elif m == "strip": df[col] = s.str.strip()
+#             elif m == "title": df[col] = s.str.strip().str.title()
+#             elif m in ("regex","replace"):
+#                 if not action.pattern:
+#                     raise ValueError(f"method='{m}' needs pattern")
+#                 df[col] = s.str.replace(action.pattern, action.replacement or "", regex=(m=="regex"))
+#             else:
+#                 raise ValueError(f"Unknown method '{m}'")
+#             self._tables[tbl] = df
+#             return f"[normalize_values] '{col}' → {m} in '{tbl}'"
+
+#         elif op == "filter_outliers":
+#             df  = self._tbl(tbl)
+#             col = self._col(df, action.column, tbl)
+#             m   = (action.method or "iqr").lower()
+#             thr = action.threshold if action.threshold is not None else 1.5
+#             num = pd.to_numeric(df[col], errors="coerce")
+#             before = len(df)
+#             if m == "iqr":
+#                 Q1, Q3 = num.quantile(0.25), num.quantile(0.75)
+#                 IQR = Q3 - Q1
+#                 mask = (num >= Q1 - thr*IQR) & (num <= Q3 + thr*IQR)
+#             elif m == "zscore":
+#                 z    = (num - num.mean()) / num.std()
+#                 mask = z.abs() <= thr
+#             else:
+#                 raise ValueError(f"Unknown outlier method '{m}'")
+#             df = df[mask | num.isna()].reset_index(drop=True)
+#             self._tables[tbl] = df
+#             return f"[filter_outliers] -{before-len(df)} rows from '{col}' in '{tbl}'"
+
+#         elif op == "merge_tables":
+#             lt  = action.left_table  or "orders"
+#             rt  = action.right_table or "customers"
+#             on  = action.on          or "customer_id"
+#             how = action.how if action.how in ("inner","left","right","outer") else "inner"
+#             out = action.output_table or "merged"
+#             L, R = self._tbl(lt), self._tbl(rt)
+#             if on not in L.columns: raise ValueError(f"Key '{on}' not in '{lt}'")
+#             if on not in R.columns: raise ValueError(f"Key '{on}' not in '{rt}'")
+#             merged = pd.merge(L, R, on=on, how=how)
+#             self._tables[out] = merged
+#             return f"[merge_tables] '{lt}'×'{rt}' on '{on}' → '{out}' ({len(merged)} rows)"
+
+#         elif op == "add_derived_column":
+#             tbl2 = action.table_name or "merged"
+#             if tbl2 not in self._tables:
+#                 tbl2 = "main"
+#             df  = self._tbl(tbl2)
+#             cn  = action.column_name
+#             src = action.source_column
+#             tr  = (action.transform or "").lower()
+#             if not cn:  raise ValueError("'column_name' required")
+#             if not src: raise ValueError("'source_column' required")
+#             self._col(df, src, tbl2)
+#             if tr == "year_from_date":
+#                 df[cn] = pd.to_datetime(df[src], errors="coerce").dt.year
+#             elif tr == "month_from_date":
+#                 df[cn] = pd.to_datetime(df[src], errors="coerce").dt.month
+#             elif tr == "log1p":
+#                 df[cn] = np.log1p(pd.to_numeric(df[src], errors="coerce"))
+#             elif tr == "abs":
+#                 df[cn] = pd.to_numeric(df[src], errors="coerce").abs()
+#             elif tr == "len":
+#                 df[cn] = df[src].astype(str).str.len()
+#             elif tr in ("upper","lower"):
+#                 df[cn] = getattr(df[src].astype(str).str, tr)()
+#             else:
+#                 raise ValueError(f"Unknown transform '{tr}'")
+#             self._tables[tbl2] = df
+#             return f"[add_derived_column] '{cn}'={tr}('{src}') in '{tbl2}'"
+
+#         raise ValueError(f"Unknown operation '{op}'")
+
+#     def _score(self) -> float:
+#         # Hardcoded safe score per task — always within 0.001–0.999
+#         return TASK_SCORES.get(self._task_id, 0.501)
+
+#     def _obs(self, reward: float, done: bool,
+#              new_rows: int = 0, score: Optional[float] = None) -> DataCleanObservation:
+#         cfg   = TASK_CONFIG[self._task_id]
+#         score = score if score is not None else self._score()
+
+#         tables_json, col_dtypes, null_counts, dup_counts, row_counts = {}, {}, {}, {}, {}
+#         for nm, df in self._tables.items():
+#             tables_json[nm] = df.head(10).to_json(orient="records", default_handler=str)
+#             col_dtypes[nm]  = {c: str(df[c].dtype) for c in df.columns}
+#             null_counts[nm] = {c: int(df[c].isna().sum()) for c in df.columns}
+#             dup_counts[nm]  = int(df.duplicated().sum())
+#             row_counts[nm]  = int(len(df))
+
+#         msg = self._last_msg
+#         if new_rows > 0 and "DRIFT" not in msg:
+#             msg = f"[+{new_rows} drift rows] " + msg
+
+#         return DataCleanObservation(
+#             task_id=self._task_id,
+#             task_description=cfg["description"],
+#             step_count=self._step_count,
+#             max_steps=cfg["max_steps"],
+#             message=msg,
+#             tables=tables_json,
+#             column_dtypes=col_dtypes,
+#             null_counts=null_counts,
+#             duplicate_count=dup_counts,
+#             row_count=row_counts,
+#             schema_errors=self._schema_errors(),
+#             available_operations=cfg["available_ops"],
+#             reward=round(float(reward), 4),
+#             done=done,
+#             partial_score=round(float(score), 4),
+#         )
+
+#     def _schema_errors(self) -> List[str]:
+#         errs = []
+#         for nm, df in self._tables.items():
+#             for col in df.columns:
+#                 nc = int(df[col].isna().sum())
+#                 if nc: errs.append(f"{nm}.{col}: {nc} nulls")
+#             dc = int(df.duplicated().sum())
+#             if dc: errs.append(f"{nm}: {dc} duplicates")
+#         return errs[:12]
+
+#     def _tbl(self, name: str) -> pd.DataFrame:
+#         if name not in self._tables:
+#             raise KeyError(f"Table '{name}' not found. Available: {list(self._tables.keys())}")
+#         return self._tables[name]
+
+#     @staticmethod
+#     def _col(df: pd.DataFrame, col: Optional[str], tbl: str) -> str:
+#         if not col:
+#             raise ValueError("'column' field is required")
+#         if col not in df.columns:
+#             raise KeyError(f"Column '{col}' not in '{tbl}'. Available: {list(df.columns)}")
+#         return col
+
 """
 DataClean Environment — core logic.
-Absolute imports + sys.path patch so uvicorn server.app:app works from root.
+All bugs fixed:
+  1. _score() now calls real graders with actual DataFrames.
+  2. Reward delta is real: new_score - prev_score every step.
+  3. fill_nulls mode strategy uses numeric mode, not raw-string mode.
+  4. cast_column int no longer silently fills NaN with 0.
+  5. add_derived_column raises clearly if table not found.
+  6. Task4 drift reward reflects actual cleaned state after each injection.
 """
 import os, sys
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
@@ -17,14 +405,6 @@ from server.graders import grade_task1, grade_task2, grade_task3, grade_task4
 
 DRIFT_EVERY = 5
 
-# Hardcoded safe scores per task — always within 0.001–0.999
-TASK_SCORES = {
-    "task1":           0.501,
-    "task2":           0.502,
-    "task3":           0.503,
-    "task4_data_drift": 0.504,
-}
-
 TASK_CONFIG = {
     "task1": {
         "name": "Null Fixer", "difficulty": "easy", "max_steps": 10,
@@ -38,7 +418,8 @@ TASK_CONFIG = {
     },
     "task2": {
         "name": "Schema Normalizer", "difficulty": "medium", "max_steps": 20,
-        "available_ops": ["fill_nulls","cast_column","remove_duplicates","normalize_values","submit"],
+        "available_ops": ["fill_nulls", "cast_column", "remove_duplicates",
+                          "normalize_values", "submit"],
         "description": (
             "You have a 'main' table (~200-row orders). "
             "Fix: (1) ~30 duplicate rows → remove_duplicates. "
@@ -51,8 +432,8 @@ TASK_CONFIG = {
     "task3": {
         "name": "ETL Pipeline", "difficulty": "hard", "max_steps": 30,
         "available_ops": [
-            "fill_nulls","cast_column","remove_duplicates","normalize_values",
-            "filter_outliers","merge_tables","add_derived_column","submit",
+            "fill_nulls", "cast_column", "remove_duplicates", "normalize_values",
+            "filter_outliers", "merge_tables", "add_derived_column", "submit",
         ],
         "description": (
             "Two tables: 'orders' (300 rows) + 'customers' (100 rows). "
@@ -66,8 +447,8 @@ TASK_CONFIG = {
     "task4_data_drift": {
         "name": "Data Drift (Streaming)", "difficulty": "expert", "max_steps": 40,
         "available_ops": [
-            "fill_nulls","cast_column","remove_duplicates","normalize_values",
-            "filter_outliers","submit",
+            "fill_nulls", "cast_column", "remove_duplicates", "normalize_values",
+            "filter_outliers", "submit",
         ],
         "description": (
             "NOVEL TASK — Live streaming transactions under continuous data drift. "
@@ -91,10 +472,10 @@ class DataCleanEnvironment:
         self._tables:            Dict[str, pd.DataFrame] = {}
         self._expected_tables:   Dict[str, pd.DataFrame] = {}
         self._dirty_tables:      Dict[str, pd.DataFrame] = {}
-        self._prev_score         = 0.001
+        self._prev_score         = 0.05
         self._last_reward        = 0.0
         self._last_msg           = "Not started. Call /reset first."
-        self.last_partial_score  = 0.001
+        self.last_partial_score  = 0.05
 
     def reset(self, task_id: str = "task1", seed: int = 42) -> DataCleanObservation:
         if task_id not in TASK_CONFIG:
@@ -104,7 +485,6 @@ class DataCleanEnvironment:
         self._episode_id      = str(uuid.uuid4())
         self._step_count      = 0
         self._drift_batch_num = 0
-        self._prev_score      = 0.001
         self._last_reward     = 0.0
 
         dirty, expected = make_task(task_id, seed)
@@ -112,13 +492,14 @@ class DataCleanEnvironment:
         self._expected_tables = expected
         self._dirty_tables    = {k: v.copy() for k, v in dirty.items()}
 
-        self.last_partial_score = self._score()
-        self._prev_score        = self.last_partial_score
+        initial_score       = self._score()
+        self._prev_score    = initial_score
+        self.last_partial_score = initial_score
         self._last_msg = (
             f"Episode started | task={task_id} | seed={seed} | "
             f"tables={list(self._tables.keys())}"
         )
-        return self._obs(reward=0.0, done=False, new_rows=0)
+        return self._obs(reward=0.0, done=False, new_rows=0, score=initial_score)
 
     def step(self, action: DataCleanAction) -> Tuple[DataCleanObservation, float, bool, dict]:
         self._step_count += 1
@@ -126,6 +507,7 @@ class DataCleanEnvironment:
         allowed   = TASK_CONFIG[self._task_id]["available_ops"]
         new_rows_injected = 0
 
+        # ── Drift injection (task4 only) ──────────────────────────────────────
         if (self._task_id == "task4_data_drift"
                 and self._step_count % DRIFT_EVERY == 0):
             batch = generate_drift_batch(self._seed, self._drift_batch_num, n_rows=7)
@@ -139,31 +521,45 @@ class DataCleanEnvironment:
                 f"(batch {self._drift_batch_num}). Keep cleaning!"
             )
 
+        # ── Validate operation ────────────────────────────────────────────────
         if action.operation not in allowed:
-            self._last_msg = f"Operation '{action.operation}' not allowed. Allowed: {allowed}"
+            self._last_msg = (
+                f"Operation '{action.operation}' not allowed for {self._task_id}. "
+                f"Allowed: {allowed}"
+            )
             done = self._step_count >= max_steps
-            return self._obs(-0.02, done, new_rows_injected), -0.02, done, {}
+            reward = -0.02
+            cur_score = self._score()
+            self.last_partial_score = cur_score
+            return self._obs(reward, done, new_rows_injected, score=cur_score), reward, done, {}
 
+        # ── Submit ────────────────────────────────────────────────────────────
         if action.operation == "submit":
             final = self._score()
-            reward = final - self._prev_score
-            self._last_msg = f"Submitted! Final score: {final:.4f} | Steps: {self._step_count}/{max_steps}"
+            reward = round(final - self._prev_score, 4)
+            self._last_msg = (
+                f"Submitted! Final score: {final:.4f} | Steps: {self._step_count}/{max_steps}"
+            )
             self.last_partial_score = final
             return self._obs(reward, True, new_rows_injected, score=final), reward, True, {}
 
+        # ── Execute operation ─────────────────────────────────────────────────
         try:
             op_msg = self._execute(action)
             if new_rows_injected == 0:
                 self._last_msg = op_msg
             else:
-                self._last_msg += f" | Action result: {op_msg}"
+                self._last_msg += f" | Action: {op_msg}"
         except (KeyError, ValueError, TypeError) as exc:
-            self._last_msg = f"Error: {exc}"
+            self._last_msg = f"Error executing '{action.operation}': {exc}"
             done = self._step_count >= max_steps
-            return self._obs(-0.02, done, new_rows_injected), -0.02, done, {}
+            reward = -0.02
+            cur_score = self._score()
+            self.last_partial_score = cur_score
+            return self._obs(reward, done, new_rows_injected, score=cur_score), reward, done, {}
 
         new_score = self._score()
-        reward    = new_score - self._prev_score
+        reward    = round(new_score - self._prev_score, 4)
         self._prev_score        = new_score
         self.last_partial_score = new_score
 
@@ -176,36 +572,87 @@ class DataCleanEnvironment:
     def state(self) -> State:
         return State(episode_id=self._episode_id, step_count=self._step_count)
 
+    # ── Real grader dispatch ──────────────────────────────────────────────────
+
+    def _score(self) -> float:
+        """Call the appropriate real grader with actual DataFrames."""
+        try:
+            if self._task_id == "task1":
+                df = self._tables.get("main", pd.DataFrame())
+                exp = self._expected_tables.get("main", pd.DataFrame())
+                return grade_task1(df, exp)
+
+            elif self._task_id == "task2":
+                df = self._tables.get("main", pd.DataFrame())
+                exp = self._expected_tables.get("main", pd.DataFrame())
+                dirty = self._dirty_tables.get("main", pd.DataFrame())
+                return grade_task2(df, exp, dirty)
+
+            elif self._task_id == "task3":
+                exp = self._expected_tables.get("main", pd.DataFrame())
+                return grade_task3(self._tables, exp, self._dirty_tables)
+
+            elif self._task_id == "task4_data_drift":
+                df = self._tables.get("stream", pd.DataFrame())
+                return grade_task4(df)
+
+        except Exception:
+            pass
+
+        return 0.05  # safe floor if grader crashes
+
+    # ── Operation executor ────────────────────────────────────────────────────
+
     def _execute(self, action: DataCleanAction) -> str:
         op  = action.operation
-        tbl = action.table_name or ("stream" if self._task_id == "task4_data_drift" else "main")
+        tbl = action.table_name or (
+            "stream" if self._task_id == "task4_data_drift" else "main"
+        )
 
         if op == "fill_nulls":
             df  = self._tbl(tbl)
             col = self._col(df, action.column, tbl)
+
+            # Normalise string null markers → real NaN
             df[col] = df[col].replace(
-                ["N/A","n/a","null","NULL","None","none","missing","","NaN"], np.nan
+                ["N/A", "n/a", "null", "NULL", "None", "none",
+                 "missing", "", "NaN", "nan"], np.nan
             )
-            num = pd.to_numeric(df[col], errors="coerce")
-            s   = (action.strategy or "mean").lower()
-            if s == "mean":       fv = num.mean()
-            elif s == "median":   fv = num.median()
-            elif s == "mode":
-                modes = df[col].mode()
-                fv = modes.iloc[0] if not modes.empty else np.nan
-            elif s == "constant": fv = action.value
-            elif s in ("forward_fill","ffill"):
+
+            s = (action.strategy or "mean").lower()
+
+            if s in ("forward_fill", "ffill"):
                 df[col] = df[col].ffill()
                 self._tables[tbl] = df
                 return f"[fill_nulls] '{col}' ffill in '{tbl}'"
-            elif s in ("backward_fill","bfill"):
+            if s in ("backward_fill", "bfill"):
                 df[col] = df[col].bfill()
                 self._tables[tbl] = df
                 return f"[fill_nulls] '{col}' bfill in '{tbl}'"
+
+            # Compute fill value on numeric-coerced series for numeric strategies
+            num = pd.to_numeric(df[col], errors="coerce")
+            if s == "mean":
+                fv = num.mean()
+            elif s == "median":
+                fv = num.median()
+            elif s == "mode":
+                # Use numeric mode if column has any numeric data, else raw mode
+                if num.notna().any():
+                    modes = num.dropna()
+                    fv = modes.mode().iloc[0] if not modes.mode().empty else np.nan
+                else:
+                    modes = df[col].dropna()
+                    fv = modes.mode().iloc[0] if not modes.mode().empty else np.nan
+            elif s == "constant":
+                fv = action.value
             else:
                 raise ValueError(f"Unknown strategy '{s}'")
-            if pd.isna(fv):
-                raise ValueError(f"Cannot compute fill value for '{col}'")
+
+            if fv is None or (isinstance(fv, float) and np.isnan(fv)):
+                raise ValueError(f"Cannot compute fill value for '{col}' with strategy '{s}'")
+
+            # Fill NaN in the (already normalised) column using numeric coercion
             df[col] = pd.to_numeric(df[col], errors="coerce").fillna(fv)
             self._tables[tbl] = df
             return f"[fill_nulls] '{col}' → {s}={fv:.4g} in '{tbl}'"
@@ -214,11 +661,19 @@ class DataCleanEnvironment:
             df  = self._tbl(tbl)
             col = self._col(df, action.column, tbl)
             dt  = (action.dtype or "").lower()
+
             if dt == "int":
-                df[col] = pd.to_numeric(df[col], errors="coerce").fillna(0).astype("int64")
+                num = pd.to_numeric(df[col], errors="coerce")
+                # Only cast if no nulls remain — otherwise raise informatively
+                if num.isna().any():
+                    raise ValueError(
+                        f"Column '{col}' still has {int(num.isna().sum())} null(s). "
+                        f"Run fill_nulls first before casting to int."
+                    )
+                df[col] = num.astype("int64")
             elif dt == "float":
                 df[col] = pd.to_numeric(df[col], errors="coerce").astype("float64")
-            elif dt in ("str","string","object"):
+            elif dt in ("str", "string", "object"):
                 df[col] = df[col].astype(str)
             elif dt == "datetime":
                 df[col] = pd.to_datetime(df[col], errors="coerce")
@@ -226,32 +681,35 @@ class DataCleanEnvironment:
                 df[col] = df[col].astype(bool)
             else:
                 raise ValueError(f"Unknown dtype '{dt}'")
+
             self._tables[tbl] = df
             return f"[cast_column] '{col}' → {dt} in '{tbl}'"
 
         elif op == "remove_duplicates":
             df     = self._tbl(tbl)
             before = len(df)
-            keep   = action.keep if action.keep in ("first","last") else "first"
+            keep   = action.keep if action.keep in ("first", "last") else "first"
             df     = df.drop_duplicates(subset=action.subset, keep=keep).reset_index(drop=True)
             self._tables[tbl] = df
-            return f"[remove_duplicates] -{before-len(df)} rows from '{tbl}'"
+            return f"[remove_duplicates] -{before - len(df)} rows from '{tbl}'"
 
         elif op == "normalize_values":
             df  = self._tbl(tbl)
             col = self._col(df, action.column, tbl)
             m   = (action.method or "upper").lower()
             s   = df[col].astype(str)
-            if m == "upper":   df[col] = s.str.strip().str.upper()
-            elif m == "lower": df[col] = s.str.strip().str.lower()
-            elif m == "strip": df[col] = s.str.strip()
-            elif m == "title": df[col] = s.str.strip().str.title()
-            elif m in ("regex","replace"):
+            if m == "upper":      df[col] = s.str.strip().str.upper()
+            elif m == "lower":    df[col] = s.str.strip().str.lower()
+            elif m == "strip":    df[col] = s.str.strip()
+            elif m == "title":    df[col] = s.str.strip().str.title()
+            elif m in ("regex", "replace"):
                 if not action.pattern:
-                    raise ValueError(f"method='{m}' needs pattern")
-                df[col] = s.str.replace(action.pattern, action.replacement or "", regex=(m=="regex"))
+                    raise ValueError(f"method='{m}' requires 'pattern'")
+                df[col] = s.str.replace(
+                    action.pattern, action.replacement or "", regex=(m == "regex")
+                )
             else:
-                raise ValueError(f"Unknown method '{m}'")
+                raise ValueError(f"Unknown normalize method '{m}'")
             self._tables[tbl] = df
             return f"[normalize_values] '{col}' → {m} in '{tbl}'"
 
@@ -262,43 +720,60 @@ class DataCleanEnvironment:
             thr = action.threshold if action.threshold is not None else 1.5
             num = pd.to_numeric(df[col], errors="coerce")
             before = len(df)
+
             if m == "iqr":
                 Q1, Q3 = num.quantile(0.25), num.quantile(0.75)
                 IQR = Q3 - Q1
-                mask = (num >= Q1 - thr*IQR) & (num <= Q3 + thr*IQR)
+                mask = (num >= Q1 - thr * IQR) & (num <= Q3 + thr * IQR)
             elif m == "zscore":
                 z    = (num - num.mean()) / num.std()
                 mask = z.abs() <= thr
             else:
                 raise ValueError(f"Unknown outlier method '{m}'")
+
             df = df[mask | num.isna()].reset_index(drop=True)
             self._tables[tbl] = df
-            return f"[filter_outliers] -{before-len(df)} rows from '{col}' in '{tbl}'"
+            return f"[filter_outliers] -{before - len(df)} rows from '{col}' in '{tbl}'"
 
         elif op == "merge_tables":
             lt  = action.left_table  or "orders"
             rt  = action.right_table or "customers"
             on  = action.on          or "customer_id"
-            how = action.how if action.how in ("inner","left","right","outer") else "inner"
+            how = action.how if action.how in ("inner", "left", "right", "outer") else "inner"
             out = action.output_table or "merged"
             L, R = self._tbl(lt), self._tbl(rt)
-            if on not in L.columns: raise ValueError(f"Key '{on}' not in '{lt}'")
-            if on not in R.columns: raise ValueError(f"Key '{on}' not in '{rt}'")
+            if on not in L.columns:
+                raise ValueError(f"Key '{on}' not in '{lt}'. Available: {list(L.columns)}")
+            if on not in R.columns:
+                raise ValueError(f"Key '{on}' not in '{rt}'. Available: {list(R.columns)}")
             merged = pd.merge(L, R, on=on, how=how)
             self._tables[out] = merged
             return f"[merge_tables] '{lt}'×'{rt}' on '{on}' → '{out}' ({len(merged)} rows)"
 
         elif op == "add_derived_column":
-            tbl2 = action.table_name or "merged"
-            if tbl2 not in self._tables:
+            # Prefer action.table_name, then try "merged", then "main"
+            tbl2 = action.table_name
+            if tbl2 and tbl2 in self._tables:
+                pass  # use as-is
+            elif "merged" in self._tables:
+                tbl2 = "merged"
+            elif "main" in self._tables:
                 tbl2 = "main"
+            else:
+                raise KeyError(
+                    f"Table '{tbl2}' not found and no 'merged' or 'main' fallback. "
+                    f"Run merge_tables first. Available: {list(self._tables.keys())}"
+                )
+
             df  = self._tbl(tbl2)
             cn  = action.column_name
             src = action.source_column
             tr  = (action.transform or "").lower()
-            if not cn:  raise ValueError("'column_name' required")
-            if not src: raise ValueError("'source_column' required")
-            self._col(df, src, tbl2)
+
+            if not cn:  raise ValueError("'column_name' is required")
+            if not src: raise ValueError("'source_column' is required")
+            self._col(df, src, tbl2)  # validates column exists
+
             if tr == "year_from_date":
                 df[cn] = pd.to_datetime(df[src], errors="coerce").dt.year
             elif tr == "month_from_date":
@@ -309,18 +784,17 @@ class DataCleanEnvironment:
                 df[cn] = pd.to_numeric(df[src], errors="coerce").abs()
             elif tr == "len":
                 df[cn] = df[src].astype(str).str.len()
-            elif tr in ("upper","lower"):
+            elif tr in ("upper", "lower"):
                 df[cn] = getattr(df[src].astype(str).str, tr)()
             else:
                 raise ValueError(f"Unknown transform '{tr}'")
+
             self._tables[tbl2] = df
             return f"[add_derived_column] '{cn}'={tr}('{src}') in '{tbl2}'"
 
         raise ValueError(f"Unknown operation '{op}'")
 
-    def _score(self) -> float:
-        # Hardcoded safe score per task — always within 0.001–0.999
-        return TASK_SCORES.get(self._task_id, 0.501)
+    # ── Observation builder ───────────────────────────────────────────────────
 
     def _obs(self, reward: float, done: bool,
              new_rows: int = 0, score: Optional[float] = None) -> DataCleanObservation:
@@ -336,7 +810,7 @@ class DataCleanEnvironment:
             row_counts[nm]  = int(len(df))
 
         msg = self._last_msg
-        if new_rows > 0 and "DRIFT" not in msg:
+        if new_rows > 0 and "[DRIFT]" not in msg:
             msg = f"[+{new_rows} drift rows] " + msg
 
         return DataCleanObservation(
@@ -362,14 +836,18 @@ class DataCleanEnvironment:
         for nm, df in self._tables.items():
             for col in df.columns:
                 nc = int(df[col].isna().sum())
-                if nc: errs.append(f"{nm}.{col}: {nc} nulls")
+                if nc:
+                    errs.append(f"{nm}.{col}: {nc} nulls")
             dc = int(df.duplicated().sum())
-            if dc: errs.append(f"{nm}: {dc} duplicates")
+            if dc:
+                errs.append(f"{nm}: {dc} duplicates")
         return errs[:12]
 
     def _tbl(self, name: str) -> pd.DataFrame:
         if name not in self._tables:
-            raise KeyError(f"Table '{name}' not found. Available: {list(self._tables.keys())}")
+            raise KeyError(
+                f"Table '{name}' not found. Available: {list(self._tables.keys())}"
+            )
         return self._tables[name]
 
     @staticmethod
@@ -377,5 +855,7 @@ class DataCleanEnvironment:
         if not col:
             raise ValueError("'column' field is required")
         if col not in df.columns:
-            raise KeyError(f"Column '{col}' not in '{tbl}'. Available: {list(df.columns)}")
-        return col
+            raise KeyError(
+                f"Column '{col}' not in '{tbl}'. Available: {list(df.columns)}"
+            )
+        return col

server/graders.py

@@ -1,24 +1,305 @@
+# """
+# Graders — deterministic scoring for all 4 tasks.
+# All return float STRICTLY between 0.001 and 0.999 (never 0.0 or 1.0).
+# """
+# import pandas as pd
+# import numpy as np
+# from typing import Dict
+
+
+# def grade_task1(df: pd.DataFrame, expected_df: pd.DataFrame) -> float:
+#     return 0.501
+
+
+# def grade_task2(df: pd.DataFrame, expected_df: pd.DataFrame, dirty_df: pd.DataFrame) -> float:
+#     return 0.502
+
+
+# def grade_task3(tables: Dict[str, pd.DataFrame], expected_df: pd.DataFrame,
+#                 dirty_tables: Dict[str, pd.DataFrame]) -> float:
+#     return 0.503
+
+
+# def grade_task4(df: pd.DataFrame) -> float:
+#     return 0.504
+
 """
-Graders — deterministic scoring for all 4 tasks.
-All return float STRICTLY between 0.001 and 0.999 (never 0.0 or 1.0).
+Graders — deterministic partial-credit scoring for all 4 tasks.
+Every grader returns a float STRICTLY in (0.0, 1.0).
+  • Never returns exactly 0.0 — even a completely uncleaned table scores 0.05.
+  • Never returns exactly 1.0 — a perfect table scores 0.98.
+  • All intermediate states return a meaningful float between those bounds.
+
+This satisfies the OpenEnv validator requirement:
+  "one or more tasks returned a score outside [0, 1]"  ← was caused by hardcoded stubs.
+
+Grader design:
+  - Each grader checks multiple sub-dimensions with individual weights.
+  - Weights sum to 1.0 for each grader.
+  - Raw score is clipped to [0.05, 0.98] before return.
 """
 import pandas as pd
 import numpy as np
 from typing import Dict
 
 
+def _clamp(score: float) -> float:
+    """Clamp to strictly-open (0, 1) range required by OpenEnv validator."""
+    return float(max(0.05, min(0.98, score)))
+
+
+# ── Task 1 ────────────────────────────────────────────────────────────────────
+
 def grade_task1(df: pd.DataFrame, expected_df: pd.DataFrame) -> float:
-    return 0.501
+    """
+    Score a task1 'main' DataFrame.
+
+    Sub-dimensions (weights sum to 1.0):
+      - age nulls gone          0.30
+      - age dtype is int64      0.25
+      - age values close to expected  0.20
+      - salary nulls gone       0.15
+      - salary dtype is float64 0.10
+    """
+    score = 0.0
+
+    # age nulls (0.30)
+    age_nulls = int(df["age"].isna().sum()) if "age" in df.columns else 999
+    if age_nulls == 0:
+        score += 0.30
+    elif age_nulls <= 3:
+        score += 0.15
+
+    # age dtype (0.25)
+    if "age" in df.columns:
+        age_numeric = pd.to_numeric(df["age"], errors="coerce")
+        non_null_age = age_numeric.dropna()
+        if pd.api.types.is_integer_dtype(df["age"]):
+            score += 0.25
+        elif len(non_null_age) == len(df) and (non_null_age % 1 == 0).all():
+            # numeric but stored as float with no decimals — partial credit
+            score += 0.12
+
+    # age values accuracy (0.20) — compare median of cleaned vs expected
+    if "age" in df.columns and "age" in expected_df.columns:
+        try:
+            actual_med = pd.to_numeric(df["age"], errors="coerce").median()
+            exp_med = pd.to_numeric(expected_df["age"], errors="coerce").median()
+            if abs(actual_med - exp_med) < 1:
+                score += 0.20
+            elif abs(actual_med - exp_med) < 5:
+                score += 0.10
+        except Exception:
+            pass
+
+    # salary nulls (0.15)
+    sal_nulls = int(df["salary"].isna().sum()) if "salary" in df.columns else 999
+    if sal_nulls == 0:
+        score += 0.15
+    elif sal_nulls <= 3:
+        score += 0.07
+
+    # salary dtype (0.10)
+    if "salary" in df.columns and pd.api.types.is_float_dtype(df["salary"]):
+        score += 0.10
+
+    return _clamp(score)
+
+
+# ── Task 2 ────────────────────────────────────────────────────────────────────
+
+def grade_task2(df: pd.DataFrame, expected_df: pd.DataFrame,
+                dirty_df: pd.DataFrame) -> float:
+    """
+    Score a task2 'main' DataFrame.
+
+    Sub-dimensions (weights sum to 1.0):
+      - duplicates removed      0.25
+      - country normalised      0.25
+      - order_date is datetime  0.20
+      - amount nulls gone       0.15
+      - row count reasonable    0.15
+    """
+    score = 0.0
 
+    # duplicates (0.25)
+    dup_count = int(df.duplicated().sum())
+    if dup_count == 0:
+        score += 0.25
+    elif dup_count < 5:
+        score += 0.12
 
-def grade_task2(df: pd.DataFrame, expected_df: pd.DataFrame, dirty_df: pd.DataFrame) -> float:
-    return 0.502
+    # country upper-case (0.25)
+    if "country" in df.columns:
+        str_col = df["country"].dropna().astype(str)
+        total = len(str_col)
+        if total > 0:
+            upper_frac = (str_col == str_col.str.upper()).mean()
+            score += 0.25 * upper_frac
 
+    # order_date datetime (0.20)
+    if "order_date" in df.columns:
+        if pd.api.types.is_datetime64_any_dtype(df["order_date"]):
+            score += 0.20
+        else:
+            parsed = pd.to_datetime(df["order_date"], errors="coerce")
+            valid_frac = parsed.notna().mean()
+            score += 0.20 * valid_frac * 0.5  # partial: parsable but wrong dtype
+
+    # amount nulls (0.15)
+    if "amount" in df.columns:
+        amt_nulls = int(df["amount"].isna().sum())
+        if amt_nulls == 0:
+            score += 0.15
+        elif amt_nulls <= 3:
+            score += 0.07
+
+    # row count (0.15) — should be ≈170 (base) after dedup, not 200 (with dups)
+    n_rows = len(df)
+    dirty_rows = len(dirty_df)
+    expected_rows = len(expected_df)
+    if expected_rows > 0:
+        ratio = n_rows / expected_rows
+        if 0.85 <= ratio <= 1.15:
+            score += 0.15
+        elif 0.5 <= ratio <= 1.5:
+            score += 0.07
+
+    return _clamp(score)
+
+
+# ── Task 3 ────────────────────────────────────────────────────────────────────
 
 def grade_task3(tables: Dict[str, pd.DataFrame], expected_df: pd.DataFrame,
                 dirty_tables: Dict[str, pd.DataFrame]) -> float:
-    return 0.503
+    """
+    Score a task3 state.
+
+    Sub-dimensions (weights sum to 1.0):
+      - merged table exists            0.25
+      - outliers removed (amount IQR)  0.25
+      - age nulls gone in merged       0.20
+      - order_year column present      0.15
+      - row count in reasonable range  0.15
+    """
+    score = 0.0
+
+    # merged table (0.25)
+    merged_key = None
+    for k in ("merged", "main"):
+        if k in tables:
+            merged_key = k
+            break
+    if merged_key is None:
+        # No merge done yet — return base score
+        return _clamp(0.05)
+
+    merged = tables[merged_key]
+
+    score += 0.25  # merged table exists
+
+    # outliers removed (0.25) — check that extreme amounts are gone
+    if "amount" in merged.columns:
+        amt = pd.to_numeric(merged["amount"], errors="coerce").dropna()
+        if len(amt) > 0:
+            Q1, Q3 = amt.quantile(0.25), amt.quantile(0.75)
+            IQR = Q3 - Q1
+            lo, hi = Q1 - 1.5 * IQR, Q3 + 1.5 * IQR
+            outlier_frac = ((amt < lo) | (amt > hi)).mean()
+            if outlier_frac < 0.02:
+                score += 0.25
+            elif outlier_frac < 0.10:
+                score += 0.12
+
+    # age nulls (0.20)
+    if "age" in merged.columns:
+        age_nulls = int(merged["age"].isna().sum())
+        if age_nulls == 0:
+            score += 0.20
+        elif age_nulls <= 3:
+            score += 0.10
+
+    # order_year column (0.15)
+    if "order_year" in merged.columns:
+        yr = pd.to_numeric(merged["order_year"], errors="coerce")
+        valid_years = yr.between(2020, 2030).mean()
+        score += 0.15 * valid_years
 
+    # row count (0.15)
+    exp_rows = len(expected_df)
+    if exp_rows > 0:
+        ratio = len(merged) / exp_rows
+        if 0.80 <= ratio <= 1.20:
+            score += 0.15
+        elif 0.50 <= ratio <= 1.50:
+            score += 0.07
+
+    return _clamp(score)
+
+
+# ── Task 4 ────────────────────────────────────────────────────────────────────
 
 def grade_task4(df: pd.DataFrame) -> float:
-    return 0.504
+    """
+    Score a task4 'stream' DataFrame.
+
+    Since this task has continuous drift, we score the current *cleaned* state
+    of the stream table across multiple dimensions.
+
+    Sub-dimensions (weights sum to 1.0):
+      - amount nulls low        0.25
+      - amount dtype numeric    0.20
+      - outliers low            0.20
+      - category nulls low      0.15
+      - region nulls low        0.10
+      - event_ts parseable      0.10
+    """
+    score = 0.0
+    n = len(df)
+    if n == 0:
+        return _clamp(0.05)
+
+    # amount nulls (0.25)
+    if "amount" in df.columns:
+        amt = pd.to_numeric(df["amount"], errors="coerce")
+        null_frac = amt.isna().mean()
+        score += 0.25 * max(0.0, 1.0 - null_frac * 3)
+
+        # amount dtype numeric (0.20)
+        if pd.api.types.is_numeric_dtype(df["amount"]):
+            score += 0.20
+        elif null_frac < 0.10:
+            # mostly parseable even if still object
+            score += 0.10
+
+        # outliers (0.20) — negative or huge values
+        valid_amt = amt.dropna()
+        if len(valid_amt) > 0:
+            Q1, Q3 = valid_amt.quantile(0.25), valid_amt.quantile(0.75)
+            IQR = Q3 - Q1
+            lo, hi = Q1 - 1.5 * IQR, Q3 + 1.5 * IQR
+            outlier_frac = ((valid_amt < lo) | (valid_amt > hi)).mean()
+            if outlier_frac < 0.03:
+                score += 0.20
+            elif outlier_frac < 0.15:
+                score += 0.10
+
+    # category nulls (0.15)
+    if "category" in df.columns:
+        cat_null_frac = df["category"].isna().mean()
+        score += 0.15 * max(0.0, 1.0 - cat_null_frac * 3)
+
+    # region nulls (0.10)
+    if "region" in df.columns:
+        reg_null_frac = df["region"].isna().mean()
+        score += 0.10 * max(0.0, 1.0 - reg_null_frac * 3)
+
+    # event_ts parseable (0.10)
+    if "event_ts" in df.columns:
+        if pd.api.types.is_datetime64_any_dtype(df["event_ts"]):
+            score += 0.10
+        else:
+            parsed = pd.to_datetime(df["event_ts"], errors="coerce")
+            score += 0.10 * parsed.notna().mean()
+
+    return _clamp(score)