plot-agent 0.1.2__py3-none-any.whl → 0.2.1__py3-none-any.whl

plot_agent/{plotly_agent.py → agent.py}

@@ -7,168 +7,21 @@ from plotly.subplots import make_subplots
 from io import StringIO
 import traceback
 import sys
-import re
-from typing import Dict, List, Optional, Any
+from typing import Dict, Optional, Any
 
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain_core.messages import AIMessage, HumanMessage
 from langchain_core.tools import Tool, StructuredTool
-from pydantic import BaseModel, Field
 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 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.
-"""
-
-
-# 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
-
-
-class PlotlyAgentExecutionEnvironment:
-    """Environment to safely execute plotly code and capture the fig object."""
-
-    def __init__(self, df: pd.DataFrame):
-        self.df = df
-        self.locals_dict = {
-            "df": df,
-            "px": px,
-            "go": go,
-            "pd": pd,
-            "np": np,
-            "plt": plt,
-            "make_subplots": make_subplots,
-        }
-        self.output = None
-        self.error = None
-        self.fig = None
-
-    def execute_code(self, generated_code: str) -> Dict[str, Any]:
-        """
-        Execute the provided code and capture the fig object if created.
-
-        Args:
-            generated_code (str): The code to execute.
-
-        Returns:
-            Dict[str, Any]: A dictionary containing the fig object, output, error, and success status.
-        """
-        self.output = None
-        self.error = None
-
-        # Capture stdout
-        old_stdout = sys.stdout
-        sys.stdout = mystdout = StringIO()
-
-        try:
-            # Execute the code
-            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. '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'."
-
-        except Exception as e:
-            self.error = f"Error executing code: {str(e)}\n{traceback.format_exc()}"
-
-        finally:
-            # Restore stdout
-            sys.stdout = old_stdout
-            captured_output = mystdout.getvalue()
-
-            if captured_output.strip():
-                if self.output:
-                    self.output += f"\nOutput:\n{captured_output}"
-                else:
-                    self.output = f"Output:\n{captured_output}"
-
-        return {
-            "fig": self.fig,
-            "output": self.output,
-            "error": self.error,
-            "success": self.error is None and self.fig is not None,
-        }
+from plot_agent.prompt import DEFAULT_SYSTEM_PROMPT
+from plot_agent.models import (
+    GeneratedCodeInput,
+    DoesFigExistInput,
+    ViewGeneratedCodeInput,
+)
+from plot_agent.execution import PlotlyAgentExecutionEnvironment
 
 
 class PlotlyAgent:
@@ -176,13 +29,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
@@ -194,6 +59,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):
         """
@@ -321,10 +190,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/execution.py

@@ -0,0 +1,79 @@
+import sys
+from io import StringIO
+import traceback
+import pandas as pd
+import plotly.express as px
+import plotly.graph_objects as go
+import numpy as np
+import matplotlib.pyplot as plt
+from plotly.subplots import make_subplots
+from typing import Dict, Any
+
+
+class PlotlyAgentExecutionEnvironment:
+    """Environment to safely execute plotly code and capture the fig object."""
+
+    def __init__(self, df: pd.DataFrame):
+        self.df = df
+        self.locals_dict = {
+            "df": df,
+            "px": px,
+            "go": go,
+            "pd": pd,
+            "np": np,
+            "plt": plt,
+            "make_subplots": make_subplots,
+        }
+        self.output = None
+        self.error = None
+        self.fig = None
+
+    def execute_code(self, generated_code: str) -> Dict[str, Any]:
+        """
+        Execute the provided code and capture the fig object if created.
+
+        Args:
+            generated_code (str): The code to execute.
+
+        Returns:
+            Dict[str, Any]: A dictionary containing the fig object, output, error, and success status.
+        """
+        self.output = None
+        self.error = None
+
+        # Capture stdout
+        old_stdout = sys.stdout
+        sys.stdout = mystdout = StringIO()
+
+        try:
+            # Execute the code
+            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. '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'."
+
+        except Exception as e:
+            self.error = f"Error executing code: {str(e)}\n{traceback.format_exc()}"
+
+        finally:
+            # Restore stdout
+            sys.stdout = old_stdout
+            captured_output = mystdout.getvalue()
+
+            if captured_output.strip():
+                if self.output:
+                    self.output += f"\nOutput:\n{captured_output}"
+                else:
+                    self.output = f"Output:\n{captured_output}"
+
+        return {
+            "fig": self.fig,
+            "output": self.output,
+            "error": self.error,
+            "success": self.error is None and self.fig is not None,
+        }

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: plot-agent
-Version: 0.1.2
+Version: 0.2.1
 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,12 +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.plotly_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({
@@ -47,10 +53,45 @@ 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
 
 - AI-powered visualization generation

plot_agent-0.2.1.dist-info/RECORD

@@ -0,0 +1,10 @@
+plot_agent/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+plot_agent/agent.py,sha256=8JhZf0r6_0HjV1dATzbiVgkf8g1fSLvOo8dLJkFzDVA,8682
+plot_agent/execution.py,sha256=B7XJCJOlOd99PUG5K9K9Al3TgvQtyta0X9LAxvodARo,2479
+plot_agent/models.py,sha256=ZOWWeYaqmnKJarXYyXnBQQ4nwUe71ae_gMul3bZXaWU,644
+plot_agent/prompt.py,sha256=HjRgbsAe8HHs8arQogvzOGQdThEWKRqQhtQyaUplxhQ,3064
+plot_agent-0.2.1.dist-info/licenses/LICENSE,sha256=A4DPih7wHrh4VMEG3p1PhorqdhjmGIo8nQdYNQL7daA,1062
+plot_agent-0.2.1.dist-info/METADATA,sha256=0GZe10CtdaXJ9qySDPFKjeCwmaTACYye6e23mCayMD8,2841
+plot_agent-0.2.1.dist-info/WHEEL,sha256=lTU6B6eIfYoiQJTZNc-fyaR6BpL6ehTzU3xGYxn2n8k,91
+plot_agent-0.2.1.dist-info/top_level.txt,sha256=KyOjpihUssx26Ra-37vKUQ71pI2qgJsHaRwXHJUhjzQ,11
+plot_agent-0.2.1.dist-info/RECORD,,

{plot_agent-0.1.2.dist-info → plot_agent-0.2.1.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.2.dist-info/RECORD

@@ -1,7 +0,0 @@
-plot_agent/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-plot_agent/plotly_agent.py,sha256=93DSwnPU7_915nsXs8SIUuL1hXFTEozBaqWQxs4HM9Y,13637
-plot_agent-0.1.2.dist-info/licenses/LICENSE,sha256=A4DPih7wHrh4VMEG3p1PhorqdhjmGIo8nQdYNQL7daA,1062
-plot_agent-0.1.2.dist-info/METADATA,sha256=GddWhRuPQ3DpujQGTANyk22oMsMASmZ2bdjz9CkZfdc,1589
-plot_agent-0.1.2.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-plot_agent-0.1.2.dist-info/top_level.txt,sha256=KyOjpihUssx26Ra-37vKUQ71pI2qgJsHaRwXHJUhjzQ,11
-plot_agent-0.1.2.dist-info/RECORD,,