plot-agent 0.2.2__tar.gz → 0.3.1__tar.gz

{plot_agent-0.2.2 → plot_agent-0.3.1}/PKG-INFO

@@ -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.2.2 → plot_agent-0.3.1}/README.md

@@ -21,7 +21,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
 
@@ -32,7 +32,7 @@ df = pd.DataFrame({
 })
 
 # Initialize the agent
-agent = PlotlyAgent()
+agent = PlotAgent()
 
 # Set the dataframe
 agent.set_df(df)

{plot_agent-0.2.2 → plot_agent-0.3.1}/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-0.3.1/plot_agent/execution.py

@@ -0,0 +1,238 @@
+"""
+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 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
+
+
+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 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
+      • 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 a dataframe.
+        """
+        self.df = df
+        # Base namespace for both globals & locals
+        self._base_ns = {
+            "__builtins__": self._SAFE_BUILTINS,
+            "df": df,
+            "pd": pd,
+            "np": np,
+            "plt": plt,
+            "px": px,
+            "go": go,
+            "make_subplots": make_subplots,
+        }
+        self.fig = None
+
+    def _validate_ast(self, node: ast.AST):
+        """
+        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.")
+
+    def execute_code(self, generated_code: str):
+        """
+        Execute the user code in a locked‑down sandbox.
+
+        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
+        """
+
+        # Copy the base namespace
+        ns = self._base_ns.copy()
+        # Purge any old `fig`
+        ns.pop("fig", None)
+
+        try:
+            # 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,
+            }
+
+        # 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:
+            # 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:
+            # Reset the timeout
+            signal.alarm(0)
+
+        # 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": fig,
+            "output": "Code executed successfully. 'fig' object was created.",
+            "error": "",
+            "success": True,
+        }

{plot_agent-0.2.2 → plot_agent-0.3.1}/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-0.2.2 → plot_agent-0.3.1}/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 → plot_agent-0.3.1}/plot_agent.egg-info/PKG-INFO

@@ -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.2.2 → plot_agent-0.3.1}/plot_agent.egg-info/SOURCES.txt

@@ -9,5 +9,4 @@ plot_agent/prompt.py
 plot_agent.egg-info/PKG-INFO
 plot_agent.egg-info/SOURCES.txt
 plot_agent.egg-info/dependency_links.txt
-plot_agent.egg-info/top_level.txt
-tests/test_plot_agent.py
+plot_agent.egg-info/top_level.txt

{plot_agent-0.2.2 → plot_agent-0.3.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "plot-agent"
-version = "0.2.2"
+version = "0.3.1"
 authors = [
     { name="andrewm4894", email="andrewm4894@gmail.com" },
 ]

plot_agent-0.2.2/plot_agent/execution.py

@@ -1,90 +0,0 @@
-import sys
-from io import StringIO
-import traceback
-import pandas as pd
-import plotly.express as px
-import plotly.graph_objects as go
-import numpy as np
-import matplotlib.pyplot as plt
-from plotly.subplots import make_subplots
-from typing import Dict, Any
-
-
-class PlotlyAgentExecutionEnvironment:
-    """
-    Environment to safely execute plotly code and capture the fig object.
-
-    Args:
-        df (pd.DataFrame): The dataframe to use for the execution environment.
-    """
-
-    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.
-        """
-        self.df = df
-        self.locals_dict = {
-            "df": df,
-            "px": px,
-            "go": go,
-            "pd": pd,
-            "np": np,
-            "plt": plt,
-            "make_subplots": make_subplots,
-        }
-        self.output = None
-        self.error = None
-        self.fig = None
-
-    def execute_code(self, generated_code: str) -> Dict[str, Any]:
-        """
-        Execute the provided code and capture the fig object if created.
-
-        Args:
-            generated_code (str): The code to execute.
-
-        Returns:
-            Dict[str, Any]: A dictionary containing the fig object, output, error, and success status.
-        """
-        self.output = None
-        self.error = None
-
-        # Capture stdout
-        old_stdout = sys.stdout
-        sys.stdout = mystdout = StringIO()
-
-        try:
-            # Execute the code
-            exec(generated_code, globals(), self.locals_dict)
-
-            # 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'."
-
-        except Exception as e:
-            self.error = f"Error executing code: {str(e)}\n{traceback.format_exc()}"
-
-        finally:
-            # Restore stdout
-            sys.stdout = old_stdout
-            captured_output = mystdout.getvalue()
-
-            if captured_output.strip():
-                if self.output:
-                    self.output += f"\nOutput:\n{captured_output}"
-                else:
-                    self.output = f"Output:\n{captured_output}"
-
-        return {
-            "fig": self.fig,
-            "output": self.output,
-            "error": self.error,
-            "success": self.error is None and self.fig is not None,
-        }

plot_agent-0.2.2/tests/test_plot_agent.py

@@ -1,471 +0,0 @@
-import pytest
-import pandas as pd
-import numpy as np
-from plot_agent.agent import PlotlyAgent
-from langchain_core.messages import HumanMessage, AIMessage
-
-
-def test_plotly_agent_initialization():
-    """Test that PlotlyAgent initializes correctly."""
-    agent = PlotlyAgent()
-    assert agent.llm is not None
-    assert agent.df is None
-    assert agent.df_info is None
-    assert agent.df_head is None
-    assert agent.sql_query is None
-    assert agent.execution_env is None
-    assert agent.chat_history == []
-    assert agent.agent_executor is None
-    assert agent.generated_code is None
-
-
-def test_set_df():
-    """Test that set_df properly sets up the dataframe and environment."""
-    # Create a sample dataframe
-    df = pd.DataFrame({"x": [1, 2, 3, 4, 5], "y": [10, 20, 30, 40, 50]})
-
-    agent = PlotlyAgent()
-    agent.set_df(df)
-
-    assert agent.df is not None
-    assert agent.df_info is not None
-    assert agent.df_head is not None
-    assert agent.execution_env is not None
-    assert agent.agent_executor is not None
-
-
-def test_execute_plotly_code():
-    """Test that execute_plotly_code works with valid code."""
-    df = pd.DataFrame({"x": [1, 2, 3, 4, 5], "y": [10, 20, 30, 40, 50]})
-
-    agent = PlotlyAgent()
-    agent.set_df(df)
-
-    # Test with valid plotly code
-    valid_code = """import plotly.express as px
-fig = px.scatter(df, x='x', y='y')"""
-
-    result = agent.execute_plotly_code(valid_code)
-    assert "Code executed successfully" in result
-    assert agent.execution_env.fig is not None
-
-
-def test_execute_plotly_code_with_error():
-    """Test that execute_plotly_code handles errors properly."""
-    df = pd.DataFrame({"x": [1, 2, 3, 4, 5], "y": [10, 20, 30, 40, 50]})
-
-    agent = PlotlyAgent()
-    agent.set_df(df)
-
-    # Test with invalid code
-    invalid_code = """import plotly.express as px
-fig = px.scatter(df, x='non_existent_column', y='y')"""
-
-    result = agent.execute_plotly_code(invalid_code)
-    assert "Error" in result
-    assert agent.execution_env.fig is None
-
-
-def test_does_fig_exist():
-    """Test that does_fig_exist correctly reports figure existence."""
-    df = pd.DataFrame({"x": [1, 2, 3, 4, 5], "y": [10, 20, 30, 40, 50]})
-
-    agent = PlotlyAgent()
-    agent.set_df(df)
-
-    # Initially no figure should exist
-    assert "No figure has been created yet" in agent.does_fig_exist()
-
-    # Create a figure
-    valid_code = """import plotly.express as px
-fig = px.scatter(df, x='x', y='y')"""
-    agent.execute_plotly_code(valid_code)
-
-    # Now a figure should exist
-    assert "A figure is available for display" in agent.does_fig_exist()
-
-
-def test_reset_conversation():
-    """Test that reset_conversation clears the chat history."""
-    agent = PlotlyAgent()
-    agent.chat_history = ["message1", "message2"]
-    agent.reset_conversation()
-    assert agent.chat_history == []
-
-
-def test_view_generated_code():
-    """Test that view_generated_code returns the last generated code."""
-    agent = PlotlyAgent()
-    test_code = "test code"
-    agent.generated_code = test_code
-    assert agent.view_generated_code() == test_code
-
-
-def test_get_figure():
-    """Test that get_figure returns the current figure if it exists."""
-    df = pd.DataFrame({"x": [1, 2, 3, 4, 5], "y": [10, 20, 30, 40, 50]})
-
-    agent = PlotlyAgent()
-    agent.set_df(df)
-
-    # Initially no figure should exist
-    assert agent.get_figure() is None
-
-    # Create a figure
-    valid_code = """import plotly.express as px
-fig = px.scatter(df, x='x', y='y')"""
-    agent.execute_plotly_code(valid_code)
-
-    # Now a figure should exist
-    assert agent.get_figure() is not None
-
-
-def test_process_message():
-    """Test that process_message updates chat history and handles responses."""
-    df = pd.DataFrame({"x": [1, 2, 3, 4, 5], "y": [10, 20, 30, 40, 50]})
-
-    agent = PlotlyAgent()
-    agent.set_df(df, sql_query="SELECT x, y FROM df")
-
-    # Test processing a message
-    response = agent.process_message("Create a scatter plot")
-
-    # Check that chat history was updated
-    assert len(agent.chat_history) == 2  # One human message and one AI message
-    assert isinstance(agent.chat_history[0], HumanMessage)
-    assert isinstance(agent.chat_history[1], AIMessage)
-    assert agent.chat_history[0].content == "Create a scatter plot"
-
-
-def test_execute_plotly_code_without_df():
-    """Test that execute_plotly_code handles the case when no dataframe is set."""
-    agent = PlotlyAgent()
-    result = agent.execute_plotly_code("some code")
-    assert "Error" in result and "No dataframe has been set" in result
-
-
-def test_set_df_with_sql_query():
-    """Test that set_df properly handles SQL query context."""
-    df = pd.DataFrame({"x": [1, 2, 3, 4, 5], "y": [10, 20, 30, 40, 50]})
-
-    sql_query = "SELECT x, y FROM table"
-    agent = PlotlyAgent()
-    agent.set_df(df, sql_query=sql_query)
-
-    assert agent.sql_query == sql_query
-
-
-def test_agent_initialization_with_custom_prompt():
-    """Test agent initialization with custom system prompt."""
-    custom_prompt = "Custom system prompt for testing"
-    agent = PlotlyAgent(system_prompt=custom_prompt)
-    assert agent.system_prompt == custom_prompt
-
-
-def test_agent_initialization_with_different_model():
-    """Test agent initialization with different model names."""
-    agent = PlotlyAgent(model="gpt-3.5-turbo")
-    assert agent.llm.model_name == "gpt-3.5-turbo"
-
-
-def test_agent_initialization_with_verbose():
-    """Test agent initialization with verbose settings."""
-    agent = PlotlyAgent(verbose=False)
-    assert agent.verbose == False
-    assert agent.agent_executor is None  # Agent executor not initialized yet
-
-
-def test_agent_initialization_with_max_iterations():
-    """Test agent initialization with different max iterations."""
-    agent = PlotlyAgent(max_iterations=5)
-    assert agent.max_iterations == 5
-
-
-def test_agent_initialization_with_early_stopping():
-    """Test agent initialization with different early stopping methods."""
-    agent = PlotlyAgent(early_stopping_method="generate")
-    assert agent.early_stopping_method == "generate"
-
-
-def test_process_empty_message():
-    """Test processing of empty messages."""
-    df = pd.DataFrame({"x": [1, 2, 3, 4, 5], "y": [10, 20, 30, 40, 50]})
-
-    agent = PlotlyAgent()
-    agent.set_df(df)
-
-    response = agent.process_message("")
-    assert len(agent.chat_history) == 2  # Should still create chat history entries
-    assert isinstance(agent.chat_history[0], HumanMessage)
-    assert isinstance(agent.chat_history[1], AIMessage)
-
-
-def test_process_message_with_code_blocks():
-    """Test processing messages that contain code blocks."""
-    df = pd.DataFrame({"x": [1, 2, 3, 4, 5], "y": [10, 20, 30, 40, 50]})
-
-    agent = PlotlyAgent()
-    agent.set_df(df)
-
-    message = "Here's some code:\n```python\nprint('test')\n```"
-    response = agent.process_message(message)
-    assert len(agent.chat_history) == 2
-    assert "```python" in agent.chat_history[0].content
-
-
-def test_execution_environment_with_different_plot_types():
-    """Test execution environment with different types of plots."""
-    df = pd.DataFrame(
-        {
-            "x": [1, 2, 3, 4, 5],
-            "y": [10, 20, 30, 40, 50],
-            "category": ["A", "B", "A", "B", "A"],
-        }
-    )
-
-    agent = PlotlyAgent()
-    agent.set_df(df)
-
-    # Test scatter plot
-    scatter_code = """import plotly.express as px
-fig = px.scatter(df, x='x', y='y')"""
-    result = agent.execute_plotly_code(scatter_code)
-    assert "Code executed successfully" in result
-    assert agent.execution_env.fig is not None
-
-    # Test bar plot
-    bar_code = """import plotly.express as px
-fig = px.bar(df, x='category', y='y')"""
-    result = agent.execute_plotly_code(bar_code)
-    assert "Code executed successfully" in result
-    assert agent.execution_env.fig is not None
-
-    # Test line plot
-    line_code = """import plotly.express as px
-fig = px.line(df, x='x', y='y')"""
-    result = agent.execute_plotly_code(line_code)
-    assert "Code executed successfully" in result
-    assert agent.execution_env.fig is not None
-
-
-def test_execution_environment_with_subplots():
-    """Test execution environment with subplots."""
-    df = pd.DataFrame(
-        {"x": [1, 2, 3, 4, 5], "y1": [10, 20, 30, 40, 50], "y2": [50, 40, 30, 20, 10]}
-    )
-
-    agent = PlotlyAgent()
-    agent.set_df(df)
-
-    subplot_code = """import plotly.subplots as sp
-import plotly.graph_objects as go
-fig = sp.make_subplots(rows=1, cols=2)
-fig.add_trace(go.Scatter(x=df['x'], y=df['y1']), row=1, col=1)
-fig.add_trace(go.Scatter(x=df['x'], y=df['y2']), row=1, col=2)"""
-
-    result = agent.execute_plotly_code(subplot_code)
-    assert "Code executed successfully" in result
-    assert agent.execution_env.fig is not None
-
-
-def test_execution_environment_with_data_preprocessing():
-    """Test execution environment with data preprocessing steps."""
-    df = pd.DataFrame({"x": [1, 2, 3, 4, 5], "y": [10, 20, 30, 40, 50]})
-
-    agent = PlotlyAgent()
-    agent.set_df(df)
-
-    preprocessing_code = """import plotly.express as px
-# Preprocessing steps
-df['y_normalized'] = (df['y'] - df['y'].min()) / (df['y'].max() - df['y'].min())
-fig = px.scatter(df, x='x', y='y_normalized')"""
-
-    result = agent.execute_plotly_code(preprocessing_code)
-    assert "Code executed successfully" in result
-    assert agent.execution_env.fig is not None
-
-
-def test_handle_syntax_error():
-    """Test handling of syntax errors in generated code."""
-    df = pd.DataFrame({"x": [1, 2, 3, 4, 5], "y": [10, 20, 30, 40, 50]})
-
-    agent = PlotlyAgent()
-    agent.set_df(df)
-
-    invalid_code = """import plotly.express as px
-fig = px.scatter(df, x='x', y='y'  # Missing closing parenthesis"""
-
-    result = agent.execute_plotly_code(invalid_code)
-    assert "Error" in result
-    assert "SyntaxError" in result
-    assert agent.execution_env.fig is None
-
-
-def test_handle_runtime_error():
-    """Test handling of runtime errors in generated code."""
-    df = pd.DataFrame({"x": [1, 2, 3, 4, 5], "y": [10, 20, 30, 40, 50]})
-
-    agent = PlotlyAgent()
-    agent.set_df(df)
-
-    error_code = """import plotly.express as px
-fig = px.scatter(df, x='x', y='y', color='non_existent_column')"""
-
-    result = agent.execute_plotly_code(error_code)
-    assert "Error" in result
-    assert agent.execution_env.fig is None
-
-
-def test_tool_interaction():
-    """Test interaction between different tools."""
-    df = pd.DataFrame({"x": [1, 2, 3, 4, 5], "y": [10, 20, 30, 40, 50]})
-
-    agent = PlotlyAgent()
-    agent.set_df(df)
-
-    # First check if figure exists (should not)
-    assert "No figure has been created yet" in agent.does_fig_exist()
-
-    # Generate and execute code
-    code = """import plotly.express as px
-fig = px.scatter(df, x='x', y='y')"""
-    result = agent.execute_plotly_code(code)
-    assert "Code executed successfully" in result
-
-    # Check if figure exists (should now exist)
-    assert "A figure is available for display" in agent.does_fig_exist()
-
-    # View the generated code
-    assert code in agent.view_generated_code()
-
-
-def test_tool_validation():
-    """Test validation of tool inputs."""
-    df = pd.DataFrame({"x": [1, 2, 3, 4, 5], "y": [10, 20, 30, 40, 50]})
-
-    agent = PlotlyAgent()
-    agent.set_df(df)
-
-    # Test with invalid code (empty string)
-    result = agent.execute_plotly_code("")
-    assert "Error" in result
-
-    # Test with invalid code (None)
-    with pytest.raises(AssertionError):
-        agent.execute_plotly_code(None)
-
-
-def test_tool_response_formatting():
-    """Test formatting of tool responses."""
-    df = pd.DataFrame({"x": [1, 2, 3, 4, 5], "y": [10, 20, 30, 40, 50]})
-
-    agent = PlotlyAgent()
-    agent.set_df(df)
-
-    # Test execute_plotly_code response format
-    code = """import plotly.express as px
-fig = px.scatter(df, x='x', y='y')"""
-    result = agent.execute_plotly_code(code)
-    assert isinstance(result, str)
-    assert "Code executed successfully" in result
-
-    # Test does_fig_exist response format
-    result = agent.does_fig_exist()
-    assert isinstance(result, str)
-    assert "figure" in result.lower()
-
-    # Test view_generated_code response format
-    result = agent.view_generated_code()
-    assert isinstance(result, str)
-    assert code in result
-
-
-def test_memory_cleanup():
-    """Test memory cleanup after multiple plot generations."""
-    df = pd.DataFrame({"x": [1, 2, 3, 4, 5], "y": [10, 20, 30, 40, 50]})
-
-    agent = PlotlyAgent()
-    agent.set_df(df)
-
-    # Generate multiple plots
-    for i in range(5):
-        code = f"""import plotly.express as px
-fig = px.scatter(df, x='x', y='y', title='Plot {i}')"""
-        result = agent.execute_plotly_code(code)
-        assert "Code executed successfully" in result
-        assert agent.execution_env.fig is not None
-
-    # Reset conversation and check memory
-    agent.reset_conversation()
-    assert len(agent.chat_history) == 0
-    assert agent.generated_code is None
-
-
-def test_large_dataframe_handling():
-    """Test handling of large dataframes."""
-    # Create a large dataframe
-    df = pd.DataFrame({"x": range(10000), "y": range(10000)})
-
-    agent = PlotlyAgent()
-    agent.set_df(df)
-
-    # Test plot generation with large dataframe
-    code = """import plotly.express as px
-fig = px.scatter(df, x='x', y='y')"""
-    result = agent.execute_plotly_code(code)
-    assert "Code executed successfully" in result
-    assert agent.execution_env.fig is not None
-
-
-def test_input_validation():
-    """Test validation of input parameters."""
-    # Test invalid dataframe input
-    with pytest.raises(AssertionError):
-        agent = PlotlyAgent()
-        agent.set_df("not a dataframe")
-
-    # Test invalid SQL query input
-    df = pd.DataFrame({"x": [1, 2, 3]})
-    agent = PlotlyAgent()
-    with pytest.raises(AssertionError):
-        agent.set_df(df, sql_query=123)  # SQL query should be string
-
-    # Test invalid message input
-    agent.set_df(df)
-    with pytest.raises(AssertionError):
-        agent.process_message(123)  # Message should be string
-
-    # Test invalid code input
-    with pytest.raises(AssertionError):
-        agent.execute_plotly_code(123)  # Code should be string
-
-
-def test_complex_plot_handling():
-    """Test handling of complex plots with multiple traces and layouts."""
-    df = pd.DataFrame(
-        {
-            "x": [1, 2, 3, 4, 5],
-            "y1": [10, 20, 30, 40, 50],
-            "y2": [50, 40, 30, 20, 10],
-            "category": ["A", "B", "A", "B", "A"],
-        }
-    )
-
-    agent = PlotlyAgent()
-    agent.set_df(df)
-
-    complex_code = """import plotly.graph_objects as go
-fig = go.Figure()
-fig.add_trace(go.Scatter(x=df['x'], y=df['y1'], name='Trace 1'))
-fig.add_trace(go.Scatter(x=df['x'], y=df['y2'], name='Trace 2'))
-fig.update_layout(
-    title='Complex Plot',
-    xaxis_title='X Axis',
-    yaxis_title='Y Axis',
-    showlegend=True,
-    template='plotly_white'
-)"""
-
-    result = agent.execute_plotly_code(complex_code)
-    assert "Code executed successfully" in result
-    assert agent.execution_env.fig is not None