training-monkey/dataoncallenv
0
1"""Main RL Environment implementation for DataOnCallEnv.2"""3 4import logging5import traceback6from models import Action, Observation, Reward, RewardBreakdown, EnvState7from tasks import get_task8from graders import grade9from database import (10 build_task1_db, build_task2_db, build_task3_db,11 run_sql, inspect_schema, check_logs, check_airflow,12 diff_report, list_tables13)14import re15 16logger = logging.getLogger(__name__)17 18 19def _safe_observation(task_id, steps_taken, done, max_steps, cost_spent, budget_remaining, result):20 """Build an Observation object, never raising."""21 try:22 return Observation(23 task_id=task_id or 0,24 result=result,25 steps_taken=steps_taken,26 done=done,27 max_steps=max_steps,28 cost_spent=cost_spent,29 budget_remaining=budget_remaining,30 )31 except Exception:32 return Observation(33 task_id=0,34 result={"result": "internal error", "success": False},35 steps_taken=0,36 done=True,37 max_steps=15,38 cost_spent=0.0,39 budget_remaining=0.0,40 )41 42 43def _error_reward(msg: str = "error occurred") -> Reward:44 """Return a safe zero-score Reward."""45 return Reward(46 score=0.0,47 breakdown=RewardBreakdown(48 diagnosis_correct=0.0,49 fix_valid=0.0,50 efficiency=0.0,51 reasoning_quality=0.0,52 investigation_quality=0.0,53 ),54 false_positive_penalty=0.0,55 )56 57 58class DataOnCallEnv:59 60 MAX_STEPS = 1561 62 # Configuration63 TOOL_COSTS = {64 "list_tables": 0.5,65 "inspect_schema": 1.0,66 "check_logs": 1.0,67 "check_airflow": 1.0,68 "run_sql": 2.0,69 "diff_report": 1.5,70 "submit": 0.0,71 }72 COST_BUDGET = 20.073 74 VALID_TOOLS = set(TOOL_COSTS.keys())75 76 # Minimum tool calls before submit() is allowed (anti-cheat)77 MIN_STEPS_BEFORE_SUBMIT = 278 79 def __init__(self):80 self.task_id = None81 self.conn = None82 self.steps_taken = 083 self.done = False84 self.actions = []85 self.last_observation = None86 self.final_answer = ""87 self.discovered_tables = set() # partial observability tracking88 self.cost_spent = 0.089 90 # ── reset() ───────────────────────────────────────────────────────────────91 92 def reset(self, task_id: int = 1) -> Observation:93 """Fresh episode. Rebuilds database. Returns scenario as first observation."""94 try:95 task_id = int(task_id)96 except (TypeError, ValueError):97 task_id = 198 99 if task_id not in (1, 2, 3):100 raise ValueError(f"task_id must be 1, 2, or 3. Got {task_id}")101 102 self.task_id = task_id103 self.steps_taken = 0104 self.done = False105 self.actions = []106 self.final_answer = ""107 self.discovered_tables = set()108 self.cost_spent = 0.0109 110 builders = {1: build_task1_db, 2: build_task2_db, 3: build_task3_db}111 self.conn = builders[task_id]()112 113 task = get_task(task_id)114 115 # PARTIAL OBSERVABILITY: no table listing in initial observation116 # Agent must call list_tables() to discover what's available117 obs = Observation(118 task_id=task_id,119 result={120 "scenario": task["description"],121 "available_tools": {122 "run_sql": "run_sql(query) — SELECT query (specify columns, no SELECT *)",123 "inspect_schema": "inspect_schema(table) — column names + types (discover tables first)",124 "check_logs": "check_logs() — dbt pipeline changelog",125 "check_airflow": "check_airflow() — Airflow DAG run history",126 "diff_report": "diff_report(d1,d2) — compare two dates, query='date1,date2'",127 "list_tables": "list_tables() — discover available tables (START HERE)",128 "submit": "submit(answer) — your final answer. CALL THIS WHEN DONE.",129 },130 "rules": {131 "step_budget": self.MAX_STEPS,132 "cost_budget": self.COST_BUDGET,133 "select_star": "SELECT * is blocked. Specify columns explicitly.",134 "discovery": "You must discover tables with list_tables() before inspecting or querying them.",135 "submit_rule": f"You must use at least {self.MIN_STEPS_BEFORE_SUBMIT} tools before submitting.",136 },137 },138 steps_taken=0,139 done=False,140 max_steps=self.MAX_STEPS,141 cost_spent=0.0,142 budget_remaining=self.COST_BUDGET,143 )144 self.last_observation = obs145 return obs146 147 # ── step() ────────────────────────────────────────────────────────────────148 149 def step(self, action: Action):150 """151 Execute one action. Returns (observation, reward, done, info).152 reward is None until done=True.153 Enforces: partial observability, query costs, anti-cheat rules.154 155 DEFENSIVE: never raises — always returns (obs, reward, done, info).156 """157 # ── Top-level safety net ─────────────────────────────────────────────158 try:159 return self._step_impl(action)160 except Exception as exc:161 logger.error("Unhandled exception in step(): %s\n%s", exc, traceback.format_exc())162 err_msg = f"Internal environment error: {exc}"163 obs = _safe_observation(164 task_id=self.task_id,165 steps_taken=self.steps_taken,166 done=True,167 max_steps=self.MAX_STEPS,168 cost_spent=self.cost_spent,169 budget_remaining=max(0.0, self.COST_BUDGET - self.cost_spent),170 result={"result": err_msg, "success": False},171 )172 self.done = True173 self.last_observation = obs174 reward = _error_reward(err_msg)175 return obs, reward, True, {}176 177 def _step_impl(self, action: Action):178 """Inner step logic. All exceptions bubble to step() for safe handling."""179 180 # ── Validate / normalise action ──────────────────────────────────────181 # Guard against missing or invalid tool182 tool = None183 try:184 tool = str(action.tool).strip() if action.tool is not None else ""185 except Exception:186 tool = ""187 188 # Guard against None / non-string query189 query = ""190 try:191 if action.query is not None:192 query = str(action.query).strip()193 except Exception:194 query = ""195 196 # Episode-over guard — return safe error obs, NOT raise197 if self.done:198 obs = _safe_observation(199 task_id=self.task_id,200 steps_taken=self.steps_taken,201 done=True,202 max_steps=self.MAX_STEPS,203 cost_spent=self.cost_spent,204 budget_remaining=max(0.0, self.COST_BUDGET - self.cost_spent),205 result={"result": "Episode is already over. Call reset() to start a new episode.", "success": False},206 )207 return obs, _error_reward("episode already done"), True, self._make_info(True)208 209 # Validate tool name up front210 if not tool or tool not in self.VALID_TOOLS:211 obs = _safe_observation(212 task_id=self.task_id,213 steps_taken=self.steps_taken,214 done=False,215 max_steps=self.MAX_STEPS,216 cost_spent=self.cost_spent,217 budget_remaining=self.COST_BUDGET - self.cost_spent,218 result={"error": f"Unknown tool '{tool}'. Valid: {sorted(self.VALID_TOOLS)}"},219 )220 self.last_observation = obs221 return obs, None, False, self._make_info(False)222 223 # Anti-cheat: block early submit224 if tool == "submit" and self.steps_taken < self.MIN_STEPS_BEFORE_SUBMIT:225 obs = _safe_observation(226 task_id=self.task_id,227 steps_taken=self.steps_taken,228 done=False,229 max_steps=self.MAX_STEPS,230 cost_spent=self.cost_spent,231 budget_remaining=self.COST_BUDGET - self.cost_spent,232 result={233 "error": (234 f"Cannot submit yet. You must investigate first. "235 f"Use at least {self.MIN_STEPS_BEFORE_SUBMIT} tools before submitting. "236 f"Steps taken so far: {self.steps_taken}."237 )238 },239 )240 self.last_observation = obs241 return obs, None, False, self._make_info(False)242 243 # Cost tracking244 tool_cost = self.TOOL_COSTS.get(tool, 0.0)245 self.cost_spent += tool_cost246 247 self.actions.append(self._safe_action_dump(action, tool, query))248 self.steps_taken += 1249 250 result = self._run_tool_safe(tool, query)251 252 # Check if budget (cost or steps) exhausted — auto-submit253 budget_exhausted = self.cost_spent >= self.COST_BUDGET254 steps_exhausted = self.steps_taken >= self.MAX_STEPS255 256 if (budget_exhausted or steps_exhausted) and not self.done:257 if not self.final_answer:258 try:259 all_reasoning = " | ".join(260 a.get("reasoning", "") for a in self.actions if a.get("reasoning")261 )262 all_queries = " | ".join(263 a.get("query", "") for a in self.actions if a.get("tool") == "run_sql"264 )265 exhaust_type = "cost budget" if budget_exhausted else "step budget"266 self.final_answer = (267 f"AUTO-SUBMITTED ({exhaust_type} exhausted). "268 f"Reasoning: {all_reasoning}. Queries run: {all_queries}"269 )270 except Exception:271 self.final_answer = "AUTO-SUBMITTED (budget exhausted)"272 273 episode_done = tool == "submit" or budget_exhausted or steps_exhausted274 self.done = episode_done275 276 obs = _safe_observation(277 task_id=self.task_id,278 steps_taken=self.steps_taken,279 done=episode_done,280 max_steps=self.MAX_STEPS,281 cost_spent=self.cost_spent,282 budget_remaining=max(0.0, self.COST_BUDGET - self.cost_spent),283 result=result,284 )285 self.last_observation = obs286 287 reward = None288 if episode_done:289 try:290 reward = grade(291 task_id=self.task_id,292 conn=self.conn,293 actions=self.actions,294 final_answer=self.final_answer,295 )296 # Ensure reward is always a valid Reward object297 if reward is None:298 reward = _error_reward("grade() returned None")299 except Exception as exc:300 logger.error("grade() raised: %s\n%s", exc, traceback.format_exc())301 reward = _error_reward(f"grading error: {exc}")302 303 return obs, reward, episode_done, self._make_info(episode_done)304 305 # ── state() ───────────────────────────────────────────────────────────────306 307 def state(self) -> EnvState:308 return EnvState(309 task_id=self.task_id or 0,310 steps_taken=self.steps_taken,311 done=self.done,312 agent_actions=self.actions,313 current_observation=self.last_observation,314 discovered_tables=sorted(self.discovered_tables),315 cost_spent=self.cost_spent,316 budget_remaining=max(0.0, self.COST_BUDGET - self.cost_spent),317 )318 319 # ── Info helper ───────────────────────────────────────────────────────────320 321 def _make_info(self, done: bool) -> dict:322 try:323 return {324 "steps_remaining": self.MAX_STEPS - self.steps_taken,325 "task_id": self.task_id,326 "done": done,327 "cost_spent": self.cost_spent,328 "budget_remaining": max(0.0, self.COST_BUDGET - self.cost_spent),329 "discovered_tables": sorted(self.discovered_tables),330 }331 except Exception:332 return {"done": done}333 334 # ── Safe action dump ──────────────────────────────────────────────────────335 336 def _safe_action_dump(self, action: Action, tool: str, query: str) -> dict:337 """Dump action to dict without crashing."""338 try:339 d = action.model_dump()340 # Ensure normalised tool/query are stored341 d["tool"] = tool342 d["query"] = query343 return d344 except Exception:345 return {"tool": tool, "query": query, "reasoning": None}346 347 # ── Partial observability checks ──────────────────────────────────────────348 349 def _check_table_access(self, query: str) -> str | None:350 """351 Check if the SQL query references any undiscovered tables.352 Returns error message if violation found, None if OK.353 """354 if not query:355 return None356 if not self.discovered_tables:357 return (358 "No tables discovered yet. Call list_tables() first to discover "359 "available tables before running SQL queries."360 )361 362 # Extract table names from the query (simple heuristic)363 q_upper = query.upper()364 tokens = re.findall(r'(?:FROM|JOIN)\s+(\w+)', q_upper, re.IGNORECASE)365 366 for token in tokens:367 if token.lower() not in {t.lower() for t in self.discovered_tables}:368 return (369 f"Table '{token.lower()}' has not been discovered yet. "370 f"Discovered tables: {sorted(self.discovered_tables)}. "371 f"Use list_tables() to discover tables first."372 )373 return None374 375 # ── Tool router ───────────────────────────────────────────────────────────376 377 def _run_tool_safe(self, tool: str, query: str) -> dict:378 """Run a tool, returning a safe dict result. Never raises."""379 try:380 return self._run_tool(tool, query)381 except Exception as exc:382 logger.error("Tool '%s' raised: %s\n%s", tool, exc, traceback.format_exc())383 return {"error": f"Tool '{tool}' failed: {exc}", "success": False}384 385 def _run_tool(self, tool: str, query: str) -> dict:386 """Route a validated tool call. Raises on internal errors (caught by _run_tool_safe)."""387 388 if tool == "list_tables":389 tables = list_tables(self.conn)390 if not isinstance(tables, list):391 tables = []392 self.discovered_tables = set(tables)393 return {394 "tables": tables,395 "message": f"Found {len(tables)} tables. Use inspect_schema(table_name) to see columns."396 }397 398 elif tool == "inspect_schema":399 if not query:400 return {"error": "inspect_schema requires a table name as query."}401 if query.lower() not in {t.lower() for t in self.discovered_tables}:402 return {403 "error": (404 f"Table '{query}' not discovered yet. Call list_tables() first. "405 f"Discovered so far: {sorted(self.discovered_tables)}"406 )407 }408 result = inspect_schema(self.conn, query)409 return result if isinstance(result, dict) else {"result": str(result)}410 411 elif tool == "run_sql":412 if not query:413 return {"error": "run_sql requires a SQL query string."}414 access_error = self._check_table_access(query)415 if access_error:416 return {"error": access_error}417 result = run_sql(self.conn, query)418 # Ensure result is always a dict, never None419 if result is None:420 return {"rows": [], "message": "Query returned no results."}421 if not isinstance(result, (dict, list)):422 return {"result": str(result)}423 return result424 425 elif tool == "check_logs":426 result = check_logs(self.conn)427 if result is None:428 return {"logs": [], "message": "No logs found."}429 return result if isinstance(result, dict) else {"result": str(result)}430 431 elif tool == "check_airflow":432 result = check_airflow(self.conn)433 if result is None:434 return {"runs": [], "message": "No Airflow runs found."}435 return result if isinstance(result, dict) else {"result": str(result)}436 437 elif tool == "diff_report":438 if not query:439 return {"error": "diff_report requires 'date1,date2' as query."}440 parts = [p.strip() for p in query.split(",")]441 if len(parts) != 2:442 return {"error": "diff_report query must be 'date1,date2'"}443 result = diff_report(self.conn, parts[0], parts[1])444 if result is None:445 return {"diff": {}, "message": "No diff data returned."}446 return result if isinstance(result, dict) else {"result": str(result)}447 448 elif tool == "submit":449 self.final_answer = query450 return {451 "message": "Answer submitted. Grading now.",452 "steps_used": self.steps_taken,453 "cost_spent": self.cost_spent,454 "your_answer_preview": query[:200] if query else "(empty answer)",455 }456 457 else:458 return {"error": f"Unknown tool '{tool}'."}