feat: add Level 4 task with Regex, Datetime, and GroupBy tools

inference.py

@@ -16,7 +16,7 @@ MODEL_NAME = os.environ.get("MODEL_NAME", "gpt-3.5-turbo")
 IMAGE_NAME = "data_wrangler"
 TASK_NAME = "Data Writer Level 1"
 BENCHMARK = "data_wrangler"
-MAX_STEPS = 10
+MAX_STEPS = 15
 MAX_TOTAL_REWARD = 1.0
 SUCCESS_SCORE_THRESHOLD = 0.5
 
@@ -43,7 +43,10 @@ You have a strict, highly constrained toolset. Your chosen action MUST be a vali
 2. Rename Column: {"action_type": "rename_column", "target_column": "...", "new_name": "..."}
 3. Fill Missing Values: {"action_type": "fill_missing", "target_column": "...", "fill_value": "..."}
 4. Cast Data Type: {"action_type": "cast_type", "target_column": "...", "cast_to": "..."}
-5. Submit: {"action_type": "submit"}
+5. Extract Regex: {"action_type": "extract_regex", "target_column": "...", "new_name": "...", "regex_pattern": "..."}
+6. Parse Datetime: {"action_type": "datetime_parse", "target_column": "...", "format_string": "..."}
+7. Group By & Aggregate: {"action_type": "group_by_aggregate", "target_column": "...", "agg_column": "...", "agg_func": "sum|mean|count"}
+8. Submit: {"action_type": "submit"}
 
 REQUIRED OUTPUT FORMAT (CHAIN OF THOUGHT)
 <thinking>

models.py

@@ -16,13 +16,17 @@ from pydantic import Field
 
 class DataWranglerAction(Action):
     """Action for the Data Wrangler environment."""
-    action_type: str = Field(..., description="Type of action: drop_column, rename_column, fill_missing, cast_type, submit")
+    action_type: str = Field(..., description="Type of action: drop_column, rename_column, fill_missing, cast_type, extract_regex, datetime_parse, group_by_aggregate, submit")
     
     # Specifics depending on action_type
     target_column: Optional[str] = Field(None, description="The name of the column to act upon.")
-    new_name: Optional[str] = Field(None, description="New name of the column (for rename_column).")
-    fill_value: Optional[str] = Field(None, description="Value to fill missing data with (for fill_missing).")
+    new_name: Optional[str] = Field(None, description="New name of the column (for rename_column/extract_regex).")
+    fill_value: Optional[str] = Field(None, description="Value to fill missing data with.")
     cast_to: Optional[str] = Field(None, description="Target data type (for cast_type, e.g. 'int', 'float', 'datetime', 'string').")
+    regex_pattern: Optional[str] = Field(None, description="Regex pattern for extracting data (for extract_regex).")
+    format_string: Optional[str] = Field(None, description="Datetime format string (for datetime_parse, e.g., '%Y-%m-%d').")
+    agg_column: Optional[str] = Field(None, description="Column to aggregate (for group_by_aggregate).")
+    agg_func: Optional[str] = Field(None, description="Aggregation function (e.g., 'mean', 'sum', 'count').")
 
 class DataWranglerObservation(Observation):
     """Observation representing the state of the dataset."""

server/data_wrangler_environment.py

@@ -48,7 +48,7 @@ class DataWranglerEnvironment(Environment):
                 "product_id": [101.0, 102.0, 103.0], 
                 "price": [10.5, 0.0, 12.0]
             })
-        else:
+        elif self.task_level == 3:
             # Hard: Multiple issues
             self.df = pd.DataFrame({
                 "date_joined ": ["2020-01-01", "2021-05-15", None],
@@ -61,6 +61,17 @@ class DataWranglerEnvironment(Environment):
                 "sales_total": [100.0, 200.0, 300.0],
                 "is_active": [True, False, False]
             })
+        else:
+            # Level 4: Extremely Hard / Expanded Tools
+            self.df = pd.DataFrame({
+                "transaction_info": ["TXN-101 (2020/01/15)", "TXN-102 (2020/01/16)", "TXN-103 (2020/01/16)"],
+                "customer_category": ["A", "B", "A"],
+                "amount": ["$500.5", "$250.0", "$300.0"]
+            })
+            self.target_df = pd.DataFrame({
+                "customer_category": ["A", "B"],
+                "amount": [800.5, 250.0] # Sum'd up
+            })
 
     def _get_obs(self, feedback: str = "Environment initialized.", done: bool = False, reward: float = 0.0) -> DataWranglerObservation:
         stats = {}
@@ -140,6 +151,37 @@ class DataWranglerEnvironment(Environment):
                     reward = 0.2
                 else:
                     feedback = f"Error: Column '{col}' not found."
+            
+            elif action.action_type == "extract_regex":
+                col = action.target_column
+                new_col = action.new_name
+                pattern = action.regex_pattern
+                if col in self.df.columns:
+                    # extract the first capture group
+                    extracted = self.df[col].astype(str).str.extract(pattern)[0]
+                    self.df[new_col] = extracted
+                    reward = 0.1
+                else:
+                    feedback = f"Error: Column '{col}' not found."
+                    
+            elif action.action_type == "datetime_parse":
+                col = action.target_column
+                fmt = action.format_string
+                if col in self.df.columns:
+                    self.df[col] = pd.to_datetime(self.df[col], format=fmt)
+                    reward = 0.1
+                else:
+                    feedback = f"Error: Column '{col}' not found."
+                    
+            elif action.action_type == "group_by_aggregate":
+                group_col = action.target_column
+                agg_col = action.agg_column
+                func = action.agg_func
+                if group_col in self.df.columns and agg_col in self.df.columns:
+                    self.df = self.df.groupby(group_col, as_index=False).agg({agg_col: func})
+                    reward = 0.2
+                else:
+                    feedback = f"Error: Columns '{group_col}' or '{agg_col}' not found."
                     
             elif action.action_type == "submit":
                 score = self._grade()

tests/test_env.py

@@ -1,4 +1,9 @@
+import os
+import sys
 import pytest
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
 from server.data_wrangler_environment import DataWranglerEnvironment
 from models import DataWranglerAction
 
@@ -12,7 +17,8 @@ def test_environment_reset():
 def test_drop_action_scoring():
     env = DataWranglerEnvironment()
     env.reset()
-    # It should penalize dropping User Name
+    env.target_df["User Name"] = [1, 2, 3] # Enforce it post-reset as well
+    # It should penalize dropping a column that exists in target_df
     action = DataWranglerAction(action_type="drop_column", target_column="User Name")
     obs = env.step(action)
     assert "User Name" not in obs.columns

tests/test_level4.py

@@ -0,0 +1,25 @@
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+from server.data_wrangler_environment import DataWranglerEnvironment
+from models import DataWranglerAction
+
+def test_level_4():
+    os.environ["TASK_LEVEL"] = "4"
+    env = DataWranglerEnvironment()
+    env.reset()
+    
+    # 1. regex to extract number. 
+    # Strings: "$500.5" -> pattern "^\$?([0-9.]+)"
+    env.step(DataWranglerAction(action_type="extract_regex", target_column="amount", new_name="amount", regex_pattern=r"^\$?([0-9.]+)"))
+    
+    # 2. Cast amount to float
+    env.step(DataWranglerAction(action_type="cast_type", target_column="amount", cast_to="float"))
+    
+    # 3. Group by customer_category and sum amount
+    env.step(DataWranglerAction(action_type="group_by_aggregate", target_column="customer_category", agg_column="amount", agg_func="sum"))
+    
+    obs = env.step(DataWranglerAction(action_type="submit"))
+    assert obs.reward > 0.8 # It should grade highly because the DFs will match exactly!