plot-agent 0.1.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

plot_agent/{plotly_agent.py → agent.py}

@@ -7,79 +7,19 @@ from plotly.subplots import make_subplots
 from io import StringIO
 import traceback
 import sys
-import re
 from typing import Dict, List, Optional, Any
 
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain_core.messages import AIMessage, HumanMessage
-from langchain_core.tools import Tool
-from pydantic import BaseModel, Field
+from langchain_core.tools import Tool, StructuredTool
 from langchain.agents import AgentExecutor, create_openai_tools_agent
 from langchain_openai import ChatOpenAI
-
-
-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
-of their pandas DataFrame (df).
-
-You have access to a pandas df with the following information:
-
-df.info():
-```plaintext
-{df_info}
-```
-
-df.head():
-```plaintext
-{df_head}
-```
-
-{sql_context}
-
-NOTES:
-- You must use the execute_plotly_code(generated_code) tool to test and run your code
-- 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 be executed 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
-- Do not use fig.show() in your code as it will be executed in a headless environment.
-
-IMPORTANT CODE FORMATTING INSTRUCTIONS:
-1. Include thorough, detailed comments in your code to explain what each section does
-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. Structure your code with proper spacing for readability
-
-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
-2. If there are errors, fix the code and run it again with execute_plotly_code(generated_code)
-3. Check that a figure object is available using does_fig_exist(). does_fig_exist() takes no arguments.
-
-IMPORTANT: The code you generate MUST be executed using the execute_plotly_code tool or no figure will be created!
-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) → 2. check that a figure object is available using does_fig_exist() → 3. (if needed) fix and execute again
-
-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.
-
-Remember that users may want to iterate on their visualizations, so be responsive to requests for changes.
-"""
-
-
-# Define input schemas for the tools
-class PlotDescriptionInput(BaseModel):
-    plot_description: str = Field(
-        ..., description="Description of the plot the user wants to create"
-    )
-
-
-class GeneratedCodeInput(BaseModel):
-    generated_code: str = Field(
-        ..., description="Python code that creates a Plotly figure"
-    )
+from plot_agent.prompt import DEFAULT_SYSTEM_PROMPT
+from plot_agent.models import (
+    GeneratedCodeInput,
+    DoesFigExistInput,
+    ViewGeneratedCodeInput,
+)
 
 
 class PlotlyAgentExecutionEnvironment:
@@ -100,13 +40,6 @@ class PlotlyAgentExecutionEnvironment:
         self.error = None
         self.fig = None
 
-    def preprocess_code(self, generated_code: str) -> str:
-        """Preprocess code to remove fig.show() calls."""
-        # Remove fig.show() calls
-        generated_code = re.sub(r"fig\.show\(\s*\)", "", generated_code)
-        generated_code = re.sub(r"fig\.show\(.*\)", "", generated_code)
-        return generated_code
-
     def execute_code(self, generated_code: str) -> Dict[str, Any]:
         """
         Execute the provided code and capture the fig object if created.
@@ -120,23 +53,20 @@ class PlotlyAgentExecutionEnvironment:
         self.output = None
         self.error = None
 
-        # Preprocess code to remove fig.show() calls
-        processed_code = self.preprocess_code(generated_code)
-
         # Capture stdout
         old_stdout = sys.stdout
         sys.stdout = mystdout = StringIO()
 
         try:
             # Execute the code
-            exec(processed_code, globals(), self.locals_dict)
+            exec(generated_code, globals(), self.locals_dict)
 
             # Check if a fig object was created
             if "fig" in self.locals_dict:
                 self.fig = self.locals_dict["fig"]
-                self.output = "Code executed successfully. Figure object was created."
+                self.output = "Code executed successfully. 'fig' object was created."
             else:
-                print(f"no fig object created: {processed_code}")
+                print(f"no fig object created: {generated_code}")
                 self.error = "Code executed without errors, but no 'fig' object was created. Make sure your code creates a variable named 'fig'."
 
         except Exception as e:
@@ -166,13 +96,25 @@ class PlotlyAgent:
     A class that uses an LLM to generate Plotly code based on a user's plot description.
     """
 
-    def __init__(self, model="gpt-4o-mini", system_prompt: Optional[str] = None):
+    def __init__(
+        self,
+        model="gpt-4o-mini",
+        system_prompt: Optional[str] = None,
+        verbose: bool = True,
+        max_iterations: int = 10,
+        early_stopping_method: str = "force",
+        handle_parsing_errors: bool = True,
+    ):
         """
         Initialize the PlotlyAgent.
 
         Args:
             model (str): The model to use for the LLM.
             system_prompt (Optional[str]): The system prompt to use for the LLM.
+            verbose (bool): Whether to print verbose output from the agent.
+            max_iterations (int): Maximum number of iterations for the agent to take.
+            early_stopping_method (str): Method to use for early stopping.
+            handle_parsing_errors (bool): Whether to handle parsing errors gracefully.
         """
         self.llm = ChatOpenAI(model=model)
         self.df = None
@@ -184,6 +126,10 @@ class PlotlyAgent:
         self.agent_executor = None
         self.generated_code = None
         self.system_prompt = system_prompt or DEFAULT_SYSTEM_PROMPT
+        self.verbose = verbose
+        self.max_iterations = max_iterations
+        self.early_stopping_method = early_stopping_method
+        self.handle_parsing_errors = handle_parsing_errors
 
     def set_df(self, df: pd.DataFrame, sql_query: Optional[str] = None):
         """
@@ -257,6 +203,12 @@ class PlotlyAgent:
         else:
             return "No figure has been created yet."
 
+    def view_generated_code(self, *args, **kwargs) -> str:
+        """
+        View the generated code.
+        """
+        return self.generated_code
+
     def _initialize_agent(self):
         """Initialize the LangChain agent with the necessary tools and prompt."""
         tools = [
@@ -266,11 +218,17 @@ class PlotlyAgent:
                 description="Execute the provided Plotly code and return the result",
                 args_schema=GeneratedCodeInput,
             ),
-            Tool.from_function(
+            StructuredTool.from_function(
                 func=self.does_fig_exist,
                 name="does_fig_exist",
-                description="Check if a figure exists and is available for display",
-                args_schema=None,
+                description="Check if a figure exists and is available for display. This tool takes no arguments.",
+                args_schema=DoesFigExistInput,
+            ),
+            StructuredTool.from_function(
+                func=self.view_generated_code,
+                name="view_generated_code",
+                description="View the generated code.",
+                args_schema=ViewGeneratedCodeInput,
             ),
         ]
 
@@ -299,10 +257,10 @@ class PlotlyAgent:
         self.agent_executor = AgentExecutor(
             agent=agent,
             tools=tools,
-            verbose=True,
-            max_iterations=10,
-            early_stopping_method="force",
-            handle_parsing_errors=True,
+            verbose=self.verbose,
+            max_iterations=self.max_iterations,
+            early_stopping_method=self.early_stopping_method,
+            handle_parsing_errors=self.handle_parsing_errors,
         )
 
     def process_message(self, user_message: str) -> str:

plot_agent/models.py

@@ -0,0 +1,26 @@
+from pydantic import BaseModel, Field
+
+
+# Define input schemas for the tools
+class PlotDescriptionInput(BaseModel):
+    plot_description: str = Field(
+        ..., description="Description of the plot the user wants to create"
+    )
+
+
+class GeneratedCodeInput(BaseModel):
+    generated_code: str = Field(
+        ..., description="Python code that creates a Plotly figure"
+    )
+
+
+class DoesFigExistInput(BaseModel):
+    """Model indicating that the does_fig_exist function takes no arguments."""
+
+    pass
+
+
+class ViewGeneratedCodeInput(BaseModel):
+    """Model indicating that the view_generated_code function takes no arguments."""
+
+    pass

plot_agent/prompt.py

@@ -0,0 +1,57 @@
+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
+of their pandas DataFrame (df).
+
+You have access to a pandas df with the following information:
+
+df.info():
+```plaintext
+{df_info}
+```
+
+df.head():
+```plaintext
+{df_head}
+```
+
+{sql_context}
+
+NOTES:
+- You must use the execute_plotly_code(generated_code) tool run your code and use the does_fig_exist() tool to check that a fig object is available for display.
+- 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.
+
+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.
+- view_generated_code() to view the generated code if need to fix it. This tool takes no arguments.
+
+IMPORTANT CODE FORMATTING INSTRUCTIONS:
+1. Include thorough, detailed comments in your code to explain what each section does.
+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.
+
+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.
+2. If there are errors, view the generated code using view_generated_code() and fix the code.
+3. Check that a figure object is available using does_fig_exist(). does_fig_exist() takes no arguments.
+4. If the figure object is not available, repeat the process until it is available.
+
+IMPORTANT: The code you generate MUST be executed using the execute_plotly_code tool or no figure will be created!
+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.
+
+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.
+
+Remember that users may want to iterate on their visualizations, so be responsive to requests for changes.
+"""

{plot_agent-0.1.1.dist-info → plot_agent-0.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: plot-agent
-Version: 0.1.1
+Version: 0.2.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
@@ -14,6 +14,9 @@ Dynamic: license-file
 
 # Plot Agent
 
+[![Tests](https://github.com/andrewm4894/plot-agent/actions/workflows/test.yml/badge.svg)](https://github.com/andrewm4894/plot-agent/actions/workflows/test.yml)
+[![PyPI version](https://badge.fury.io/py/plot-agent.svg)](https://badge.fury.io/py/plot-agent)
+
 An AI-powered data visualization assistant that helps users create Plotly visualizations in Python.
 
 ## Installation
@@ -26,11 +29,15 @@ pip install plot-agent
 
 ## Usage
 
-Here's a simple example of how to use Plot Agent:
+See more examples in [/examples/](https://nbviewer.org/github/andrewm4894/plot-agent/tree/main/examples/) (via nbviewer so that can see the charts easily).
+
+Here's a simple minimal example of how to use Plot Agent:
 
 ```python
 import pandas as pd
-from plot_agent import PlotlyAgent
+from plot_agent.agent import PlotlyAgent
+
+# ensure OPENAI_API_KEY is set and available for langchain
 
 # Create a sample dataframe
 df = pd.DataFrame({
@@ -46,6 +53,43 @@ agent.set_df(df)
 
 # Process a visualization request
 response = agent.process_message("Create a line plot of x vs y")
+
+# Print generated code
+print(agent.generated_code)
+
+# Get fig
+fig = agent.get_figure()
+fig.show()
+```
+
+`agent.generated_code`:
+
+```python
+import pandas as pd
+import plotly.graph_objects as go
+
+# Creating a line plot of x vs y
+# Create a figure object
+fig = go.Figure()
+
+# Add a line trace to the figure
+fig.add_trace(
+    go.Scatter(
+        x=df['x'],  # The x values
+        y=df['y'],  # The y values
+        mode='lines+markers',  # Display both lines and markers
+        name='Line Plot',  # Name of the trace
+        line=dict(color='blue', width=2)  # Specify line color and width
+    )
+)
+
+# Adding titles and labels
+fig.update_layout(
+    title='Line Plot of x vs y',  # Plot title
+    xaxis_title='x',  # x-axis label
+    yaxis_title='y',  # y-axis label
+    template='plotly_white'  # A clean layout
+)
 ```
 
 ## Features

plot_agent-0.2.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+plot_agent/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+plot_agent/agent.py,sha256=zukHGVexMXOHHjtmb7dNB3l2iEpC0rq0a6UUpVPLe14,10845
+plot_agent/models.py,sha256=ZOWWeYaqmnKJarXYyXnBQQ4nwUe71ae_gMul3bZXaWU,644
+plot_agent/prompt.py,sha256=HjRgbsAe8HHs8arQogvzOGQdThEWKRqQhtQyaUplxhQ,3064
+plot_agent-0.2.0.dist-info/licenses/LICENSE,sha256=A4DPih7wHrh4VMEG3p1PhorqdhjmGIo8nQdYNQL7daA,1062
+plot_agent-0.2.0.dist-info/METADATA,sha256=1_3wlHu5E1_gONv4uckGTc4X9F7hHzXuQRv9EjAHIVI,2841
+plot_agent-0.2.0.dist-info/WHEEL,sha256=lTU6B6eIfYoiQJTZNc-fyaR6BpL6ehTzU3xGYxn2n8k,91
+plot_agent-0.2.0.dist-info/top_level.txt,sha256=KyOjpihUssx26Ra-37vKUQ71pI2qgJsHaRwXHJUhjzQ,11
+plot_agent-0.2.0.dist-info/RECORD,,

{plot_agent-0.1.1.dist-info → plot_agent-0.2.0.dist-info}/WHEEL

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

plot_agent-0.1.1.dist-info/RECORD

@@ -1,7 +0,0 @@
-plot_agent/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-plot_agent/plotly_agent.py,sha256=5ddz0zyMRwqsdLzIDZwJtyyX72zqTjyp06FiJl1WEqg,12686
-plot_agent-0.1.1.dist-info/licenses/LICENSE,sha256=A4DPih7wHrh4VMEG3p1PhorqdhjmGIo8nQdYNQL7daA,1062
-plot_agent-0.1.1.dist-info/METADATA,sha256=gGVDZymtbirRKJFJjmuCLJhClQZ8sfedNpFqgKPDUwE,1539
-plot_agent-0.1.1.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-plot_agent-0.1.1.dist-info/top_level.txt,sha256=KyOjpihUssx26Ra-37vKUQ71pI2qgJsHaRwXHJUhjzQ,11
-plot_agent-0.1.1.dist-info/RECORD,,