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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: plot-agent
-Version: 0.3.1
+Version: 0.4.0
 Summary: An AI-powered data visualization assistant using Plotly
 Author-email: andrewm4894 <andrewm4894@gmail.com>
 Project-URL: Homepage, https://github.com/andrewm4894/plot-agent
@@ -19,6 +19,8 @@ Dynamic: license-file
 
 An AI-powered data visualization assistant that helps users create Plotly visualizations in Python.
 
+Built on LangGraph with tool-calling to reliably execute generated Plotly code in a sandbox and keep the current `fig` in sync.
+
 ## Installation
 
 You can install the package using pip:
@@ -37,7 +39,7 @@ Here's a simple minimal example of how to use Plot Agent:
 import pandas as pd
 from plot_agent.agent import PlotAgent
 
-# ensure OPENAI_API_KEY is set and available for langchain
+# ensure OPENAI_API_KEY is set (env or .env); optional debug via PLOT_AGENT_DEBUG=1
 
 # Create a sample dataframe
 df = pd.DataFrame({
@@ -92,19 +94,72 @@ fig.update_layout(
 )
 ```
 
+## How it works
+
+```mermaid
+flowchart TD
+    A[User message] --> B{LangGraph ReAct Agent}
+    subgraph Tools
+      T1[execute_plotly_code<br/>- runs code in sandbox<br/>- returns success/fig/error]
+      T2[does_fig_exist]
+      T3[view_generated_code]
+    end
+    B -- tool call --> T1
+    T1 -- result --> B
+    B -- optional --> T2
+    B -- optional --> T3
+    B --> C[AI response]
+    C --> D{Agent wrapper}
+    D -- persist messages --> B
+    D -- extract code blocks --> E[Sandbox execution]
+    E --> F[fig]
+    F --> G[get_figure]
+```
+
+- The LangGraph agent plans and decides when to call tools.
+- The wrapper persists full graph messages between turns and executes any returned code blocks to keep `fig` updated.
+- A safe execution environment runs code with an allowlist and a main-thread-only timeout.
+
 ## Features
 
 - AI-powered visualization generation
 - Support for various Plotly chart types
 - Automatic data preprocessing
 - Interactive visualization capabilities
-- Integration with LangChain for advanced AI capabilities
+- LangGraph-based tool calling and control flow
+- Debug logging via `PlotAgent(debug=True)` or `PLOT_AGENT_DEBUG=1`
 
 ## Requirements
 
 - Python 3.8 or higher
 - Dependencies are automatically installed with the package
 
+## Development
+
+- Run unit tests:
+
+```bash
+make test
+```
+
+- Execute all example notebooks:
+
+```bash
+make run-examples
+```
+
+- Execute with debug logs enabled:
+
+```bash
+make run-examples-debug
+```
+
+- Quick CLI repro that prints evolving code each step:
+
+```bash
+make run-example-script
+```
+
 ## License
 
 This project is licensed under the MIT License - see the LICENSE file for details. 

{plot_agent-0.3.1 → plot_agent-0.4.0}/README.md

@@ -5,6 +5,8 @@
 
 An AI-powered data visualization assistant that helps users create Plotly visualizations in Python.
 
+Built on LangGraph with tool-calling to reliably execute generated Plotly code in a sandbox and keep the current `fig` in sync.
+
 ## Installation
 
 You can install the package using pip:
@@ -23,7 +25,7 @@ Here's a simple minimal example of how to use Plot Agent:
 import pandas as pd
 from plot_agent.agent import PlotAgent
 
-# ensure OPENAI_API_KEY is set and available for langchain
+# ensure OPENAI_API_KEY is set (env or .env); optional debug via PLOT_AGENT_DEBUG=1
 
 # Create a sample dataframe
 df = pd.DataFrame({
@@ -78,19 +80,72 @@ fig.update_layout(
 )
 ```
 
+## How it works
+
+```mermaid
+flowchart TD
+    A[User message] --> B{LangGraph ReAct Agent}
+    subgraph Tools
+      T1[execute_plotly_code<br/>- runs code in sandbox<br/>- returns success/fig/error]
+      T2[does_fig_exist]
+      T3[view_generated_code]
+    end
+    B -- tool call --> T1
+    T1 -- result --> B
+    B -- optional --> T2
+    B -- optional --> T3
+    B --> C[AI response]
+    C --> D{Agent wrapper}
+    D -- persist messages --> B
+    D -- extract code blocks --> E[Sandbox execution]
+    E --> F[fig]
+    F --> G[get_figure]
+```
+
+- The LangGraph agent plans and decides when to call tools.
+- The wrapper persists full graph messages between turns and executes any returned code blocks to keep `fig` updated.
+- A safe execution environment runs code with an allowlist and a main-thread-only timeout.
+
 ## Features
 
 - AI-powered visualization generation
 - Support for various Plotly chart types
 - Automatic data preprocessing
 - Interactive visualization capabilities
-- Integration with LangChain for advanced AI capabilities
+- LangGraph-based tool calling and control flow
+- Debug logging via `PlotAgent(debug=True)` or `PLOT_AGENT_DEBUG=1`
 
 ## Requirements
 
 - Python 3.8 or higher
 - Dependencies are automatically installed with the package
 
+## Development
+
+- Run unit tests:
+
+```bash
+make test
+```
+
+- Execute all example notebooks:
+
+```bash
+make run-examples
+```
+
+- Execute with debug logs enabled:
+
+```bash
+make run-examples-debug
+```
+
+- Quick CLI repro that prints evolving code each step:
+
+```bash
+make run-example-script
+```
+
 ## License
 
 This project is licensed under the MIT License - see the LICENSE file for details. 

plot_agent-0.4.0/plot_agent/agent.py

@@ -0,0 +1,407 @@
+"""
+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
+import os
+import re
+import logging
+from typing import Optional
+from dotenv import load_dotenv
+
+from langchain_core.messages import AIMessage, HumanMessage
+from langchain_core.tools import Tool, StructuredTool
+from langgraph.prebuilt import create_react_agent
+from langchain_openai import ChatOpenAI
+
+from plot_agent.prompt import DEFAULT_SYSTEM_PROMPT
+from plot_agent.models import (
+    GeneratedCodeInput,
+    DoesFigExistInput,
+    ViewGeneratedCodeInput,
+)
+from plot_agent.execution import PlotAgentExecutionEnvironment
+
+
+class PlotAgent:
+    """
+    A class that uses an LLM to generate Plotly code based on a user's plot description.
+    """
+
+    def __init__(
+        self,
+        model: str = "gpt-4o-mini",
+        system_prompt: Optional[str] = None,
+        verbose: bool = True,
+        max_iterations: int = 10,
+        early_stopping_method: str = "force",
+        handle_parsing_errors: bool = True,
+        llm_temperature: float = 0.0,
+        llm_timeout: int = 60,
+        llm_max_retries: int = 1,
+        debug: bool = False,
+    ):
+        """
+        Initialize the PlotAgent.
+
+        Args:
+            model (str): The model to use for the LLM.
+            system_prompt (Optional[str]): The system prompt to use for the LLM.
+            verbose (bool): Whether to print verbose output from the agent.
+            max_iterations (int): Maximum number of iterations for the agent to take.
+            early_stopping_method (str): Method to use for early stopping.
+            handle_parsing_errors (bool): Whether to handle parsing errors gracefully.
+        """
+        # Load .env if present, then require a valid API key
+        load_dotenv()
+        openai_api_key = os.getenv("OPENAI_API_KEY")
+        if not openai_api_key:
+            raise RuntimeError(
+                "OPENAI_API_KEY is not set. Provide it via environment or a .env file."
+            )
+        self.debug = debug or os.getenv("PLOT_AGENT_DEBUG") == "1"
+
+        # Configure logger
+        self._logger = logging.getLogger("plot_agent")
+        if self.debug:
+            self._logger.setLevel(logging.DEBUG)
+            if not self._logger.handlers:
+                handler = logging.StreamHandler()
+                handler.setFormatter(
+                    logging.Formatter(
+                        "%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+                        datefmt="%H:%M:%S",
+                    )
+                )
+                self._logger.addHandler(handler)
+
+        self.llm = ChatOpenAI(
+            model=model,
+            temperature=llm_temperature,
+            timeout=llm_timeout,
+            max_retries=llm_max_retries,
+        )
+        self.df = None
+        self.df_info = None
+        self.df_head = None
+        self.sql_query = None
+        self.execution_env = None
+        self.chat_history = []
+        # Internal graph-native message history, including tool messages
+        self._graph_messages = []
+        self.agent_executor = None
+        self.generated_code = None
+        self.system_prompt = system_prompt or DEFAULT_SYSTEM_PROMPT
+        self.verbose = verbose
+        self.max_iterations = max_iterations
+        self.early_stopping_method = early_stopping_method
+        self.handle_parsing_errors = handle_parsing_errors
+
+    def set_df(self, df: pd.DataFrame, sql_query: Optional[str] = None):
+        """
+        Set the dataframe and capture its schema and sample.
+
+        Args:
+            df (pd.DataFrame): The pandas dataframe to set.
+            sql_query (Optional[str]): The SQL query used to generate the dataframe.
+
+        Returns:
+            None
+        """
+
+        # Check df
+        assert isinstance(df, pd.DataFrame), "The dataframe must be a pandas dataframe."
+        assert not df.empty, "The dataframe must not be empty."
+
+        if sql_query:
+            assert isinstance(sql_query, str), "The SQL query must be a string."
+
+        self.df = df
+
+        # Capture df.info() output
+        buffer = StringIO()
+        df.info(buf=buffer)
+        self.df_info = buffer.getvalue()
+
+        # Capture df.head() as string representation
+        self.df_head = df.head().to_string()
+
+        # Store SQL query if provided
+        self.sql_query = sql_query
+
+        # Initialize execution environment
+        self.execution_env = PlotAgentExecutionEnvironment(df)
+
+        # Initialize the agent with tools
+        self._initialize_agent()
+        # Reset graph messages for a fresh session with this dataframe
+        self._graph_messages = []
+        if self.debug:
+            self._logger.debug("set_df() initialized execution environment and graph")
+
+    def execute_plotly_code(self, generated_code: str) -> str:
+        """
+        Execute the provided Plotly code and return the result.
+
+        Args:
+            generated_code (str): The Plotly code to execute.
+
+        Returns:
+            str: The result of the execution.
+        """
+        assert isinstance(generated_code, str), "The generated code must be a string."
+
+        if not self.execution_env:
+            return "Error: No dataframe has been set. Please set a dataframe first."
+
+        # Store this as the last generated code
+        self.generated_code = generated_code
+
+        # Execute the generated code
+        code_execution_result = self.execution_env.execute_code(generated_code)
+
+        # Extract the results from the code execution
+        code_execution_success = code_execution_result.get("success", False)
+        code_execution_output = code_execution_result.get("output", "")
+        code_execution_error = code_execution_result.get("error", "")
+
+        # Check if the code executed successfully
+        if code_execution_success:
+            return f"Success: {code_execution_output}"
+        else:
+            return f"Error: {code_execution_error}\n{code_execution_output}"
+
+    def does_fig_exist(self, *args, **kwargs) -> str:
+        """
+        Check if a figure object is available for display.
+
+        Args:
+            *args: Any positional arguments (ignored)
+            **kwargs: Any keyword arguments (ignored)
+
+        Returns:
+            str: A message indicating whether a figure is available for display.
+        """
+        if not self.execution_env:
+            return "No execution environment has been initialized. Please set a dataframe first."
+
+        if self.execution_env.fig is not None:
+            return "A figure is available for display."
+        else:
+            return "No figure has been created yet."
+
+    def view_generated_code(self, *args, **kwargs) -> str:
+        """
+        View the generated code.
+        """
+        return self.generated_code or ""
+
+    def _initialize_agent(self):
+        """Initialize a LangGraph ReAct agent with tools and keep API compatibility."""
+
+        # Initialize the tools
+        tools = [
+            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."
+                ),
+                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."
+                ),
+                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."
+                ),
+                args_schema=ViewGeneratedCodeInput,
+            ),
+        ]
+
+        # Prepare system prompt with dataframe information
+        sql_context = ""
+        if self.sql_query:
+            sql_context = (
+                "In case it is useful to help with the data understanding, the df was generated using the following SQL query:\n"
+                f"```sql\n{self.sql_query}\n```"
+            )
+
+        # Store formatted system instructions for the graph state modifier
+        self._system_message_content = self.system_prompt.format(
+            df_info=self.df_info,
+            df_head=self.df_head,
+            sql_context=sql_context,
+        )
+
+        # Create a ReAct agent graph with the provided tools and system prompt
+        self._graph = create_react_agent(
+            self.llm,
+            tools,
+            prompt=self._system_message_content,
+            debug=self.debug,
+        )
+
+        # Backwards-compatibility: expose under the old attribute name
+        self.agent_executor = self._graph
+
+    def process_message(self, user_message: str) -> str:
+        """Process a user message and return the agent's response."""
+        assert isinstance(user_message, str), "The user message must be a string."
+
+        if not self.agent_executor:
+            return "Please set a dataframe first using set_df() method."
+
+        # Add user message to outward-facing chat history
+        self.chat_history.append(HumanMessage(content=user_message))
+
+        # Reset generated_code
+        self.generated_code = None
+
+        # Short-circuit empty inputs to avoid graph recursion
+        if user_message.strip() == "":
+            ai_content = (
+                "Please provide a non-empty plotting request (e.g., 'scatter x vs y')."
+            )
+            self.chat_history.append(AIMessage(content=ai_content))
+            if self.debug:
+                self._logger.debug("empty message received; returning guidance without invoking graph")
+            return ai_content
+
+        # Short-circuit messages that are primarily raw code blocks without a visualization request
+        if "```" in user_message and not re.search(
+            r"\b(plot|chart|graph|visuali(s|z)e|figure|subplot|heatmap|bar|line|scatter)\b",
+            user_message,
+            flags=re.IGNORECASE,
+        ):
+            ai_content = (
+                "I see a code snippet. Please describe the visualization you want (e.g., 'line chart of y over x')."
+            )
+            self.chat_history.append(AIMessage(content=ai_content))
+            if self.debug:
+                self._logger.debug("code-only message received; returning guidance without invoking graph")
+            return ai_content
+
+        # Build graph messages (includes tool call/observation history)
+        graph_messages = [*self._graph_messages, HumanMessage(content=user_message)]
+        if self.debug:
+            self._logger.debug(f"process_message() user: {user_message}")
+            self._logger.debug(f"graph message count before invoke: {len(graph_messages)}")
+        # Invoke the LangGraph agent
+        result = self.agent_executor.invoke(
+            {"messages": graph_messages},
+            config={"recursion_limit": self.max_iterations},
+        )
+
+        # Extract the latest AI message from the returned messages
+        ai_messages = [m for m in result.get("messages", []) if isinstance(m, AIMessage)]
+        ai_content = ai_messages[-1].content if ai_messages else ""
+
+        # Persist full graph messages for future context
+        self._graph_messages = result.get("messages", [])
+        if self.debug:
+            self._logger.debug(f"graph message count after invoke: {len(self._graph_messages)}")
+
+        # Add agent response to outward-facing chat history
+        self.chat_history.append(AIMessage(content=ai_content))
+
+        # If the agent didn't execute the code via tool, but we have prior generated_code, execute it
+        if self.execution_env and self.execution_env.fig is None and self.generated_code is not None:
+            if self.debug:
+                self._logger.debug("executing stored generated_code because no fig exists yet")
+            exec_result = self.execution_env.execute_code(self.generated_code)
+            if self.debug:
+                self._logger.debug(f"execution result success={exec_result.get('success')} error={exec_result.get('error')!r}")
+
+        # If the assistant returned code in the message, execute it to update the figure
+        code_executed = False
+        if self.execution_env and isinstance(ai_content, str):
+            extracted_code = None
+            if "```python" in ai_content:
+                parts = ai_content.split("```python", 1)
+                extracted_code = parts[1].split("```", 1)[0].strip() if len(parts) > 1 else None
+            elif "```" in ai_content:
+                # Fallback: extract first generic fenced code block
+                parts = ai_content.split("```", 1)
+                if len(parts) > 1:
+                    extracted_code = parts[1].split("```", 1)[0].strip()
+            if extracted_code:
+                if (self.generated_code or "").strip() != extracted_code:
+                    self.generated_code = extracted_code
+                    if self.debug:
+                        self._logger.debug("executing code extracted from AI message")
+                    exec_result = self.execution_env.execute_code(extracted_code)
+                    if self.debug:
+                        self._logger.debug(f"execution result success={exec_result.get('success')} error={exec_result.get('error')!r}")
+                    code_executed = True
+
+        # If still no figure and no code was executed, run one guided retry to force tool usage
+        if self.execution_env and self.execution_env.fig is None and not code_executed:
+            if self.debug:
+                self._logger.debug("guided retry: prompting model to use execute_plotly_code tool")
+            guided_messages = [
+                *self._graph_messages,
+                HumanMessage(
+                    content=(
+                        "Please use the execute_plotly_code(generated_code) tool with the FULL code to "
+                        "create a variable named 'fig', then call does_fig_exist(). Return the final "
+                        "code in a fenced ```python block."
+                    )
+                ),
+            ]
+            retry_result = self.agent_executor.invoke(
+                {"messages": guided_messages},
+                config={"recursion_limit": max(3, self.max_iterations // 2)},
+            )
+            self._graph_messages = retry_result.get("messages", [])
+            retry_ai_messages = [
+                m for m in self._graph_messages if isinstance(m, AIMessage)
+            ]
+            retry_content = retry_ai_messages[-1].content if retry_ai_messages else ""
+            if isinstance(retry_content, str):
+                if "```python" in retry_content:
+                    parts = retry_content.split("```python", 1)
+                    retry_code = (
+                        parts[1].split("```", 1)[0].strip() if len(parts) > 1 else None
+                    )
+                elif "```" in retry_content:
+                    parts = retry_content.split("```", 1)
+                    retry_code = (
+                        parts[1].split("```", 1)[0].strip() if len(parts) > 1 else None
+                    )
+                else:
+                    retry_code = None
+                if retry_code:
+                    if (self.generated_code or "").strip() != retry_code:
+                        self.generated_code = retry_code
+                        if self.debug:
+                            self._logger.debug("executing code extracted from guided retry response")
+                        exec_result = self.execution_env.execute_code(retry_code)
+                        if self.debug:
+                            self._logger.debug(f"execution result success={exec_result.get('success')} error={exec_result.get('error')!r}")
+
+        return ai_content if isinstance(ai_content, str) else str(ai_content)
+
+    def get_figure(self):
+        """Return the current figure if one exists."""
+        if self.execution_env and self.execution_env.fig:
+            return self.execution_env.fig
+        return None
+
+    def reset_conversation(self):
+        """Reset the conversation history."""
+        self.chat_history = []
+        self.generated_code = None

{plot_agent-0.3.1 → plot_agent-0.4.0}/plot_agent/execution.py

@@ -11,6 +11,7 @@ Security features:
 import ast
 import builtins
 import signal
+import threading
 import traceback
 from io import StringIO
 import contextlib
@@ -183,9 +184,16 @@ class PlotAgentExecutionEnvironment:
                 "success": False,
             }
 
-        # Set a timeout
-        signal.signal(signal.SIGALRM, _timeout_handler)
-        signal.alarm(self.TIMEOUT_SECONDS)
+        # Set a timeout only if running on the main thread; signals are not supported in worker threads
+        timeout_set = False
+        try:
+            if threading.current_thread() is threading.main_thread():
+                signal.signal(signal.SIGALRM, _timeout_handler)
+                signal.alarm(self.TIMEOUT_SECONDS)
+                timeout_set = True
+        except Exception:
+            # If setting the signal handler fails (e.g., not in main thread), proceed without timeout
+            timeout_set = False
 
         # Execute the code
         out_buf, err_buf = StringIO(), StringIO()
@@ -215,8 +223,12 @@ class PlotAgentExecutionEnvironment:
                 "success": False,
             }
         finally:
-            # Reset the timeout
-            signal.alarm(0)
+            # Reset the timeout if it was set
+            if timeout_set:
+                try:
+                    signal.alarm(0)
+                except Exception:
+                    pass
 
         # Get the `fig`
         fig = ns.get("fig")

{plot_agent-0.3.1 → plot_agent-0.4.0}/plot_agent.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: plot-agent
-Version: 0.3.1
+Version: 0.4.0
 Summary: An AI-powered data visualization assistant using Plotly
 Author-email: andrewm4894 <andrewm4894@gmail.com>
 Project-URL: Homepage, https://github.com/andrewm4894/plot-agent
@@ -19,6 +19,8 @@ Dynamic: license-file
 
 An AI-powered data visualization assistant that helps users create Plotly visualizations in Python.
 
+Built on LangGraph with tool-calling to reliably execute generated Plotly code in a sandbox and keep the current `fig` in sync.
+
 ## Installation
 
 You can install the package using pip:
@@ -37,7 +39,7 @@ Here's a simple minimal example of how to use Plot Agent:
 import pandas as pd
 from plot_agent.agent import PlotAgent
 
-# ensure OPENAI_API_KEY is set and available for langchain
+# ensure OPENAI_API_KEY is set (env or .env); optional debug via PLOT_AGENT_DEBUG=1
 
 # Create a sample dataframe
 df = pd.DataFrame({
@@ -92,19 +94,72 @@ fig.update_layout(
 )
 ```
 
+## How it works
+
+```mermaid
+flowchart TD
+    A[User message] --> B{LangGraph ReAct Agent}
+    subgraph Tools
+      T1[execute_plotly_code<br/>- runs code in sandbox<br/>- returns success/fig/error]
+      T2[does_fig_exist]
+      T3[view_generated_code]
+    end
+    B -- tool call --> T1
+    T1 -- result --> B
+    B -- optional --> T2
+    B -- optional --> T3
+    B --> C[AI response]
+    C --> D{Agent wrapper}
+    D -- persist messages --> B
+    D -- extract code blocks --> E[Sandbox execution]
+    E --> F[fig]
+    F --> G[get_figure]
+```
+
+- The LangGraph agent plans and decides when to call tools.
+- The wrapper persists full graph messages between turns and executes any returned code blocks to keep `fig` updated.
+- A safe execution environment runs code with an allowlist and a main-thread-only timeout.
+
 ## Features
 
 - AI-powered visualization generation
 - Support for various Plotly chart types
 - Automatic data preprocessing
 - Interactive visualization capabilities
-- Integration with LangChain for advanced AI capabilities
+- LangGraph-based tool calling and control flow
+- Debug logging via `PlotAgent(debug=True)` or `PLOT_AGENT_DEBUG=1`
 
 ## Requirements
 
 - Python 3.8 or higher
 - Dependencies are automatically installed with the package
 
+## Development
+
+- Run unit tests:
+
+```bash
+make test
+```
+
+- Execute all example notebooks:
+
+```bash
+make run-examples
+```
+
+- Execute with debug logs enabled:
+
+```bash
+make run-examples-debug
+```
+
+- Quick CLI repro that prints evolving code each step:
+
+```bash
+make run-example-script
+```
+
 ## License
 
 This project is licensed under the MIT License - see the LICENSE file for details. 

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

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

plot_agent-0.3.1/plot_agent/agent.py

@@ -1,269 +0,0 @@
-"""
-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
-
-from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
-from langchain_core.messages import AIMessage, HumanMessage
-from langchain_core.tools import Tool, StructuredTool
-from langchain.agents import AgentExecutor, create_openai_tools_agent
-from langchain_openai import ChatOpenAI
-
-from plot_agent.prompt import DEFAULT_SYSTEM_PROMPT
-from plot_agent.models import (
-    GeneratedCodeInput,
-    DoesFigExistInput,
-    ViewGeneratedCodeInput,
-)
-from plot_agent.execution import PlotAgentExecutionEnvironment
-
-
-class PlotAgent:
-    """
-    A class that uses an LLM to generate Plotly code based on a user's plot description.
-    """
-
-    def __init__(
-        self,
-        model="gpt-4o-mini",
-        system_prompt: Optional[str] = None,
-        verbose: bool = True,
-        max_iterations: int = 10,
-        early_stopping_method: str = "force",
-        handle_parsing_errors: bool = True,
-    ):
-        """
-        Initialize the PlotAgent.
-
-        Args:
-            model (str): The model to use for the LLM.
-            system_prompt (Optional[str]): The system prompt to use for the LLM.
-            verbose (bool): Whether to print verbose output from the agent.
-            max_iterations (int): Maximum number of iterations for the agent to take.
-            early_stopping_method (str): Method to use for early stopping.
-            handle_parsing_errors (bool): Whether to handle parsing errors gracefully.
-        """
-        self.llm = ChatOpenAI(model=model)
-        self.df = None
-        self.df_info = None
-        self.df_head = None
-        self.sql_query = None
-        self.execution_env = None
-        self.chat_history = []
-        self.agent_executor = None
-        self.generated_code = None
-        self.system_prompt = system_prompt or DEFAULT_SYSTEM_PROMPT
-        self.verbose = verbose
-        self.max_iterations = max_iterations
-        self.early_stopping_method = early_stopping_method
-        self.handle_parsing_errors = handle_parsing_errors
-
-    def set_df(self, df: pd.DataFrame, sql_query: Optional[str] = None):
-        """
-        Set the dataframe and capture its schema and sample.
-
-        Args:
-            df (pd.DataFrame): The pandas dataframe to set.
-            sql_query (Optional[str]): The SQL query used to generate the dataframe.
-
-        Returns:
-            None
-        """
-
-        # Check df
-        assert isinstance(df, pd.DataFrame), "The dataframe must be a pandas dataframe."
-        assert not df.empty, "The dataframe must not be empty."
-
-        if sql_query:
-            assert isinstance(sql_query, str), "The SQL query must be a string."
-
-        self.df = df
-
-        # Capture df.info() output
-        buffer = StringIO()
-        df.info(buf=buffer)
-        self.df_info = buffer.getvalue()
-
-        # Capture df.head() as string representation
-        self.df_head = df.head().to_string()
-
-        # Store SQL query if provided
-        self.sql_query = sql_query
-
-        # Initialize execution environment
-        self.execution_env = PlotAgentExecutionEnvironment(df)
-
-        # Initialize the agent with tools
-        self._initialize_agent()
-
-    def execute_plotly_code(self, generated_code: str) -> str:
-        """
-        Execute the provided Plotly code and return the result.
-
-        Args:
-            generated_code (str): The Plotly code to execute.
-
-        Returns:
-            str: The result of the execution.
-        """
-        assert isinstance(generated_code, str), "The generated code must be a string."
-
-        if not self.execution_env:
-            return "Error: No dataframe has been set. Please set a dataframe first."
-
-        # Store this as the last generated code
-        self.generated_code = generated_code
-
-        # Execute the generated code
-        code_execution_result = self.execution_env.execute_code(generated_code)
-
-        # Extract the results from the code execution
-        code_execution_success = code_execution_result.get("success", False)
-        code_execution_output = code_execution_result.get("output", "")
-        code_execution_error = code_execution_result.get("error", "")
-
-        # Check if the code executed successfully
-        if code_execution_success:
-            return f"Success: {code_execution_output}"
-        else:
-            return f"Error: {code_execution_error}\n{code_execution_output}"
-
-    def does_fig_exist(self, *args, **kwargs) -> str:
-        """
-        Check if a figure object is available for display.
-
-        Args:
-            *args: Any positional arguments (ignored)
-            **kwargs: Any keyword arguments (ignored)
-
-        Returns:
-            str: A message indicating whether a figure is available for display.
-        """
-        if not self.execution_env:
-            return "No execution environment has been initialized. Please set a dataframe first."
-
-        if self.execution_env.fig is not None:
-            return "A figure is available for display."
-        else:
-            return "No figure has been created yet."
-
-    def view_generated_code(self, *args, **kwargs) -> str:
-        """
-        View the generated code.
-        """
-        return self.generated_code
-
-    def _initialize_agent(self):
-        """Initialize the LangChain agent with the necessary tools and prompt."""
-
-        # Initialize the tools
-        tools = [
-            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."
-                ),
-                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."
-                ),
-                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."
-                ),
-                args_schema=ViewGeneratedCodeInput,
-            ),
-        ]
-
-        # Create system prompt with dataframe information
-        sql_context = ""
-        if self.sql_query:
-            sql_context = f"In case it is useful to help with the data understanding, the df was generated using the following SQL query:\n```sql\n{self.sql_query}\n```"
-
-        prompt = ChatPromptTemplate.from_messages(
-            [
-                (
-                    "system",
-                    self.system_prompt.format(
-                        df_info=self.df_info,
-                        df_head=self.df_head,
-                        sql_context=sql_context,
-                    ),
-                ),
-                MessagesPlaceholder(variable_name="chat_history"),
-                ("human", "{input}"),
-                MessagesPlaceholder(variable_name="agent_scratchpad"),
-            ]
-        )
-
-        agent = create_openai_tools_agent(self.llm, tools, prompt)
-        self.agent_executor = AgentExecutor(
-            agent=agent,
-            tools=tools,
-            verbose=self.verbose,
-            max_iterations=self.max_iterations,
-            early_stopping_method=self.early_stopping_method,
-            handle_parsing_errors=self.handle_parsing_errors,
-        )
-
-    def process_message(self, user_message: str) -> str:
-        """Process a user message and return the agent's response."""
-        assert isinstance(user_message, str), "The user message must be a string."
-
-        if not self.agent_executor:
-            return "Please set a dataframe first using set_df() method."
-
-        # Add user message to chat history
-        self.chat_history.append(HumanMessage(content=user_message))
-
-        # Reset generated_code
-        self.generated_code = None
-
-        # Get response from agent
-        response = self.agent_executor.invoke(
-            {"input": user_message, "chat_history": self.chat_history}
-        )
-
-        # Add agent response to chat history
-        self.chat_history.append(AIMessage(content=response["output"]))
-
-        # If the agent didn't execute the code, but did generate code, execute it directly
-        if self.execution_env.fig is None and self.generated_code is not None:
-            self.execution_env.execute_code(self.generated_code)
-
-        # If we can extract code from the response when no code was executed, try that too
-        if self.execution_env.fig is None and "```python" in response["output"]:
-            code_blocks = response["output"].split("```python")
-            if len(code_blocks) > 1:
-                generated_code = code_blocks[1].split("```")[0].strip()
-                self.execution_env.execute_code(generated_code)
-
-        # Return the agent's response
-        return response["output"]
-
-    def get_figure(self):
-        """Return the current figure if one exists."""
-        if self.execution_env and self.execution_env.fig:
-            return self.execution_env.fig
-        return None
-
-    def reset_conversation(self):
-        """Reset the conversation history."""
-        self.chat_history = []
-        self.generated_code = None