Modules/Code_Interpreter.py

from __future__ import annotations

import os
import sys
from io import StringIO
from typing import Annotated

import gradio as gr
from ._docstrings import autodoc
from .File_System import ROOT_DIR

from app import _log_call_end, _log_call_start, _truncate_for_log

# Single source of truth for the LLM-facing tool description
TOOL_SUMMARY = (
    "Execute Python code from the tool root; returns captured stdout or the exception text."
)


@autodoc(
    summary=TOOL_SUMMARY,
)
def Code_Interpreter(code: Annotated[str, "Python source code to run; stdout is captured and returned."]) -> str:
    _log_call_start("Code_Interpreter", code=_truncate_for_log(code or "", 300))
    if code is None:
        result = "No code provided."
        _log_call_end("Code_Interpreter", result)
        return result
    
    # Pre-import commonly used libraries to make them available
    pre_imports = """
try:
    import pandas as pd
    import numpy as np
    import matplotlib.pyplot as plt
    import seaborn as sns
    import io
    import requests
    import json
    from datetime import datetime, timedelta
    import re
except ImportError as e:
    print(f"Warning: Could not import some libraries: {e}")

"""
    
    # Combine pre-imports with user code
    full_code = pre_imports + "\n" + code
    
    old_stdout = sys.stdout
    old_cwd = os.getcwd()
    redirected_output = sys.stdout = StringIO()
    try:
        os.chdir(ROOT_DIR)
        exec(full_code, globals())
        result = redirected_output.getvalue()
        if not result:
            result = "Code executed successfully (no output)"
    except Exception as exc:  # pylint: disable=broad-except
        import traceback
        result = f"Error: {exc}\n\nFull traceback:\n{traceback.format_exc()}"
    finally:
        sys.stdout = old_stdout
        try:
            os.chdir(old_cwd)
        except Exception:
            pass
    _log_call_end("Code_Interpreter", _truncate_for_log(result))
    return result


def build_interface() -> gr.Interface:
    return gr.Interface(
        fn=Code_Interpreter,
        inputs=gr.Code(label="Python Code", language="python"),
        outputs=gr.Textbox(label="Output", lines=5, max_lines=20),
        title="Code Interpreter",
        description="<div style=\"text-align:center\">Execute Python code and see the output.</div>",
        api_description=TOOL_SUMMARY,
        flagging_mode="never",
    )


__all__ = ["Code_Interpreter", "build_interface"]