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

plot_agent/__init__.py

@@ -0,0 +1,5 @@
+"""plot-agent: LLM-powered Plotly visualization agent."""
+
+from plot_agent.agent import PlotAgent
+
+__all__ = ["PlotAgent"]

plot_agent/agent.py

@@ -4,12 +4,15 @@ This module contains the PlotAgent class, which is used to generate Plotly code
 
 import pandas as pd
 from io import StringIO
-from typing import Optional
+import os
+import re
+import logging
+from typing import List, Optional, Union
+from dotenv import load_dotenv
 
-from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
-from langchain_core.messages import AIMessage, HumanMessage
+from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
 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
@@ -17,9 +20,18 @@ from plot_agent.models import (
     GeneratedCodeInput,
     DoesFigExistInput,
     ViewGeneratedCodeInput,
+    CheckPlotOutputsInput,
 )
 from plot_agent.execution import PlotAgentExecutionEnvironment
 
+# Optional PostHog integration
+try:
+    from posthog import Posthog
+    from posthog.ai.langchain import CallbackHandler as PostHogCallbackHandler
+    POSTHOG_AVAILABLE = True
+except ImportError:
+    POSTHOG_AVAILABLE = False
+
 
 class PlotAgent:
     """
@@ -28,12 +40,17 @@ 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,
+        include_plot_image: bool = False,
     ):
         """
         Initialize the PlotAgent.
@@ -45,14 +62,116 @@ class PlotAgent:
             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.
+            llm_temperature (float): Temperature for LLM sampling.
+            llm_timeout (int): Timeout in seconds for LLM calls.
+            llm_max_retries (int): Maximum retries for LLM calls.
+            debug (bool): Enable debug logging.
+            include_plot_image (bool): Generate PNG image of plots for PostHog analytics.
         """
-        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)
+
+        # Initialize PostHog for LLM analytics (optional)
+        self.posthog_client = None
+        self.posthog_callback_handler = None
+        posthog_enabled = os.getenv("POSTHOG_ENABLED", "false").lower() == "true"
+
+        # Enable PostHog multimodal capture if include_plot_image is True
+        if include_plot_image and posthog_enabled:
+            os.environ["_INTERNAL_LLMA_MULTIMODAL"] = "true"
+
+        if posthog_enabled:
+            if not POSTHOG_AVAILABLE:
+                self._logger.warning(
+                    "PostHog is enabled but the posthog package is not installed. "
+                    "Install it with: pip install posthog"
+                )
+            else:
+                posthog_api_key = os.getenv("POSTHOG_API_KEY")
+                posthog_host = os.getenv("POSTHOG_HOST", "https://app.posthog.com")
+
+                if not posthog_api_key:
+                    self._logger.warning(
+                        "POSTHOG_ENABLED is true but POSTHOG_API_KEY is not set. "
+                        "PostHog tracking will be disabled."
+                    )
+                else:
+                    try:
+                        # Build super_properties for session tracking
+                        super_properties = {}
+
+                        # Add session ID from environment if provided
+                        ai_session_id = os.getenv("POSTHOG_AI_SESSION_ID")
+                        if ai_session_id:
+                            super_properties["$ai_session_id"] = ai_session_id
+
+                        # Initialize PostHog client with super_properties
+                        self.posthog_client = Posthog(
+                            posthog_api_key,
+                            host=posthog_host,
+                            super_properties=super_properties
+                        )
+
+                        # Build callback handler config
+                        callback_config = {"client": self.posthog_client}
+
+                        # Add optional distinct_id
+                        distinct_id = os.getenv("POSTHOG_DISTINCT_ID")
+                        if distinct_id:
+                            callback_config["distinct_id"] = distinct_id
+
+                        # Add privacy mode setting
+                        privacy_mode = os.getenv("POSTHOG_PRIVACY_MODE", "false").lower() == "true"
+                        callback_config["privacy_mode"] = privacy_mode
+
+                        self.posthog_callback_handler = PostHogCallbackHandler(**callback_config)
+
+                        if self.debug:
+                            session_info = f"session_id={ai_session_id}" if ai_session_id else "no session"
+                            self._logger.debug(
+                                f"PostHog LLM analytics initialized (host={posthog_host}, "
+                                f"distinct_id={distinct_id or 'anonymous'}, "
+                                f"privacy_mode={privacy_mode}, {session_info})"
+                            )
+                    except Exception as e:
+                        self._logger.error(f"Failed to initialize PostHog: {e}")
+                        self.posthog_client = None
+                        self.posthog_callback_handler = None
+
+        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
@@ -60,6 +179,7 @@ class PlotAgent:
         self.max_iterations = max_iterations
         self.early_stopping_method = early_stopping_method
         self.handle_parsing_errors = handle_parsing_errors
+        self.include_plot_image = include_plot_image
 
     def set_df(self, df: pd.DataFrame, sql_query: Optional[str] = None):
         """
@@ -94,10 +214,16 @@ class PlotAgent:
         self.sql_query = sql_query
 
         # Initialize execution environment
-        self.execution_env = PlotAgentExecutionEnvironment(df)
+        self.execution_env = PlotAgentExecutionEnvironment(
+            df, include_plot_image=self.include_plot_image
+        )
 
         # 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:
         """
@@ -150,14 +276,52 @@ class PlotAgent:
         else:
             return "No figure has been created yet."
 
+    def check_plot_outputs(self, *args, **kwargs) -> str:
+        """
+        Check if all required plot outputs (fig, plot_title, plot_summary) are available.
+
+        Args:
+            *args: Any positional arguments (ignored)
+            **kwargs: Any keyword arguments (ignored)
+
+        Returns:
+            str: A message indicating which plot outputs are available.
+        """
+        if not self.execution_env:
+            return "No execution environment has been initialized. Please set a dataframe first."
+
+        available = []
+        missing = []
+        
+        if self.execution_env.fig is not None:
+            available.append("fig")
+        else:
+            missing.append("fig")
+            
+        if self.execution_env.plot_title is not None:
+            available.append("plot_title")
+        else:
+            missing.append("plot_title")
+            
+        if self.execution_env.plot_summary is not None:
+            available.append("plot_summary")
+        else:
+            missing.append("plot_summary")
+        
+        if not missing:
+            return "All required plot outputs are available: fig, plot_title, and plot_summary."
+        else:
+            status = f"Available: {', '.join(available) if available else 'none'}. Missing: {', '.join(missing)}."
+            return status
+
     def view_generated_code(self, *args, **kwargs) -> str:
         """
         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 = [
@@ -189,39 +353,43 @@ class PlotAgent:
                 ),
                 args_schema=ViewGeneratedCodeInput,
             ),
+            StructuredTool.from_function(
+                func=self.check_plot_outputs,
+                name="check_plot_outputs",
+                description=(
+                    "Check if all required plot outputs (fig, plot_title, plot_summary) are available. "
+                    "This tool takes no arguments and returns the status of all plot outputs."
+                ),
+                args_schema=CheckPlotOutputsInput,
+            ),
         ]
 
-        # 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."
@@ -229,33 +397,156 @@ 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)}")
+
+        # Build config with optional PostHog callback
+        invoke_config = {"recursion_limit": self.max_iterations}
+        if self.posthog_callback_handler:
+            invoke_config["callbacks"] = [self.posthog_callback_handler]
+
+        # Invoke the LangGraph agent
+        result = self.agent_executor.invoke(
+            {"messages": graph_messages},
+            config=invoke_config,
         )
 
-        # 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."
+                    )
+                ),
+            ]
+            # Build config with optional PostHog callback for retry
+            retry_config = {"recursion_limit": max(3, self.max_iterations // 2)}
+            if self.posthog_callback_handler:
+                retry_config["callbacks"] = [self.posthog_callback_handler]
+
+            retry_result = self.agent_executor.invoke(
+                {"messages": guided_messages},
+                config=retry_config,
+            )
+            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}")
+
+        # Run verification step with image if enabled and figure exists
+        # This sends the plot image to the LLM for verification, which gets captured by PostHog
+        if (
+            self.include_plot_image
+            and self.posthog_callback_handler
+            and self.execution_env
+            and self.execution_env.fig is not None
+        ):
+            self._verify_plot_with_image(user_message)
+
+        return ai_content if isinstance(ai_content, str) else str(ai_content)
 
     def get_figure(self):
         """Return the current figure if one exists."""
@@ -263,7 +554,100 @@ class PlotAgent:
             return self.execution_env.fig
         return None
 
+    def get_plot_title(self):
+        """Return the current plot title if one exists."""
+        if self.execution_env and self.execution_env.plot_title:
+            return self.execution_env.plot_title
+        return None
+
+    def get_plot_summary(self):
+        """Return the current plot summary if one exists."""
+        if self.execution_env and self.execution_env.plot_summary:
+            return self.execution_env.plot_summary
+        return None
+
+    def get_plot_image_base64(self) -> Optional[str]:
+        """Return the current plot image as base64-encoded data URI."""
+        if self.execution_env and self.execution_env.plot_image_base64:
+            return self.execution_env.plot_image_base64
+        return None
+
+    def _verify_plot_with_image(self, user_request: str) -> Optional[str]:
+        """
+        Send the generated plot image to the LLM for verification.
+
+        This step serves two purposes:
+        1. Verifies the plot matches the user's request
+        2. Captures the plot image in PostHog LLM traces (via multimodal message)
+
+        Args:
+            user_request: The original user request for context.
+
+        Returns:
+            The LLM's verification response, or None if verification fails.
+        """
+        plot_image = self.get_plot_image_base64()
+        if not plot_image:
+            return None
+
+        # Build multimodal message with the plot image
+        human_content: List[Union[dict, str]] = [
+            {
+                "type": "text",
+                "text": (
+                    f"Please verify this generated plot matches the user's request.\n\n"
+                    f"User request: {user_request}\n\n"
+                    f"Generated code:\n```python\n{self.generated_code or 'N/A'}\n```\n\n"
+                    f"Plot title: {self.get_plot_title() or 'N/A'}\n"
+                    f"Plot summary: {self.get_plot_summary() or 'N/A'}\n\n"
+                    f"Respond with a brief confirmation that the plot is correct, "
+                    f"or note any issues you see."
+                ),
+            },
+            {
+                "type": "image_url",
+                "image_url": {"url": plot_image},
+            },
+        ]
+
+        messages = [
+            SystemMessage(
+                content=(
+                    "You are a plot verification assistant. "
+                    "Review the generated plot image and verify it matches the user's request. "
+                    "Be concise in your response."
+                )
+            ),
+            HumanMessage(content=human_content),
+        ]
+
+        try:
+            # Build config with PostHog callback to capture this verification step
+            invoke_config = {}
+            if self.posthog_callback_handler:
+                invoke_config["callbacks"] = [self.posthog_callback_handler]
+
+            if self.debug:
+                self._logger.debug("Running plot verification with image")
+
+            # Call LLM directly (not through agent) for verification
+            response = self.llm.invoke(messages, config=invoke_config)
+            verification_result = response.content if hasattr(response, "content") else str(response)
+
+            if self.debug:
+                self._logger.debug(f"Plot verification result: {verification_result[:200]}...")
+
+            return verification_result
+        except Exception as e:
+            self._logger.warning(f"Plot verification failed: {e}")
+            return None
+
     def reset_conversation(self):
         """Reset the conversation history."""
         self.chat_history = []
         self.generated_code = None
+        if self.execution_env:
+            self.execution_env.fig = None
+            self.execution_env.plot_title = None
+            self.execution_env.plot_summary = None
+            self.execution_env.plot_image_base64 = None

plot_agent/execution.py

@@ -9,11 +9,15 @@ Security features:
   • Enforce a 60 second timeout via signal.alarm
 """
 import ast
+import base64
 import builtins
+import logging
 import signal
+import threading
 import traceback
 from io import StringIO
 import contextlib
+from typing import Optional
 
 import pandas as pd
 import numpy as np
@@ -110,11 +114,17 @@ class PlotAgentExecutionEnvironment:
         "__import__": _safe_import,
     }
 
-    def __init__(self, df: pd.DataFrame):
+    def __init__(self, df: pd.DataFrame, include_plot_image: bool = False):
         """
         Initialize the execution environment with a dataframe.
+
+        Args:
+            df: The pandas dataframe to use for plotting.
+            include_plot_image: If True, generate a PNG image of the plot after execution.
         """
         self.df = df
+        self.include_plot_image = include_plot_image
+        self._logger = logging.getLogger("plot_agent.execution")
         # Base namespace for both globals & locals
         self._base_ns = {
             "__builtins__": self._SAFE_BUILTINS,
@@ -127,6 +137,30 @@ class PlotAgentExecutionEnvironment:
             "make_subplots": make_subplots,
         }
         self.fig = None
+        self.plot_title = None
+        self.plot_summary = None
+        self.plot_image_base64 = None
+
+    def _generate_plot_png(self, fig, width: int = 800, height: int = 600) -> Optional[str]:
+        """
+        Generate PNG as base64 data URI from a Plotly figure.
+
+        Args:
+            fig: The Plotly figure to convert.
+            width: Image width in pixels.
+            height: Image height in pixels.
+
+        Returns:
+            Base64-encoded data URI string, or None if generation fails.
+        """
+        try:
+            import plotly.io as pio
+            img_bytes = pio.to_image(fig, format='png', width=width, height=height)
+            b64 = base64.b64encode(img_bytes).decode('utf-8')
+            return f"data:image/png;base64,{b64}"
+        except Exception as e:
+            self._logger.warning(f"Failed to generate plot PNG: {e}")
+            return None
 
     def _validate_ast(self, node: ast.AST):
         """
@@ -166,8 +200,10 @@ class PlotAgentExecutionEnvironment:
 
         # Copy the base namespace
         ns = self._base_ns.copy()
-        # Purge any old `fig`
+        # Purge any old variables
         ns.pop("fig", None)
+        ns.pop("plot_title", None)
+        ns.pop("plot_summary", None)
 
         try:
             # Parse the generated code
@@ -178,14 +214,23 @@ class PlotAgentExecutionEnvironment:
             # If the code is rejected on safety grounds, return an error
             return {
                 "fig": None,
+                "plot_title": None,
+                "plot_summary": 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)
+        # 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()
@@ -201,6 +246,8 @@ class PlotAgentExecutionEnvironment:
             tb = traceback.format_exc()
             return {
                 "fig": None,
+                "plot_title": None,
+                "plot_summary": None,
                 "output": out_buf.getvalue(),
                 "error": f"Code execution timed out: {te}\n{tb}",
                 "success": False,
@@ -210,29 +257,78 @@ class PlotAgentExecutionEnvironment:
             tb = traceback.format_exc()
             return {
                 "fig": None,
+                "plot_title": None,
+                "plot_summary": None,
                 "output": out_buf.getvalue(),
                 "error": f"Error executing code: {e}\n{tb}",
                 "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`
+        # Get the variables
         fig = ns.get("fig")
+        plot_title = ns.get("plot_title")
+        plot_summary = ns.get("plot_summary")
+
+        # Store the variables
         self.fig = fig
+        self.plot_title = plot_title
+        self.plot_summary = plot_summary
+
+        # Generate PNG if enabled and figure exists
+        self.plot_image_base64 = None
+        if self.include_plot_image and fig is not None:
+            self.plot_image_base64 = self._generate_plot_png(fig)
+
+        # Validate required variables
+        missing_vars = []
         if fig is None:
+            missing_vars.append("fig")
+        if plot_title is None:
+            missing_vars.append("plot_title")
+        if plot_summary is None:
+            missing_vars.append("plot_summary")
+        
+        if missing_vars:
             return {
-                "fig": None,
+                "fig": fig,
+                "plot_title": plot_title,
+                "plot_summary": plot_summary,
+                "output": out_buf.getvalue(),
+                "error": f"Missing required variables: {', '.join(missing_vars)}. Please create variables named: {', '.join(missing_vars)}.",
+                "success": False,
+            }
+        
+        # Validate that plot_title and plot_summary are strings
+        validation_errors = []
+        if not isinstance(plot_title, str):
+            validation_errors.append("plot_title must be a string")
+        if not isinstance(plot_summary, str):
+            validation_errors.append("plot_summary must be a string")
+        
+        if validation_errors:
+            return {
+                "fig": fig,
+                "plot_title": plot_title,
+                "plot_summary": plot_summary,
                 "output": out_buf.getvalue(),
-                "error": "No `fig` created. Assign your figure to a variable named `fig`.",
+                "error": f"Validation errors: {'; '.join(validation_errors)}.",
                 "success": False,
             }
 
         # Return the result
         return {
             "fig": fig,
-            "output": "Code executed successfully. 'fig' object was created.",
+            "plot_title": plot_title,
+            "plot_summary": plot_summary,
+            "plot_image_base64": self.plot_image_base64,
+            "output": "Code executed successfully. 'fig', 'plot_title', and 'plot_summary' objects were created.",
             "error": "",
             "success": True,
         }

plot_agent/models.py

@@ -31,3 +31,9 @@ class ViewGeneratedCodeInput(BaseModel):
     """Model indicating that the view_generated_code function takes no arguments."""
 
     pass
+
+
+class CheckPlotOutputsInput(BaseModel):
+    """Model indicating that the check_plot_outputs function takes no arguments."""
+
+    pass

plot_agent/prompt.py

@@ -26,10 +26,12 @@ NOTES:
 - You must paste the full code, not just a reference to the code.
 - You must not use fig.show() in your code as it will ultimately be executed elsewhere in a headless environment.
 - If you need to do any data cleaning or wrangling, do it in the code before generating the plotly code as preprocessing steps assume the data is in the pandas 'df' object.
+- Your code MUST create three variables: 'fig' (Plotly figure), 'plot_title' (string), and 'plot_summary' (string).
 
 TOOLS:
 - execute_plotly_code(generated_code) to execute the generated code.
 - does_fig_exist() to check that a fig object is available for display. This tool takes no arguments.
+- check_plot_outputs() to check if all required outputs (fig, plot_title, plot_summary) are available. This tool takes no arguments.
 - view_generated_code() to view the generated code if need to fix it. This tool takes no arguments.
 
 IMPORTANT CODE FORMATTING INSTRUCTIONS:
@@ -37,6 +39,9 @@ IMPORTANT CODE FORMATTING INSTRUCTIONS:
 2. Use descriptive variable names.
 3. DO NOT include fig.show() in your code - the visualization will be rendered externally.
 4. Ensure your code creates a variable named 'fig' that contains the Plotly figure object.
+5. You MUST also create two string variables:
+   - 'plot_title': A concise, descriptive title for the plot (string)
+   - 'plot_summary': A brief summary explaining what the plot shows and any key insights (string)
 
 When a user asks for a visualization:
 1. YOU MUST ALWAYS use the execute_plotly_code(generated_code) tool to test and run your code.
@@ -48,11 +53,11 @@ IMPORTANT: The code you generate MUST be executed using the execute_plotly_code
 YOU MUST CALL execute_plotly_code WITH THE FULL CODE, NOT JUST A REFERENCE TO THE CODE.
 
 YOUR WORKFLOW MUST BE:
-1. execute_plotly_code(generated_code) to make sure the code is ran and a figure object is created.
-2. check that a figure object is available using does_fig_exist() to make sure the figure object was created.
-3. if there are errors, view the generated code using view_generated_code() to see what went wrong.
-4. fix the code and execute it again with execute_plotly_code(generated_code) to make sure the figure object is created.
-5. repeat until the figure object is available.
+1. execute_plotly_code(generated_code) to make sure the code is ran and a figure object, plot_title, and plot_summary are created.
+2. use check_plot_outputs() to verify that all required outputs (fig, plot_title, plot_summary) are available.
+3. if there are errors or missing outputs, view the generated code using view_generated_code() to see what went wrong.
+4. fix the code and execute it again with execute_plotly_code(generated_code) to make sure all required outputs are created.
+5. repeat until all outputs (figure object, plot_title, and plot_summary) are available.
 
 Always return the final working code (with all the comments) to the user along with an explanation of what the visualization shows.
 Make sure to follow best practices for data visualization, such as appropriate chart types, labels, and colors.

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

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

@@ -0,0 +1,10 @@
+plot_agent/__init__.py,sha256=hDmKvr4Hqe-NWpSGX8B36eWjBBkxq-JChIep4gN9pXc,123
+plot_agent/agent.py,sha256=VcctFGizXPn5E3l-fD0QK_LHP8rXE6jjpZv65WE6lXY,27063
+plot_agent/execution.py,sha256=oeiYy7dLF5krOK_vt0zoNB5C7mbliAQq9veVUYtBGfo,11350
+plot_agent/models.py,sha256=0ca5absK0L3kOyZhkVm0V-iyVVBEQ7cw_DCililYTzg,996
+plot_agent/prompt.py,sha256=WdXDhoTLdYZSmCay5K-eBLrJZrRHqwyDxfYrpsxUa4k,3684
+plot_agent-0.5.0.dist-info/licenses/LICENSE,sha256=A4DPih7wHrh4VMEG3p1PhorqdhjmGIo8nQdYNQL7daA,1062
+plot_agent-0.5.0.dist-info/METADATA,sha256=1sY052ANI3VRnsY794K7W2P2C1UiBs1RWFDxVDOpem8,4152
+plot_agent-0.5.0.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+plot_agent-0.5.0.dist-info/top_level.txt,sha256=KyOjpihUssx26Ra-37vKUQ71pI2qgJsHaRwXHJUhjzQ,11
+plot_agent-0.5.0.dist-info/RECORD,,

{plot_agent-0.3.1.dist-info → plot_agent-0.5.0.dist-info}/WHEEL

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

plot_agent-0.3.1.dist-info/RECORD

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