Upload environment.py

server/environment.py

@@ -214,8 +214,46 @@ class SQLQualityEnvironment:
         Returns a result preview in the observation. Does NOT end the episode.
         Awards small step_reward (0.05) for queries that target the right area.
         """
+        import re
+
         sql = (action.sql or "").strip()
 
+        # Strip comments and leading whitespace to get the first keyword
+        sql_stripped = re.sub(r"--[^\n]*", "", sql)  # remove single-line comments
+        sql_stripped = re.sub(r"/\*.*?\*/", "", sql_stripped, flags=re.DOTALL)  # block comments
+        first_keyword = sql_stripped.strip().split()[0].upper() if sql_stripped.strip() else ""
+
+        ALLOWED_READ_KEYWORDS = {"SELECT", "WITH", "EXPLAIN", "PRAGMA"}
+        if first_keyword not in ALLOWED_READ_KEYWORDS:
+            rejection_feedback = (
+                f"QUERY action only accepts read-only SQL. "
+                f"You submitted a '{first_keyword}' statement. "
+                f"To modify data, use the FIX action (action_type: 'fix'). "
+                f"To find the answer, use action_type: 'submit'."
+            )
+            # Add to trajectory as a step but with zero reward and clear feedback
+            self._trajectory.append(
+                TrajectoryStep(
+                    step_number=len(self._trajectory) + 1,
+                    action_type=ActionType.QUERY,
+                    sql_submitted=sql[:200],
+                    result_preview=rejection_feedback,
+                    step_reward=0.0,
+                )
+            )
+            self._step_count = getattr(self, "_step_count", 0) + 1
+            obs = self._build_observation(result_preview=rejection_feedback)
+            return StepResponse(
+                observation=obs,
+                reward=SQLReward(
+                    score=self._cumulative_reward,
+                    step_reward=0.0,
+                    feedback=rejection_feedback,
+                ),
+                done=False,
+                info={"action_type": "query", "error": "non_select_rejected", "step": len(self._trajectory)},
+            )
+
         # Split multiple statements by semicolon, but preserve semicolons in strings
         statements = []
         current = ""