assathe/data-analyst-agent2

1# code_executor.py2import sys3import traceback4import io5from contextlib import redirect_stdout, redirect_stderr6import subprocess7import os8 9def install_required_packages():10    """Install required packages if not available."""11    required_packages = [12        'pandas', 'numpy', 'matplotlib', 'seaborn', 'plotly', 13        'requests', 'beautifulsoup4', 'duckdb', 'lxml', 'html5lib'14    ]15    16    for package in required_packages:17        try:18            __import__(package.replace('-', '_'))19        except ImportError:20            try:21                subprocess.check_call([sys.executable, '-m', 'pip', 'install', package])22            except subprocess.CalledProcessError:23                print(f"Warning: Could not install {package}")24 25def execute_user_code(code: str, globals_dict: dict = None) -> dict:26    """27    Execute user-provided Python code safely.28    29    Args:30        code: Python code to execute31        globals_dict: Global variables to make available to the code32    33    Returns:34        dict: Execution result with 'success', 'output', 'error', and 'globals' keys35    """36    37    if globals_dict is None:38        globals_dict = {}39    40    # Add common imports and utilities41    exec_globals = {42        '__builtins__': __builtins__,43        'print': print,44        **globals_dict45    }46    47    # Capture stdout and stderr48    stdout_capture = io.StringIO()49    stderr_capture = io.StringIO()50    51    try:52        # Install packages if needed (for local development)53        try:54            install_required_packages()55        except:56            pass  # Continue even if package installation fails57        58        with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture):59            # Execute the code60            exec(code, exec_globals)61        62        return {63            'success': True,64            'output': stdout_capture.getvalue(),65            'error': stderr_capture.getvalue() if stderr_capture.getvalue() else None,66            'globals': exec_globals67        }68        69    except Exception as e:70        error_traceback = traceback.format_exc()71        stderr_output = stderr_capture.getvalue()72        73        full_error = f"Execution Error: {str(e)}\n"74        if stderr_output:75            full_error += f"Stderr: {stderr_output}\n"76        full_error += f"Traceback:\n{error_traceback}"77        78        return {79            'success': False,80            'output': stdout_capture.getvalue(),81            'error': full_error,82            'globals': exec_globals83        }84