plot-agent 0.2.2__py3-none-any.whl → 0.3.1__py3-none-any.whl

plot_agent/agent.py

@@ -1,3 +1,7 @@
+"""
+This module contains the PlotAgent class, which is used to generate Plotly code based on a user's plot description.
+"""
+
 import pandas as pd
 from io import StringIO
 from typing import Optional
@@ -14,10 +18,10 @@ from plot_agent.models import (
     DoesFigExistInput,
     ViewGeneratedCodeInput,
 )
-from plot_agent.execution import PlotlyAgentExecutionEnvironment
+from plot_agent.execution import PlotAgentExecutionEnvironment
 
 
-class PlotlyAgent:
+class PlotAgent:
     """
     A class that uses an LLM to generate Plotly code based on a user's plot description.
     """
@@ -32,7 +36,7 @@ class PlotlyAgent:
         handle_parsing_errors: bool = True,
     ):
         """
-        Initialize the PlotlyAgent.
+        Initialize the PlotAgent.
 
         Args:
             model (str): The model to use for the LLM.
@@ -90,7 +94,7 @@ class PlotlyAgent:
         self.sql_query = sql_query
 
         # Initialize execution environment
-        self.execution_env = PlotlyAgentExecutionEnvironment(df)
+        self.execution_env = PlotAgentExecutionEnvironment(df)
 
         # Initialize the agent with tools
         self._initialize_agent()
@@ -123,7 +127,7 @@ class PlotlyAgent:
 
         # Check if the code executed successfully
         if code_execution_success:
-            return f"Code executed successfully! A figure object was created.\n{code_execution_output}"
+            return f"Success: {code_execution_output}"
         else:
             return f"Error: {code_execution_error}\n{code_execution_output}"
 
@@ -160,19 +164,29 @@ class PlotlyAgent:
             Tool.from_function(
                 func=self.execute_plotly_code,
                 name="execute_plotly_code",
-                description="Execute the provided Plotly code and return a result indicating if the code executed successfully and if a figure object was created.",
+                description=(
+                    "Execute the provided Plotly code and return a result indicating "
+                    "if the code executed successfully and if a figure object was created."
+                ),
                 args_schema=GeneratedCodeInput,
             ),
             StructuredTool.from_function(
                 func=self.does_fig_exist,
                 name="does_fig_exist",
-                description="Check if a figure exists and is available for display. This tool takes no arguments and returns a string indicating if a figure is available for display or not.",
+                description=(
+                    "Check if a figure exists and is available for display. "
+                    "This tool takes no arguments and returns a string indicating "
+                    "if a figure is available for display or not."
+                ),
                 args_schema=DoesFigExistInput,
             ),
             StructuredTool.from_function(
                 func=self.view_generated_code,
                 name="view_generated_code",
-                description="View the generated code. This tool takes no arguments and returns the generated code as a string.",
+                description=(
+                    "View the generated code. "
+                    "This tool takes no arguments and returns the generated code as a string."
+                ),
                 args_schema=ViewGeneratedCodeInput,
             ),
         ]

plot_agent/execution.py

@@ -1,90 +1,238 @@
-import sys
-from io import StringIO
+"""
+This module contains the PlotAgentExecutionEnvironment class, which is used to safely execute LLM‑generated plotting code and capture `fig`.
+
+Security features:
+  • Only allow imports from a fixed list of packages (pandas, numpy,
+    matplotlib, plotly, sklearn)
+  • AST scan rejects any import outside that list and any __dunder__ access
+  • Sandbox builtins to include only a minimal safe set + our _safe_import
+  • Enforce a 60 second timeout via signal.alarm
+"""
+import ast
+import builtins
+import signal
 import traceback
+from io import StringIO
+import contextlib
+
 import pandas as pd
-import plotly.express as px
-import plotly.graph_objects as go
 import numpy as np
 import matplotlib.pyplot as plt
+import plotly.express as px
+import plotly.graph_objects as go
 from plotly.subplots import make_subplots
-from typing import Dict, Any
 
 
-class PlotlyAgentExecutionEnvironment:
+def _timeout_handler(signum, frame):
+    raise TimeoutError("Code execution timed out")
+
+
+# List of allowed modules
+_ALLOWED_MODULES = {
+    "pandas",
+    "numpy",
+    "matplotlib",
+    "plotly",
+    "sklearn",
+    "scipy",
+}
+
+
+# Wrap the real __import__ so only our allowlist can get in
+_orig_import = builtins.__import__
+
+
+def _safe_import(name, globals=None, locals=None, fromlist=(), level=0):
+    """
+    Wrap the real __import__ so only our allowlist can get in
+    """
+    root = name.split(".", 1)[0]
+    if root in _ALLOWED_MODULES:
+        return _orig_import(name, globals, locals, fromlist, level)
+    # If the module is not in the allowlist, raise an ImportError
+    raise ImportError(f"Import of module '{name}' is not allowed.")
+
+
+class PlotAgentExecutionEnvironment:
     """
-    Environment to safely execute plotly code and capture the fig object.
+    Environment to safely execute LLM‑generated plotting code and capture `fig`.
 
-    Args:
-        df (pd.DataFrame): The dataframe to use for the execution environment.
+    Security features:
+      • Only allow imports from a fixed list of packages (pandas, numpy,
+        matplotlib, plotly, sklearn)
+      • AST scan rejects any import outside that list and any __dunder__ access
+      • Sandbox builtins to include only a minimal safe set + our _safe_import
+      • Enforce a 60 second timeout via signal.alarm
+      • Capture both stdout & stderr
+      • Purge any old `fig` between runs
     """
 
+    TIMEOUT_SECONDS = 60
+
+    # A lean set of builtins, plus our safe-import hook
+    _SAFE_BUILTINS = {
+        "abs": abs,
+        "all": all,
+        "any": any,
+        "bin": bin,
+        "bool": bool,
+        "chr": chr,
+        "dict": dict,
+        "divmod": divmod,
+        "enumerate": enumerate,
+        "float": float,
+        "int": int,
+        "len": len,
+        "list": list,
+        "map": map,
+        "max": max,
+        "min": min,
+        "next": next,
+        "pow": pow,
+        "print": print,
+        "range": range,
+        "reversed": reversed,
+        "round": round,
+        "set": set,
+        "str": str,
+        "sum": sum,
+        "tuple": tuple,
+        "zip": zip,
+        # basic exceptions so user code can raise/catch
+        "BaseException": BaseException,
+        "Exception": Exception,
+        "ValueError": ValueError,
+        "TypeError": TypeError,
+        "NameError": NameError,
+        "IndexError": IndexError,
+        "KeyError": KeyError,
+        # our import guard
+        "__import__": _safe_import,
+    }
+
     def __init__(self, df: pd.DataFrame):
         """
-        Initialize the execution environment with the given dataframe.
-
-        Args:
-            df (pd.DataFrame): The dataframe to use for the execution environment.
+        Initialize the execution environment with a dataframe.
         """
         self.df = df
-        self.locals_dict = {
+        # Base namespace for both globals & locals
+        self._base_ns = {
+            "__builtins__": self._SAFE_BUILTINS,
             "df": df,
-            "px": px,
-            "go": go,
             "pd": pd,
             "np": np,
             "plt": plt,
+            "px": px,
+            "go": go,
             "make_subplots": make_subplots,
         }
-        self.output = None
-        self.error = None
         self.fig = None
 
-    def execute_code(self, generated_code: str) -> Dict[str, Any]:
+    def _validate_ast(self, node: ast.AST):
         """
-        Execute the provided code and capture the fig object if created.
+        Walk the AST and enforce:
+         • any Import/ImportFrom must be from _ALLOWED_MODULES
+         • no __dunder__ attribute access
+        """
+        # Walk the AST and enforce:
+        for child in ast.walk(node):
+            # Check for imports
+            if isinstance(child, ast.Import):
+                # Check for imports
+                for alias in child.names:
+                    root = alias.name.split(".", 1)[0]
+                    # Check if the module is in the allowlist
+                    if root not in _ALLOWED_MODULES:
+                        raise ValueError(f"Import of '{alias.name}' is not allowed.")
+            # Check for import-froms
+            elif isinstance(child, ast.ImportFrom):
+                root = (child.module or "").split(".", 1)[0]
+                if root not in _ALLOWED_MODULES:
+                    raise ValueError(f"Import-from of '{child.module}' is not allowed.")
+            # Check for dunder attribute access
+            elif isinstance(child, ast.Attribute) and child.attr.startswith("__"):
+                raise ValueError("Access to dunder attributes is forbidden.")
 
-        Args:
-            generated_code (str): The code to execute.
+    def execute_code(self, generated_code: str):
+        """
+        Execute the user code in a locked‑down sandbox.
 
-        Returns:
-            Dict[str, Any]: A dictionary containing the fig object, output, error, and success status.
+        Returns a dict with:
+          - fig: The figure if created, else None
+          - output: Captured stdout
+          - error: Captured stderr or exception text
+          - success: True if fig was produced and no errors
         """
-        self.output = None
-        self.error = None
 
-        # Capture stdout
-        old_stdout = sys.stdout
-        sys.stdout = mystdout = StringIO()
+        # Copy the base namespace
+        ns = self._base_ns.copy()
+        # Purge any old `fig`
+        ns.pop("fig", None)
 
         try:
-            # Execute the code
-            exec(generated_code, globals(), self.locals_dict)
+            # Parse the generated code
+            tree = ast.parse(generated_code)
+            # Validate the AST
+            self._validate_ast(tree)
+        except Exception as e:
+            # If the code is rejected on safety grounds, return an error
+            return {
+                "fig": None,
+                "output": "",
+                "error": f"Code rejected on safety grounds: {e}",
+                "success": False,
+            }
 
-            # Check if a fig object was created
-            if "fig" in self.locals_dict:
-                self.fig = self.locals_dict["fig"]
-                self.output = "Code executed successfully. 'fig' object was created."
-            else:
-                print(f"no fig object created: {generated_code}")
-                self.error = "Code executed without errors, but no 'fig' object was created. Make sure your code creates a variable named 'fig'."
+        # Set a timeout
+        signal.signal(signal.SIGALRM, _timeout_handler)
+        signal.alarm(self.TIMEOUT_SECONDS)
 
+        # Execute the code
+        out_buf, err_buf = StringIO(), StringIO()
+        try:
+            # Redirect stdout and stderr
+            with contextlib.redirect_stdout(out_buf), contextlib.redirect_stderr(
+                err_buf
+            ):
+                # Execute the code
+                exec(generated_code, ns, ns)
+        except TimeoutError as te:
+            # If the code execution timed out, return an error
+            tb = traceback.format_exc()
+            return {
+                "fig": None,
+                "output": out_buf.getvalue(),
+                "error": f"Code execution timed out: {te}\n{tb}",
+                "success": False,
+            }
         except Exception as e:
-            self.error = f"Error executing code: {str(e)}\n{traceback.format_exc()}"
-
+            # If there was an error, return an error
+            tb = traceback.format_exc()
+            return {
+                "fig": None,
+                "output": out_buf.getvalue(),
+                "error": f"Error executing code: {e}\n{tb}",
+                "success": False,
+            }
         finally:
-            # Restore stdout
-            sys.stdout = old_stdout
-            captured_output = mystdout.getvalue()
+            # Reset the timeout
+            signal.alarm(0)
 
-            if captured_output.strip():
-                if self.output:
-                    self.output += f"\nOutput:\n{captured_output}"
-                else:
-                    self.output = f"Output:\n{captured_output}"
+        # Get the `fig`
+        fig = ns.get("fig")
+        self.fig = fig
+        if fig is None:
+            return {
+                "fig": None,
+                "output": out_buf.getvalue(),
+                "error": "No `fig` created. Assign your figure to a variable named `fig`.",
+                "success": False,
+            }
 
+        # Return the result
         return {
-            "fig": self.fig,
-            "output": self.output,
-            "error": self.error,
-            "success": self.error is None and self.fig is not None,
+            "fig": fig,
+            "output": "Code executed successfully. 'fig' object was created.",
+            "error": "",
+            "success": True,
         }

plot_agent/models.py

@@ -1,14 +1,21 @@
+"""
+This module contains the models for the PlotAgent.
+"""
+
 from pydantic import BaseModel, Field
 
 
-# Define input schemas for the tools
 class PlotDescriptionInput(BaseModel):
+    """Model indicating that the plot_description function takes a plot_description argument."""
+
     plot_description: str = Field(
         ..., description="Description of the plot the user wants to create"
     )
 
 
 class GeneratedCodeInput(BaseModel):
+    """Model indicating that the generated_code function takes a generated_code argument."""
+
     generated_code: str = Field(
         ..., description="Python code that creates a Plotly figure"
     )
@@ -23,4 +30,4 @@ class DoesFigExistInput(BaseModel):
 class ViewGeneratedCodeInput(BaseModel):
     """Model indicating that the view_generated_code function takes no arguments."""
 
-    pass
+    pass

plot_agent/prompt.py

@@ -1,3 +1,7 @@
+"""
+This module contains the prompts for the PlotAgent.
+"""
+
 DEFAULT_SYSTEM_PROMPT = """
 You are an expert data visualization assistant that helps users create Plotly visualizations in Python.
 Your job is to generate Python and Plotly code based on the user's request that will create the desired visualization

{plot_agent-0.2.2.dist-info → plot_agent-0.3.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: plot-agent
-Version: 0.2.2
+Version: 0.3.1
 Summary: An AI-powered data visualization assistant using Plotly
 Author-email: andrewm4894 <andrewm4894@gmail.com>
 Project-URL: Homepage, https://github.com/andrewm4894/plot-agent
@@ -35,7 +35,7 @@ Here's a simple minimal example of how to use Plot Agent:
 
 ```python
 import pandas as pd
-from plot_agent.agent import PlotlyAgent
+from plot_agent.agent import PlotAgent
 
 # ensure OPENAI_API_KEY is set and available for langchain
 
@@ -46,7 +46,7 @@ df = pd.DataFrame({
 })
 
 # Initialize the agent
-agent = PlotlyAgent()
+agent = PlotAgent()
 
 # Set the dataframe
 agent.set_df(df)

plot_agent-0.3.1.dist-info/RECORD

@@ -0,0 +1,10 @@
+plot_agent/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+plot_agent/agent.py,sha256=sIG8GMS2A8TP_3kRxbgefn-yZM4K_7niZQR6cJrhl4s,9872
+plot_agent/execution.py,sha256=lQNyPzphPIdMQXxQkaf_g6oDZsU3dgF0or0ysKJm6FM,7537
+plot_agent/models.py,sha256=THdGGGfGmRZ5rtgXvjPcQxFRRTZVFoADEHI_lsMVha8,860
+plot_agent/prompt.py,sha256=5hBlF7jdMrj6MiGEL7YmSDWFUfiCXyIZfZtf3NstKoo,3125
+plot_agent-0.3.1.dist-info/licenses/LICENSE,sha256=A4DPih7wHrh4VMEG3p1PhorqdhjmGIo8nQdYNQL7daA,1062
+plot_agent-0.3.1.dist-info/METADATA,sha256=zkpeWRWczA_CzH7mahNtEuIvumBOhXaNGTiAcUIOQZQ,2837
+plot_agent-0.3.1.dist-info/WHEEL,sha256=lTU6B6eIfYoiQJTZNc-fyaR6BpL6ehTzU3xGYxn2n8k,91
+plot_agent-0.3.1.dist-info/top_level.txt,sha256=KyOjpihUssx26Ra-37vKUQ71pI2qgJsHaRwXHJUhjzQ,11
+plot_agent-0.3.1.dist-info/RECORD,,

plot_agent-0.2.2.dist-info/RECORD

@@ -1,10 +0,0 @@
-plot_agent/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-plot_agent/agent.py,sha256=mflUj-vE_x9_W7XTJ3-GgYFfvg3uXV3wkrUw7wRPlVM,9592
-plot_agent/execution.py,sha256=BBKBVGQDrg7BPWvPbLWOjkAAFRlfyu28QxClXuspd8o,2772
-plot_agent/models.py,sha256=ZOWWeYaqmnKJarXYyXnBQQ4nwUe71ae_gMul3bZXaWU,644
-plot_agent/prompt.py,sha256=HjRgbsAe8HHs8arQogvzOGQdThEWKRqQhtQyaUplxhQ,3064
-plot_agent-0.2.2.dist-info/licenses/LICENSE,sha256=A4DPih7wHrh4VMEG3p1PhorqdhjmGIo8nQdYNQL7daA,1062
-plot_agent-0.2.2.dist-info/METADATA,sha256=dLTSBMjvxg0U63r-EbN2tfI7-eZGYexq854L0kN8M6M,2841
-plot_agent-0.2.2.dist-info/WHEEL,sha256=lTU6B6eIfYoiQJTZNc-fyaR6BpL6ehTzU3xGYxn2n8k,91
-plot_agent-0.2.2.dist-info/top_level.txt,sha256=KyOjpihUssx26Ra-37vKUQ71pI2qgJsHaRwXHJUhjzQ,11
-plot_agent-0.2.2.dist-info/RECORD,,