feat: enhance REPL functionality with improved plot capture and download utilities

1. support_tools.py -- Replaced the standalone _captured_plots
   list, _capture_matplotlib_plots(), and
  _apply_matplotlib_patches() (~90 lines) with a single
  module-level _global_plot_capture = PlotCapture() instance.
  The backward-compat functions
  get_captured_plots()/clear_captured_plots() now delegate to
  it. No more duplicated capture logic between the class and the
   old module-level functions.
  2. execution.py -- Extracted _run_script_in_tempfile() shared
  helper. run_r_code and run_bash_script had identical
  temp-file-write → subprocess → cleanup boilerplate. Also
  consolidated _LANGUAGE_MARKERS dict as single source of truth
  for both detect_code_language() and strip_code_markers().
  3. download.py -- Extracted _download_file() to replace the
  inline download_with_progress() closure that was duplicated
  inside check_and_download_s3_files, and the near-identical
  download code in download_and_unzip. Also improved temp file
  cleanup with finally blocks.
  4. formatting.py -- Replaced 20-line manual word-wrapping with
   textwrap.fill(). Removed a stray print(message) debug
  statement.
  5. tool_bridge.py -- Replaced per-key loop with
  namespace.update(). Removed the dead
  builtins._BaseAgent_custom_functions hack (no callers).
  6. __init__.py -- Stopped re-exporting private helpers
  _parse_docstring/_type_to_string. Tests import them directly
  from schema.py, so no breakage.

  One minor behavioral change (non-breaking):

  run_bash_script no longer passes shell=True to subprocess.run.
   Since the script has a #!/bin/bash shebang and is chmod 755,
  direct execution is actually more correct. The old env/cwd
  explicit args were also redundant since subprocess.run
  inherits them from the parent process.

Author: BinglanLi

BaseAgent/tools/support_tools.py

@@ -1,14 +1,13 @@
 import base64
+import importlib
+import inspect
 import io
 import sys
 from io import StringIO
 
 # Create a persistent namespace that will be shared across all executions
 _persistent_namespace = {}
 
-# Global list to store captured plots (module-level fallback for backward compat)
-_captured_plots = []
-
 
 class PlotCapture:
     """Per-agent matplotlib plot capture with isolated state.
@@ -101,6 +100,11 @@ def clear(self) -> None:
         self._plots = []
 
 
+# Module-level PlotCapture for backward-compat callers that use run_python_repl
+# without supplying a namespace (i.e., direct callers, not BaseAgent instances).
+_global_plot_capture = PlotCapture()
+
+
 def run_python_repl(command: str, namespace: dict | None = None) -> str:
     """
     Executes the provided Python command in a persistent environment and returns the output.
@@ -128,7 +132,7 @@ def execute_in_repl(command: str) -> str:
             # (backward compat). Per-instance callers apply PlotCapture.apply_patches()
             # themselves before invoking run_python_repl.
             if namespace is None:
-                _apply_matplotlib_patches()
+                _global_plot_capture.apply_patches()
 
             exec(command, _ns)
             output = mystdout.getvalue()
@@ -143,109 +147,18 @@ def execute_in_repl(command: str) -> str:
     return execute_in_repl(command)
 
 
-def _capture_matplotlib_plots():
-    """
-    Capture any matplotlib plots that might have been generated during execution.
-    """
-    global _captured_plots
-    try:
-        import matplotlib.pyplot as plt
-
-        # Check if there are any active figures
-        if plt.get_fignums():
-            for fig_num in plt.get_fignums():
-                fig = plt.figure(fig_num)
-
-                # Save figure to base64
-                buffer = io.BytesIO()
-                fig.savefig(buffer, format="png", dpi=150, bbox_inches="tight")
-                buffer.seek(0)
-
-                # Convert to base64
-                image_data = base64.b64encode(buffer.getvalue()).decode("utf-8")
-                plot_data = f"data:image/png;base64,{image_data}"
-
-                # Add to captured plots if not already there
-                if plot_data not in _captured_plots:
-                    _captured_plots.append(plot_data)
-
-                # Close the figure to free memory
-                plt.close(fig)
-
-    except ImportError:
-        # matplotlib not available
-        pass
-    except Exception as e:
-        print(f"Warning: Could not capture matplotlib plots: {e}")
-
-
-def _apply_matplotlib_patches():
-    """
-    Apply simple monkey patches to matplotlib functions to automatically capture plots.
-    """
-    try:
-        import matplotlib.pyplot as plt
-
-        # Only patch if matplotlib is available and not already patched
-        if hasattr(plt, "_base_agent_patched"):
-            return
-
-        # Store original functions
-        original_show = plt.show
-        original_savefig = plt.savefig
-
-        def show_with_capture(*args, **kwargs):
-            """Enhanced show function that captures plots before displaying them."""
-            # Capture any plots before showing
-            _capture_matplotlib_plots()
-            # Print a message to indicate plot was generated
-            print("Plot generated and displayed")
-            # Call the original show function
-            return original_show(*args, **kwargs)
-
-        def savefig_with_capture(*args, **kwargs):
-            """Enhanced savefig function that captures plots after saving them."""
-            # Get the filename from args if provided
-            filename = args[0] if args else kwargs.get("fname", "unknown")
-            # Call the original savefig function
-            result = original_savefig(*args, **kwargs)
-            # Capture the plot after saving
-            _capture_matplotlib_plots()
-            # Print a message to indicate plot was saved
-            print(f"Plot saved to: {filename}")
-            return result
-
-        # Replace functions with enhanced versions
-        plt.show = show_with_capture
-        plt.savefig = savefig_with_capture
-
-        # Mark as patched to avoid double-patching
-        plt._BaseAgent_patched = True
-
-    except ImportError:
-        # matplotlib not available
-        pass
-    except Exception as e:
-        print(f"Warning: Could not apply matplotlib patches: {e}")
-
-
 def get_captured_plots():
-    """
-    Get all captured matplotlib plots.
+    """Get all captured matplotlib plots (backward-compat wrapper).
 
     Returns:
         list: A list of captured matplotlib plots
     """
-    global _captured_plots
-    return _captured_plots.copy()
+    return _global_plot_capture.get_plots()
 
 
 def clear_captured_plots():
-    """
-    Clear all captured matplotlib plots.
-    """
-    global _captured_plots
-    _captured_plots = []
+    """Clear all captured matplotlib plots (backward-compat wrapper)."""
+    _global_plot_capture.clear()
 
 
 def read_function_source_code(function_name: str) -> str:
@@ -258,24 +171,14 @@ def read_function_source_code(function_name: str) -> str:
         str: The source code of the function
 
     """
-    import importlib
-    import inspect
-
     # Split the function name into module path and function name
     parts = function_name.split(".")
     module_path = ".".join(parts[:-1])
     func_name = parts[-1]
 
     try:
-        # Import the module
         module = importlib.import_module(module_path)
-
-        # Get the function object from the module
         function = getattr(module, func_name)
-
-        # Get the source code of the function
-        source_code = inspect.getsource(function)
-
-        return source_code
+        return inspect.getsource(function)
     except (ImportError, AttributeError) as e:
-        return f"Error: Could not find function '{function_name}'. Details: {str(e)}"
+        return f"Error: Could not find function '{function_name}'. Details: {str(e)}"

BaseAgent/utils/__init__.py

@@ -8,8 +8,6 @@
 from BaseAgent.utils.schema import (
     extract_schema_from_function,
     function_to_api_schema,
-    _parse_docstring,
-    _type_to_string,
     enhance_description_with_llm,
     enhance_parameters_with_llm,
 )
@@ -37,8 +35,6 @@
     # schema
     "extract_schema_from_function",
     "function_to_api_schema",
-    "_parse_docstring",
-    "_type_to_string",
     "enhance_description_with_llm",
     "enhance_parameters_with_llm",
     # formatting

BaseAgent/utils/download.py

@@ -9,6 +9,44 @@
 import tqdm
 
 
+def _download_file(url: str, dest_path: str, desc: str) -> bool:
+    """Stream-download *url* to *dest_path* with a tqdm progress bar.
+
+    Args:
+        url: URL to download.
+        dest_path: Local path to write the downloaded bytes to.
+        desc: Label shown in the progress bar.
+
+    Returns:
+        ``True`` on success, ``False`` on any error (error is printed).
+    """
+    try:
+        response = requests.get(url, stream=True)
+        response.raise_for_status()
+        total_size = int(response.headers.get("content-length", 0))
+        with open(dest_path, "wb") as f:
+            with tqdm.tqdm(
+                total=total_size or None,
+                unit="B",
+                unit_scale=True,
+                desc=desc,
+                ncols=80,
+            ) as pbar:
+                for chunk in response.iter_content(chunk_size=8192):
+                    if chunk:
+                        f.write(chunk)
+                        pbar.update(len(chunk))
+        return True
+    except Exception as e:
+        print(f"✗ Failed to download {desc}: {e}")
+        if os.path.exists(dest_path):
+            try:
+                os.remove(dest_path)
+            except OSError:
+                pass
+        return False
+
+
 def check_or_create_path(path=None):
     # Set a default path if none is provided
     if path is None:
@@ -36,35 +74,30 @@ def download_and_unzip(url: str, dest_dir: str) -> str:
         The path to the extracted directory, or an error message.
 
     """
+    tmp_zip_path = None
     try:
         os.makedirs(dest_dir, exist_ok=True)
         print(f"Downloading from {url} ...")
-        with requests.get(url, stream=True) as r:
-            r.raise_for_status()
-            total_size = int(r.headers.get("content-length", 0))
-            chunk_size = 8192
-            with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp_file:
-                with tqdm.tqdm(
-                    total=total_size / (1024**3),
-                    unit="GB",
-                    unit_scale=True,
-                    desc="Downloading",
-                    ncols=80,
-                ) as pbar:
-                    for chunk in r.iter_content(chunk_size=chunk_size):
-                        if chunk:
-                            tmp_file.write(chunk)
-                            pbar.update(len(chunk) / (1024**3))
-                tmp_zip_path = tmp_file.name
+        with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp_file:
+            tmp_zip_path = tmp_file.name
+
+        if not _download_file(url, tmp_zip_path, "Downloading"):
+            return f"Error: Failed to download {url}"
+
         print(f"Downloaded to {tmp_zip_path}. Extracting...")
         with zipfile.ZipFile(tmp_zip_path, "r") as zip_ref:
             zip_ref.extractall(dest_dir)
-        os.unlink(tmp_zip_path)
         print(f"Extraction complete to {dest_dir}")
         return dest_dir
     except Exception as e:
         print(f"Error downloading or extracting zip: {e}")
         return f"Error: {e}"
+    finally:
+        if tmp_zip_path and os.path.exists(tmp_zip_path):
+            try:
+                os.unlink(tmp_zip_path)
+            except OSError:
+                pass
 
 
 def check_and_download_s3_files(
@@ -81,56 +114,19 @@ def check_and_download_s3_files(
     Returns:
         Dictionary mapping file names to download success status
     """
-
     os.makedirs(local_data_lake_path, exist_ok=True)
     download_results = {}
 
-    def download_with_progress(url: str, file_path: str, desc: str) -> bool:
-        """Download file with progress bar."""
-        try:
-            response = requests.get(url, stream=True)
-            response.raise_for_status()
-
-            total_size = int(response.headers.get("content-length", 0))
-
-            with open(file_path, "wb") as f:
-                if total_size > 0:
-                    with tqdm.tqdm(total=total_size, unit="B", unit_scale=True, desc=desc, ncols=80) as pbar:
-                        for chunk in response.iter_content(chunk_size=8192):
-                            if chunk:
-                                f.write(chunk)
-                                pbar.update(len(chunk))
-                else:
-                    for chunk in response.iter_content(chunk_size=8192):
-                        if chunk:
-                            f.write(chunk)
-            return True
-        except Exception as e:
-            print(f"✗ Failed to download {desc}: {e}")
-            if os.path.exists(file_path):
-                try:
-                    os.remove(file_path)
-                except OSError:
-                    pass
-            return False
-
-    def cleanup_file(file_path: str):
-        """Clean up file if it exists."""
-        if os.path.exists(file_path):
-            try:
-                os.remove(file_path)
-            except OSError:
-                pass
-
     # Handle benchmark folder (download as zip)
     if folder == "benchmark":
         print(f"Downloading entire {folder} folder structure...")
         s3_zip_url = urljoin(s3_bucket_url + "/", folder + ".zip")
+        tmp_zip_path = None
+        try:
+            with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp_zip:
+                tmp_zip_path = tmp_zip.name
 
-        with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp_zip:
-            tmp_zip_path = tmp_zip.name
-
-            if download_with_progress(s3_zip_url, tmp_zip_path, f"{folder}.zip"):
+            if _download_file(s3_zip_url, tmp_zip_path, f"{folder}.zip"):
                 print(f"Extracting {folder}.zip...")
                 try:
                     with zipfile.ZipFile(tmp_zip_path, "r") as zip_ref:
@@ -140,10 +136,14 @@ def cleanup_file(file_path: str):
                 except Exception as e:
                     print(f"✗ Error extracting {folder}.zip: {e}")
                     download_results = dict.fromkeys(expected_files, False)
-                finally:
-                    cleanup_file(tmp_zip_path)
             else:
                 download_results = dict.fromkeys(expected_files, False)
+        finally:
+            if tmp_zip_path and os.path.exists(tmp_zip_path):
+                try:
+                    os.unlink(tmp_zip_path)
+                except OSError:
+                    pass
 
         return download_results
 
@@ -158,7 +158,7 @@ def cleanup_file(file_path: str):
         s3_file_url = urljoin(s3_bucket_url + "/" + folder + "/", filename)
         print(f"Downloading {filename} from {folder}...")
 
-        if download_with_progress(s3_file_url, local_file_path, filename):
+        if _download_file(s3_file_url, local_file_path, filename):
             print(f"✓ Successfully downloaded: {filename}")
             download_results[filename] = True
         else: