src/managers/execution/python_executor.py

"""
Python Executor - Persistent Python execution environment.
Similar to a Jupyter kernel: maintains state and tool namespace across calls.
"""

import io
import sys
from typing import Any, Dict

__all__ = ["PythonExecutor"]

# Cap on stdout captured per execution. Large prints (e.g. full dataframes)
# get re-sent to the LLM on every subsequent step, so an uncapped print can
# multiply token cost across the rest of the run.
MAX_OUTPUT_CHARS = 4000


class PythonExecutor:
    """
    Persistent Python execution environment.
    Maintains namespace (variables, functions) across multiple execute() calls.
    """

    def __init__(self):
        self.namespace = {}
        self.reset_environment()

    def reset_environment(self):
        """Reset the execution environment to a clean state."""
        self.namespace = {"__builtins__": __builtins__}

        try:
            import pandas as pd
            import numpy as np
            import os
            from pathlib import Path
            self.namespace.update({"pd": pd, "np": np, "os": os, "Path": Path})

            def write_text_file(filepath, content):
                """Write text content to a file."""
                with open(filepath, 'w') as f:
                    f.write(content)
                return f"File written to {filepath}"

            def write_dataframe_to_csv(df, filepath, **kwargs):
                """Write pandas DataFrame to CSV file."""
                df.to_csv(filepath, **kwargs)
                return f"DataFrame written to {filepath}"

            def write_dataframe_to_tsv(df, filepath, **kwargs):
                """Write pandas DataFrame to TSV file."""
                df.to_csv(filepath, sep='\t', **kwargs)
                return f"DataFrame written to {filepath}"

            def create_directory(dirpath):
                """Create directory if it doesn't exist."""
                Path(dirpath).mkdir(parents=True, exist_ok=True)
                return f"Directory created: {dirpath}"

            self.namespace.update({
                "write_text_file": write_text_file,
                "write_dataframe_to_csv": write_dataframe_to_csv,
                "write_dataframe_to_tsv": write_dataframe_to_tsv,
                "create_directory": create_directory,
            })

        except ImportError as e:
            print(f"Warning: Could not import library: {e}")

    def send_functions(self, functions: Dict[str, Any]):
        """Inject functions into the execution namespace."""
        self.namespace.update(functions)

    def send_variables(self, variables: Dict[str, Any]):
        """Inject variables into the execution namespace."""
        if variables:
            self.namespace.update(variables)

    def __call__(self, code: str) -> Any:
        return self.execute(code)

    def execute(self, code: str) -> Any:
        """Execute Python code in the persistent namespace, returning stdout output."""
        try:
            old_stdout = sys.stdout
            sys.stdout = captured_output = io.StringIO()
            try:
                # Executing LLM-generated code is this agent's core purpose; the
                # accepted risk is bounded by running in an isolated HF Space with
                # no secrets beyond the model API key.
                exec(code, self.namespace)  # nosec B102
                result = captured_output.getvalue().strip()
                if not result:
                    return "Code executed successfully"
                if len(result) > MAX_OUTPUT_CHARS:
                    omitted = len(result) - MAX_OUTPUT_CHARS
                    result = (
                        result[:MAX_OUTPUT_CHARS]
                        + f"\n...[output truncated, {omitted} more characters omitted. "
                          f"Print a smaller summary (e.g. .head(), .describe(), or specific columns) "
                          f"if you need to inspect this further.]"
                    )
                return result
            finally:
                sys.stdout = old_stdout
        except Exception as e:
            return f"Error: {str(e)}"