Spaces:
Running
Running
| import logging | |
| import pandas as pd | |
| import json | |
| import ast | |
| import traceback | |
| from typing import Dict, Any, Tuple | |
| from core.llm import chat | |
| logger = logging.getLogger(__name__) | |
| def generate_etl_script(df: pd.DataFrame, user_instructions: str = "") -> str: | |
| """ | |
| Generates a Python Pandas script to clean the data. | |
| """ | |
| sample_data = df.head(3).to_csv(index=False) | |
| dtypes = df.dtypes.to_dict() | |
| missing_info = df.isnull().sum().to_dict() | |
| prompt = f"""You are an Expert Data Engineer. | |
| Write a Python script using pandas to clean this dataset. | |
| DATASET INFO: | |
| Columns and Types: {dtypes} | |
| Missing Values: {missing_info} | |
| Sample Data: | |
| {sample_data} | |
| USER INSTRUCTIONS (if any): {user_instructions} | |
| RULES: | |
| 1. The script MUST define a function called `clean_data(df)` that takes a pandas DataFrame and returns a pandas DataFrame. | |
| 2. Fix missing values, convert bad data types, drop obvious duplicates, and remove extreme outliers. | |
| 3. Use ONLY the 'pandas' and 'numpy' libraries. Do NOT use any other libraries (no os, sys, subprocess, etc). | |
| 4. Output ONLY the raw python code. No markdown formatting, no explanation. Just the raw code. | |
| 5. Example format: | |
| import pandas as pd | |
| import numpy as np | |
| def clean_data(df): | |
| df = df.drop_duplicates() | |
| df['Age'] = df['Age'].fillna(df['Age'].mean()) | |
| return df | |
| """ | |
| try: | |
| response = chat(prompt, temperature=0.1, max_tokens=1000) | |
| # Clean up response if it wrapped in markdown | |
| if '```python' in response: | |
| response = response.split('```python')[1].split('```')[0].strip() | |
| elif '```' in response: | |
| response = response.split('```')[1].split('```')[0].strip() | |
| return response | |
| except Exception as e: | |
| logger.error(f"Failed to generate ETL script: {e}") | |
| return "def clean_data(df):\n return df" | |
| def safe_execute_etl(df: pd.DataFrame, code_str: str) -> Tuple[bool, pd.DataFrame, str]: | |
| """ | |
| Safely executes the AI-generated pandas script on the DataFrame. | |
| Returns (success, cleaned_df, error_message). | |
| """ | |
| # 1. Security Check: Parse AST to ensure no dangerous imports/calls | |
| try: | |
| tree = ast.parse(code_str) | |
| for node in ast.walk(tree): | |
| if isinstance(node, ast.Import) or isinstance(node, ast.ImportFrom): | |
| for alias in node.names: | |
| module_name = alias.name.split('.')[0] | |
| if module_name not in ['pandas', 'numpy']: | |
| return False, df, f"Security Violation: Import of '{module_name}' is not allowed. Only pandas and numpy are permitted." | |
| if isinstance(node, ast.Call): | |
| if isinstance(node.func, ast.Name): | |
| if node.func.id in ['exec', 'eval', 'open', '__import__']: | |
| return False, df, f"Security Violation: Function '{node.func.id}' is not allowed." | |
| except SyntaxError as e: | |
| return False, df, f"Syntax Error in AI code: {e}" | |
| # 2. Execute Code safely | |
| local_env = { | |
| 'pd': pd, | |
| 'np': __import__('numpy') | |
| } | |
| try: | |
| # Execute the function definition into our local_env | |
| exec(code_str, {}, local_env) | |
| if 'clean_data' not in local_env: | |
| return False, df, "Function 'clean_data' was not defined by the AI." | |
| # Run the function on a copy of the dataframe | |
| df_copy = df.copy() | |
| cleaned_df = local_env['clean_data'](df_copy) | |
| if not isinstance(cleaned_df, pd.DataFrame): | |
| return False, df, "The 'clean_data' function did not return a DataFrame." | |
| return True, cleaned_df, code_str | |
| except Exception as e: | |
| error_msg = traceback.format_exc() | |
| logger.error(f"ETL Execution failed: {error_msg}") | |
| return False, df, f"Runtime Error: {e}" | |