plot-agent 0.3.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

plot_agent/agent.py

@@ -1,11 +1,18 @@
+"""
+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.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 langgraph.prebuilt import create_react_agent
 from langchain_openai import ChatOpenAI
 
 from plot_agent.prompt import DEFAULT_SYSTEM_PROMPT
@@ -24,12 +31,16 @@ class PlotAgent:
 
     def __init__(
         self,
-        model="gpt-4o-mini",
+        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.
@@ -42,13 +53,43 @@ class PlotAgent:
             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)
+        # 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
@@ -94,6 +135,10 @@ class PlotAgent:
 
         # 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:
         """
@@ -123,7 +168,7 @@ class PlotAgent:
 
         # 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}"
 
@@ -150,64 +195,69 @@ class PlotAgent:
         """
         View the generated code.
         """
-        return self.generated_code
+        return self.generated_code or ""
 
     def _initialize_agent(self):
-        """Initialize the LangChain agent with the necessary tools and prompt."""
+        """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.",
+                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,
             ),
         ]
 
-        # Create system prompt with dataframe information
+        # Prepare 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"),
-            ]
+            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,
         )
 
-        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,
+        # 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."
@@ -215,33 +265,135 @@ class PlotAgent:
         if not self.agent_executor:
             return "Please set a dataframe first using set_df() method."
 
-        # Add user message to chat history
+        # Add user message to outward-facing 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}
+        # 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},
         )
 
-        # 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"]
+        # 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."""

plot_agent/execution.py

@@ -1,90 +1,250 @@
-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 threading
 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
+
+
+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):
+        """
+        Walk the AST and enforce:
+         • any Import/ImportFrom must be from _ALLOWED_MODULES
+         • no __dunder__ attribute access
         """
-        Execute the provided code and capture the fig object if created.
+        # 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 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()
+        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 if it was set
+            if timeout_set:
+                try:
+                    signal.alarm(0)
+                except Exception:
+                    pass
 
-            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,3 +1,7 @@
+"""
+This module contains the models for the PlotAgent.
+"""
+
 from pydantic import BaseModel, Field
 
 

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.3.0.dist-info → plot_agent-0.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: plot-agent
-Version: 0.3.0
+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.4.0.dist-info/RECORD

@@ -0,0 +1,10 @@
+plot_agent/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+plot_agent/agent.py,sha256=9N33lBkXe3PSYP6SsQkX2X2CAu7n2UjkRZap8hgxwzo,16733
+plot_agent/execution.py,sha256=FaMWKyFxQewrVx5tpHKLI7WafE9Q2ogvXhOVYZ4G3hw,8086
+plot_agent/models.py,sha256=THdGGGfGmRZ5rtgXvjPcQxFRRTZVFoADEHI_lsMVha8,860
+plot_agent/prompt.py,sha256=5hBlF7jdMrj6MiGEL7YmSDWFUfiCXyIZfZtf3NstKoo,3125
+plot_agent-0.4.0.dist-info/licenses/LICENSE,sha256=A4DPih7wHrh4VMEG3p1PhorqdhjmGIo8nQdYNQL7daA,1062
+plot_agent-0.4.0.dist-info/METADATA,sha256=iNH_2qd_k_Jp3RyWiFRitY4Vwx4ZHFNwMSFFodzMcfE,4152
+plot_agent-0.4.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+plot_agent-0.4.0.dist-info/top_level.txt,sha256=KyOjpihUssx26Ra-37vKUQ71pI2qgJsHaRwXHJUhjzQ,11
+plot_agent-0.4.0.dist-info/RECORD,,

{plot_agent-0.3.0.dist-info → plot_agent-0.4.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (78.1.1)
+Generator: setuptools (80.9.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

plot_agent-0.3.0.dist-info/RECORD

@@ -1,10 +0,0 @@
-plot_agent/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-plot_agent/agent.py,sha256=FPHma_u_77y3c5nBrCPS0VR8AWz-NFFRqURn0xHph-M,9584
-plot_agent/execution.py,sha256=Ewnz6Pb2EWM0pDAmeVoJ5RmhY5gmO1F7iOvERkj2D28,2770
-plot_agent/models.py,sha256=MkAaSELr54RfGRONKfiqCRA2ghlRbzTe5L5krVuqMsk,800
-plot_agent/prompt.py,sha256=HjRgbsAe8HHs8arQogvzOGQdThEWKRqQhtQyaUplxhQ,3064
-plot_agent-0.3.0.dist-info/licenses/LICENSE,sha256=A4DPih7wHrh4VMEG3p1PhorqdhjmGIo8nQdYNQL7daA,1062
-plot_agent-0.3.0.dist-info/METADATA,sha256=Cqw5gPpT2OM0fPaV_XLFJChEjsAFrVKfiwvxWxdVphM,2837
-plot_agent-0.3.0.dist-info/WHEEL,sha256=lTU6B6eIfYoiQJTZNc-fyaR6BpL6ehTzU3xGYxn2n8k,91
-plot_agent-0.3.0.dist-info/top_level.txt,sha256=KyOjpihUssx26Ra-37vKUQ71pI2qgJsHaRwXHJUhjzQ,11
-plot_agent-0.3.0.dist-info/RECORD,,