Spaces:
Sleeping
Sleeping
| """ | |
| Cleaning Engine for DataWranglerEnv. | |
| Parses text commands from the agent and executes data cleaning operations | |
| on the working DataFrame. Returns text results for the observation. | |
| Commands: | |
| Diagnostic: help, view, profile, profile_column, find_missing, | |
| find_duplicates, find_outliers | |
| Cleaning: fill_missing, remove_duplicates, fix_dtype, replace, | |
| regex_replace, standardize, remove_rows, clip, | |
| rename_column, drop_column, sort | |
| Special: validate, submit | |
| """ | |
| import re | |
| from typing import Any, Dict, Optional, Tuple | |
| import numpy as np | |
| import pandas as pd | |
| from .dataset_generator import COMMANDS_HELP | |
| class CleaningEngine: | |
| """Parses agent text commands and applies cleaning operations to a DataFrame.""" | |
| def __init__(self, df: pd.DataFrame): | |
| self.df = df | |
| def execute(self, command_str: str) -> Tuple[str, bool]: | |
| """Parse and execute a text command. | |
| Args: | |
| command_str: Raw text command from the agent | |
| Returns: | |
| Tuple of (response_text, data_was_modified) | |
| """ | |
| command_str = command_str.strip() | |
| if not command_str: | |
| return "Error: Empty command. Type 'help' for available commands.", False | |
| parts = self._parse_command(command_str) | |
| cmd = parts[0].lower() | |
| args = parts[1:] | |
| dispatch = { | |
| "help": self._cmd_help, | |
| "view": self._cmd_view, | |
| "profile": self._cmd_profile, | |
| "profile_column": self._cmd_profile_column, | |
| "find_missing": self._cmd_find_missing, | |
| "find_duplicates": self._cmd_find_duplicates, | |
| "find_outliers": self._cmd_find_outliers, | |
| "fill_missing": self._cmd_fill_missing, | |
| "remove_duplicates": self._cmd_remove_duplicates, | |
| "fix_dtype": self._cmd_fix_dtype, | |
| "replace": self._cmd_replace, | |
| "regex_replace": self._cmd_regex_replace, | |
| "standardize": self._cmd_standardize, | |
| "remove_rows": self._cmd_remove_rows, | |
| "clip": self._cmd_clip, | |
| "rename_column": self._cmd_rename_column, | |
| "drop_column": self._cmd_drop_column, | |
| "sort": self._cmd_sort, | |
| "validate": self._cmd_validate, | |
| "submit": self._cmd_submit, | |
| } | |
| handler = dispatch.get(cmd) | |
| if handler is None: | |
| suggestions = [c for c in dispatch.keys() if c.startswith(cmd[:3])] if len(cmd) >= 3 else [] | |
| msg = f"Error: Unknown command '{cmd}'." | |
| if suggestions: | |
| msg += f" Did you mean: {', '.join(suggestions)}?" | |
| msg += " Type 'help' for available commands." | |
| return msg, False | |
| try: | |
| return handler(args) | |
| except Exception as e: | |
| return f"Error executing '{cmd}': {str(e)}", False | |
| def _parse_command(self, command_str: str) -> list: | |
| """Parse command string, handling quoted arguments.""" | |
| parts = [] | |
| current = "" | |
| in_quotes = False | |
| quote_char = None | |
| for char in command_str: | |
| if char in ('"', "'") and not in_quotes: | |
| in_quotes = True | |
| quote_char = char | |
| elif char == quote_char and in_quotes: | |
| in_quotes = False | |
| quote_char = None | |
| elif char == " " and not in_quotes: | |
| if current: | |
| parts.append(current) | |
| current = "" | |
| else: | |
| current += char | |
| if current: | |
| parts.append(current) | |
| return parts | |
| # ── Diagnostic commands (read-only) ────────────────────────────────── | |
| def _cmd_help(self, args: list) -> Tuple[str, bool]: | |
| return COMMANDS_HELP, False | |
| def _cmd_view(self, args: list) -> Tuple[str, bool]: | |
| n = 10 | |
| if args: | |
| try: | |
| n = int(args[0]) | |
| except ValueError: | |
| return "Error: 'view' expects an integer argument. Usage: view [N]", False | |
| n = min(n, 50) | |
| result = self.df.head(n).to_string(max_colwidth=30) | |
| return f"Showing first {n} rows ({len(self.df)} total):\n\n{result}", False | |
| def _cmd_profile(self, args: list) -> Tuple[str, bool]: | |
| lines = [] | |
| lines.append(f"Dataset Shape: {self.df.shape[0]} rows × {self.df.shape[1]} columns") | |
| lines.append(f"\nColumns:") | |
| lines.append(f"{'Column':<20} {'Type':<12} {'Non-Null':<10} {'Missing':<10} {'Missing%':<10} {'Unique':<8}") | |
| lines.append("-" * 70) | |
| for col in self.df.columns: | |
| dtype = str(self.df[col].dtype) | |
| non_null = self.df[col].notna().sum() | |
| missing = self.df[col].isna().sum() | |
| missing_pct = f"{(missing / len(self.df) * 100):.1f}%" | |
| unique = self.df[col].nunique() | |
| lines.append(f"{col:<20} {dtype:<12} {non_null:<10} {missing:<10} {missing_pct:<10} {unique:<8}") | |
| n_dupes = self.df.duplicated().sum() | |
| lines.append(f"\nDuplicate rows: {n_dupes} ({n_dupes / len(self.df) * 100:.1f}%)") | |
| return "\n".join(lines), False | |
| def _cmd_profile_column(self, args: list) -> Tuple[str, bool]: | |
| if not args: | |
| return "Error: Usage: profile_column COLUMN_NAME", False | |
| col = args[0] | |
| if col not in self.df.columns: | |
| return f"Error: Column '{col}' not found. Available: {', '.join(self.df.columns)}", False | |
| lines = [f"Profile for column '{col}':"] | |
| series = self.df[col] | |
| lines.append(f" Type: {series.dtype}") | |
| lines.append(f" Non-null: {series.notna().sum()} / {len(series)}") | |
| lines.append(f" Missing: {series.isna().sum()} ({series.isna().mean() * 100:.1f}%)") | |
| lines.append(f" Unique values: {series.nunique()}") | |
| if pd.api.types.is_numeric_dtype(series): | |
| desc = series.describe() | |
| lines.append(f" Min: {desc.get('min', 'N/A')}") | |
| lines.append(f" Max: {desc.get('max', 'N/A')}") | |
| lines.append(f" Mean: {desc.get('mean', 'N/A'):.2f}" if pd.notna(desc.get('mean')) else " Mean: N/A") | |
| lines.append(f" Std: {desc.get('std', 'N/A'):.2f}" if pd.notna(desc.get('std')) else " Std: N/A") | |
| lines.append(f" Median: {desc.get('50%', 'N/A')}") | |
| else: | |
| top_values = series.dropna().value_counts().head(10) | |
| lines.append(f" Top values:") | |
| for val, count in top_values.items(): | |
| lines.append(f" '{val}': {count}") | |
| return "\n".join(lines), False | |
| def _cmd_find_missing(self, args: list) -> Tuple[str, bool]: | |
| missing = self.df.isnull().sum() | |
| missing = missing[missing > 0] | |
| if missing.empty: | |
| return "No missing values found! The dataset is complete.", False | |
| lines = ["Missing values by column:"] | |
| lines.append(f"{'Column':<25} {'Count':<8} {'Percentage':<10}") | |
| lines.append("-" * 43) | |
| for col, count in missing.sort_values(ascending=False).items(): | |
| pct = f"{count / len(self.df) * 100:.1f}%" | |
| lines.append(f"{col:<25} {count:<8} {pct:<10}") | |
| lines.append(f"\nTotal missing cells: {missing.sum()}") | |
| return "\n".join(lines), False | |
| def _cmd_find_duplicates(self, args: list) -> Tuple[str, bool]: | |
| subset = None | |
| if args: | |
| subset = [c.strip() for c in args[0].split(",")] | |
| invalid = [c for c in subset if c not in self.df.columns] | |
| if invalid: | |
| return f"Error: Unknown columns: {invalid}. Available: {list(self.df.columns)}", False | |
| dupes = self.df[self.df.duplicated(subset=subset, keep=False)] | |
| n_dupes = self.df.duplicated(subset=subset, keep="first").sum() | |
| if n_dupes == 0: | |
| cols_desc = f" (on columns: {subset})" if subset else "" | |
| return f"No duplicate rows found{cols_desc}.", False | |
| lines = [f"Found {n_dupes} duplicate rows (keeping first occurrence):"] | |
| if len(dupes) <= 20: | |
| lines.append(dupes.to_string(max_colwidth=25)) | |
| else: | |
| lines.append(f"Showing first 10 of {len(dupes)} duplicate entries:") | |
| lines.append(dupes.head(10).to_string(max_colwidth=25)) | |
| return "\n".join(lines), False | |
| def _cmd_find_outliers(self, args: list) -> Tuple[str, bool]: | |
| if not args: | |
| return "Error: Usage: find_outliers COLUMN_NAME", False | |
| col = args[0] | |
| if col not in self.df.columns: | |
| return f"Error: Column '{col}' not found. Available: {', '.join(self.df.columns)}", False | |
| try: | |
| numeric_col = pd.to_numeric(self.df[col], errors="coerce") | |
| except Exception: | |
| return f"Error: Column '{col}' cannot be converted to numeric for outlier detection.", False | |
| q1 = numeric_col.quantile(0.25) | |
| q3 = numeric_col.quantile(0.75) | |
| iqr = q3 - q1 | |
| lower = q1 - 1.5 * iqr | |
| upper = q3 + 1.5 * iqr | |
| outliers_mask = (numeric_col < lower) | (numeric_col > upper) | |
| n_outliers = outliers_mask.sum() | |
| if n_outliers == 0: | |
| return f"No outliers found in '{col}' (IQR method, bounds: [{lower:.2f}, {upper:.2f}]).", False | |
| lines = [f"Found {n_outliers} outliers in '{col}' (IQR method):"] | |
| lines.append(f" Q1: {q1:.2f}, Q3: {q3:.2f}, IQR: {iqr:.2f}") | |
| lines.append(f" Lower bound: {lower:.2f}") | |
| lines.append(f" Upper bound: {upper:.2f}") | |
| outlier_values = numeric_col[outliers_mask].dropna() | |
| lines.append(f" Outlier values: {list(outlier_values.head(15).values)}") | |
| return "\n".join(lines), False | |
| # ── Cleaning commands (modify data) ────────────────────────────────── | |
| def _cmd_fill_missing(self, args: list) -> Tuple[str, bool]: | |
| if len(args) < 2: | |
| return "Error: Usage: fill_missing COLUMN STRATEGY [VALUE]\n Strategies: mean, median, mode, constant, forward_fill", False | |
| col = args[0] | |
| strategy = args[1].lower() | |
| if col not in self.df.columns: | |
| return f"Error: Column '{col}' not found. Available: {', '.join(self.df.columns)}", False | |
| n_before = self.df[col].isna().sum() | |
| if n_before == 0: | |
| return f"No missing values in '{col}'. Nothing to fill.", False | |
| if strategy == "mean": | |
| try: | |
| fill_val = pd.to_numeric(self.df[col], errors="coerce").mean() | |
| self.df[col] = pd.to_numeric(self.df[col], errors="coerce").fillna(fill_val) | |
| except Exception: | |
| return f"Error: Cannot compute mean for non-numeric column '{col}'.", False | |
| elif strategy == "median": | |
| try: | |
| fill_val = pd.to_numeric(self.df[col], errors="coerce").median() | |
| self.df[col] = pd.to_numeric(self.df[col], errors="coerce").fillna(fill_val) | |
| except Exception: | |
| return f"Error: Cannot compute median for non-numeric column '{col}'.", False | |
| elif strategy == "mode": | |
| mode_val = self.df[col].mode() | |
| if mode_val.empty: | |
| return f"Error: No mode found for '{col}'.", False | |
| self.df[col] = self.df[col].fillna(mode_val.iloc[0]) | |
| elif strategy == "constant": | |
| if len(args) < 3: | |
| return "Error: 'constant' strategy requires a VALUE. Usage: fill_missing COL constant VALUE", False | |
| fill_val = args[2] | |
| self.df[col] = self.df[col].fillna(fill_val) | |
| elif strategy == "forward_fill": | |
| self.df[col] = self.df[col].ffill() | |
| else: | |
| return f"Error: Unknown strategy '{strategy}'. Use: mean, median, mode, constant, forward_fill", False | |
| n_after = self.df[col].isna().sum() | |
| filled = n_before - n_after | |
| return f"Filled {filled} missing values in '{col}' using strategy '{strategy}'. Remaining: {n_after}", True | |
| def _cmd_remove_duplicates(self, args: list) -> Tuple[str, bool]: | |
| subset = None | |
| keep = "first" | |
| if args: | |
| subset = [c.strip() for c in args[0].split(",")] | |
| invalid = [c for c in subset if c not in self.df.columns] | |
| if invalid: | |
| return f"Error: Unknown columns: {invalid}. Available: {list(self.df.columns)}", False | |
| if len(args) > 1: | |
| keep = args[1].lower() | |
| if keep not in ("first", "last", "none", "false"): | |
| return f"Error: keep must be 'first', 'last', or 'none'. Got: '{keep}'", False | |
| if keep == "none": | |
| keep = False | |
| n_before = len(self.df) | |
| self.df = self.df.drop_duplicates(subset=subset, keep=keep).reset_index(drop=True) | |
| n_after = len(self.df) | |
| removed = n_before - n_after | |
| if removed == 0: | |
| return "No duplicate rows found to remove.", False | |
| return f"Removed {removed} duplicate rows. Dataset: {n_before} → {n_after} rows.", True | |
| def _cmd_fix_dtype(self, args: list) -> Tuple[str, bool]: | |
| if len(args) < 2: | |
| return "Error: Usage: fix_dtype COLUMN TYPE (int/float/str/datetime)", False | |
| col = args[0] | |
| target = args[1].lower() | |
| if col not in self.df.columns: | |
| return f"Error: Column '{col}' not found.", False | |
| before_type = str(self.df[col].dtype) | |
| errors = 0 | |
| if target in ("int", "int64"): | |
| self.df[col] = self.df[col].astype(str).str.replace(r'[^\d.\-]', '', regex=True) | |
| numeric = pd.to_numeric(self.df[col], errors="coerce") | |
| errors = numeric.isna().sum() - self.df[col].isna().sum() | |
| self.df[col] = numeric.astype("Int64") | |
| elif target in ("float", "float64"): | |
| self.df[col] = self.df[col].astype(str).str.replace(r'[^\d.\-]', '', regex=True) | |
| self.df[col] = pd.to_numeric(self.df[col], errors="coerce") | |
| errors = self.df[col].isna().sum() | |
| elif target in ("str", "string", "object"): | |
| self.df[col] = self.df[col].astype(str) | |
| elif target in ("datetime", "date"): | |
| self.df[col] = pd.to_datetime(self.df[col], errors="coerce", infer_datetime_format=True) | |
| errors = self.df[col].isna().sum() | |
| else: | |
| return f"Error: Unknown type '{target}'. Use: int, float, str, datetime", False | |
| return f"Converted '{col}' from {before_type} → {target}. Coercion errors: {errors}", True | |
| def _cmd_replace(self, args: list) -> Tuple[str, bool]: | |
| if len(args) < 3: | |
| return "Error: Usage: replace COLUMN OLD_VALUE NEW_VALUE", False | |
| col = args[0] | |
| old_val = args[1] | |
| new_val = args[2] | |
| if col not in self.df.columns: | |
| return f"Error: Column '{col}' not found.", False | |
| mask = self.df[col].astype(str) == old_val | |
| n_matches = mask.sum() | |
| if n_matches == 0: | |
| return f"No matches found for '{old_val}' in column '{col}'.", False | |
| self.df.loc[mask, col] = new_val | |
| return f"Replaced {n_matches} occurrences of '{old_val}' with '{new_val}' in '{col}'.", True | |
| def _cmd_regex_replace(self, args: list) -> Tuple[str, bool]: | |
| """Regex-based replacement within a column.""" | |
| if len(args) < 3: | |
| return "Error: Usage: regex_replace COLUMN PATTERN REPLACEMENT", False | |
| col = args[0] | |
| pattern = args[1] | |
| replacement = args[2] | |
| if col not in self.df.columns: | |
| return f"Error: Column '{col}' not found.", False | |
| try: | |
| before_vals = self.df[col].astype(str).copy() | |
| self.df[col] = self.df[col].astype(str).str.replace(pattern, replacement, regex=True) | |
| n_changed = (before_vals != self.df[col].astype(str)).sum() | |
| except re.error as e: | |
| return f"Error: Invalid regex pattern '{pattern}': {e}", False | |
| if n_changed == 0: | |
| return f"No matches for pattern '{pattern}' in column '{col}'.", False | |
| return f"Regex replaced {n_changed} values in '{col}' (pattern: '{pattern}' → '{replacement}').", True | |
| def _cmd_standardize(self, args: list) -> Tuple[str, bool]: | |
| if len(args) < 2: | |
| return "Error: Usage: standardize COLUMN METHOD (lowercase/uppercase/titlecase/strip)", False | |
| col = args[0] | |
| method = args[1].lower() | |
| if col not in self.df.columns: | |
| return f"Error: Column '{col}' not found.", False | |
| before_uniq = self.df[col].nunique() | |
| if method == "lowercase": | |
| self.df[col] = self.df[col].astype(str).str.lower() | |
| elif method == "uppercase": | |
| self.df[col] = self.df[col].astype(str).str.upper() | |
| elif method == "titlecase": | |
| self.df[col] = self.df[col].astype(str).str.title() | |
| elif method == "strip": | |
| self.df[col] = self.df[col].astype(str).str.strip() | |
| else: | |
| return f"Error: Unknown method '{method}'. Use: lowercase, uppercase, titlecase, strip", False | |
| after_uniq = self.df[col].nunique() | |
| consolidated = before_uniq - after_uniq | |
| return f"Standardized '{col}' using {method}. Unique values: {before_uniq} → {after_uniq} (consolidated {consolidated}).", True | |
| def _cmd_remove_rows(self, args: list) -> Tuple[str, bool]: | |
| if len(args) < 3: | |
| return "Error: Usage: remove_rows COLUMN CONDITION VALUE\n Conditions: equals, not_equals, less_than, greater_than, contains", False | |
| col = args[0] | |
| condition = args[1].lower() | |
| value = args[2] | |
| if col not in self.df.columns: | |
| return f"Error: Column '{col}' not found.", False | |
| n_before = len(self.df) | |
| if condition == "equals": | |
| mask = self.df[col].astype(str) == value | |
| elif condition == "not_equals": | |
| mask = self.df[col].astype(str) != value | |
| elif condition == "less_than": | |
| try: | |
| val = float(value) | |
| mask = pd.to_numeric(self.df[col], errors="coerce") < val | |
| except ValueError: | |
| return f"Error: '{value}' is not a valid number for less_than.", False | |
| elif condition == "greater_than": | |
| try: | |
| val = float(value) | |
| mask = pd.to_numeric(self.df[col], errors="coerce") > val | |
| except ValueError: | |
| return f"Error: '{value}' is not a valid number for greater_than.", False | |
| elif condition == "contains": | |
| mask = self.df[col].astype(str).str.contains(value, case=False, na=False) | |
| else: | |
| return f"Error: Unknown condition '{condition}'. Use: equals, not_equals, less_than, greater_than, contains", False | |
| n_removed = mask.sum() | |
| if n_removed == 0: | |
| return f"No rows match condition '{col} {condition} {value}'.", False | |
| self.df = self.df[~mask].reset_index(drop=True) | |
| return f"Removed {n_removed} rows where {col} {condition} {value}. Dataset: {n_before} → {len(self.df)} rows.", True | |
| def _cmd_clip(self, args: list) -> Tuple[str, bool]: | |
| if len(args) < 3: | |
| return "Error: Usage: clip COLUMN LOWER UPPER", False | |
| col = args[0] | |
| if col not in self.df.columns: | |
| return f"Error: Column '{col}' not found.", False | |
| try: | |
| lower = float(args[1]) | |
| upper = float(args[2]) | |
| except ValueError: | |
| return "Error: LOWER and UPPER must be numbers.", False | |
| numeric_col = pd.to_numeric(self.df[col], errors="coerce") | |
| n_clipped = ((numeric_col < lower) | (numeric_col > upper)).sum() | |
| self.df[col] = numeric_col.clip(lower=lower, upper=upper) | |
| return f"Clipped {n_clipped} values in '{col}' to [{lower}, {upper}].", True | |
| def _cmd_rename_column(self, args: list) -> Tuple[str, bool]: | |
| """Rename a column.""" | |
| if len(args) < 2: | |
| return "Error: Usage: rename_column OLD_NAME NEW_NAME", False | |
| old_name = args[0] | |
| new_name = args[1] | |
| if old_name not in self.df.columns: | |
| return f"Error: Column '{old_name}' not found.", False | |
| if new_name in self.df.columns: | |
| return f"Error: Column '{new_name}' already exists.", False | |
| self.df = self.df.rename(columns={old_name: new_name}) | |
| return f"Renamed column '{old_name}' → '{new_name}'.", True | |
| def _cmd_drop_column(self, args: list) -> Tuple[str, bool]: | |
| """Drop a column from the dataset.""" | |
| if not args: | |
| return "Error: Usage: drop_column COLUMN_NAME", False | |
| col = args[0] | |
| if col not in self.df.columns: | |
| return f"Error: Column '{col}' not found.", False | |
| self.df = self.df.drop(columns=[col]) | |
| return f"Dropped column '{col}'. Remaining columns: {len(self.df.columns)}", True | |
| def _cmd_sort(self, args: list) -> Tuple[str, bool]: | |
| """Sort dataset by a column.""" | |
| if not args: | |
| return "Error: Usage: sort COLUMN [asc|desc]", False | |
| col = args[0] | |
| if col not in self.df.columns: | |
| return f"Error: Column '{col}' not found.", False | |
| ascending = True | |
| if len(args) > 1 and args[1].lower() == "desc": | |
| ascending = False | |
| self.df = self.df.sort_values(by=col, ascending=ascending, na_position="last").reset_index(drop=True) | |
| direction = "ascending" if ascending else "descending" | |
| return f"Sorted dataset by '{col}' ({direction}).", True | |
| # ── Special commands ───────────────────────────────────────────────── | |
| def _cmd_validate(self, args: list) -> Tuple[str, bool]: | |
| return "__VALIDATE__", False | |
| def _cmd_submit(self, args: list) -> Tuple[str, bool]: | |
| return "__SUBMIT__", False | |