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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: plot-agent
-Version: 0.4.0
+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

plot_agent-0.5.0/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-0.4.0 → plot_agent-0.5.0}/plot_agent/agent.py

@@ -7,10 +7,10 @@ from io import StringIO
 import os
 import re
 import logging
-from typing import Optional
+from typing import List, Optional, Union
 from dotenv import load_dotenv
 
-from langchain_core.messages import AIMessage, HumanMessage
+from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
 from langchain_core.tools import Tool, StructuredTool
 from langgraph.prebuilt import create_react_agent
 from langchain_openai import ChatOpenAI
@@ -20,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:
     """
@@ -41,6 +50,7 @@ class PlotAgent:
         llm_timeout: int = 60,
         llm_max_retries: int = 1,
         debug: bool = False,
+        include_plot_image: bool = False,
     ):
         """
         Initialize the PlotAgent.
@@ -52,6 +62,11 @@ 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.
         """
         # Load .env if present, then require a valid API key
         load_dotenv()
@@ -76,6 +91,73 @@ class PlotAgent:
                 )
                 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,
@@ -97,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):
         """
@@ -131,7 +214,9 @@ 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()
@@ -191,6 +276,44 @@ 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.
@@ -230,6 +353,15 @@ 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,
+            ),
         ]
 
         # Prepare system prompt with dataframe information
@@ -300,10 +432,16 @@ class PlotAgent:
         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={"recursion_limit": self.max_iterations},
+            config=invoke_config,
         )
 
         # Extract the latest AI message from the returned messages
@@ -362,9 +500,14 @@ class PlotAgent:
                     )
                 ),
             ]
+            # 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={"recursion_limit": max(3, self.max_iterations // 2)},
+                config=retry_config,
             )
             self._graph_messages = retry_result.get("messages", [])
             retry_ai_messages = [
@@ -393,6 +536,16 @@ class PlotAgent:
                         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):
@@ -401,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-0.4.0 → plot_agent-0.5.0}/plot_agent/execution.py

@@ -9,12 +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
@@ -111,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,
@@ -128,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):
         """
@@ -167,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
@@ -179,6 +214,8 @@ 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,
@@ -209,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,
@@ -218,6 +257,8 @@ 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,
@@ -230,21 +271,64 @@ class PlotAgentExecutionEnvironment:
                 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-0.4.0 → plot_agent-0.5.0}/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-0.4.0 → plot_agent-0.5.0}/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.4.0 → plot_agent-0.5.0}/plot_agent.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: plot-agent
-Version: 0.4.0
+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

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

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