plot-agent 0.1.1__tar.gz → 0.2.0__tar.gz

{plot_agent-0.1.1 → plot_agent-0.2.0}/PKG-INFO

@@ -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/README.md

@@ -0,0 +1,96 @@
+# 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
+
+You can install the package using pip:
+
+```bash
+pip install plot-agent
+```
+
+## Usage
+
+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.agent import PlotlyAgent
+
+# ensure OPENAI_API_KEY is set and available for langchain
+
+# Create a sample dataframe
+df = pd.DataFrame({
+    'x': [1, 2, 3, 4, 5],
+    'y': [10, 20, 30, 40, 50]
+})
+
+# Initialize the agent
+agent = PlotlyAgent()
+
+# Set the dataframe
+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
+- Support for various Plotly chart types
+- Automatic data preprocessing
+- Interactive visualization capabilities
+- Integration with LangChain for advanced AI capabilities
+
+## Requirements
+
+- Python 3.8 or higher
+- Dependencies are automatically installed with the package
+
+## License
+
+This project is licensed under the MIT License - see the LICENSE file for details. 

plot_agent-0.1.1/plot_agent/plotly_agent.py → plot_agent-0.2.0/plot_agent/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-0.2.0/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-0.2.0/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 → plot_agent-0.2.0}/plot_agent.egg-info/PKG-INFO

@@ -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.1.1 → plot_agent-0.2.0}/plot_agent.egg-info/SOURCES.txt

@@ -2,8 +2,11 @@ LICENSE
 README.md
 pyproject.toml
 plot_agent/__init__.py
-plot_agent/plotly_agent.py
+plot_agent/agent.py
+plot_agent/models.py
+plot_agent/prompt.py
 plot_agent.egg-info/PKG-INFO
 plot_agent.egg-info/SOURCES.txt
 plot_agent.egg-info/dependency_links.txt
-plot_agent.egg-info/top_level.txt
+plot_agent.egg-info/top_level.txt
+tests/test_plotly_agent.py

{plot_agent-0.1.1 → plot_agent-0.2.0}/pyproject.toml

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

plot_agent-0.2.0/tests/test_plotly_agent.py

@@ -0,0 +1,152 @@
+import pytest
+import pandas as pd
+import numpy as np
+from plot_agent.agent import PlotlyAgent
+from langchain_core.messages import HumanMessage, AIMessage
+
+def test_plotly_agent_initialization():
+    """Test that PlotlyAgent initializes correctly."""
+    agent = PlotlyAgent()
+    assert agent.llm is not None
+    assert agent.df is None
+    assert agent.df_info is None
+    assert agent.df_head is None
+    assert agent.sql_query is None
+    assert agent.execution_env is None
+    assert agent.chat_history == []
+    assert agent.agent_executor is None
+    assert agent.generated_code is None
+
+def test_set_df():
+    """Test that set_df properly sets up the dataframe and environment."""
+    # Create a sample dataframe
+    df = pd.DataFrame({
+        'x': [1, 2, 3, 4, 5],
+        'y': [10, 20, 30, 40, 50]
+    })
+    
+    agent = PlotlyAgent()
+    agent.set_df(df)
+    
+    assert agent.df is not None
+    assert agent.df_info is not None
+    assert agent.df_head is not None
+    assert agent.execution_env is not None
+    assert agent.agent_executor is not None
+
+def test_execute_plotly_code():
+    """Test that execute_plotly_code works with valid code."""
+    df = pd.DataFrame({
+        'x': [1, 2, 3, 4, 5],
+        'y': [10, 20, 30, 40, 50]
+    })
+    
+    agent = PlotlyAgent()
+    agent.set_df(df)
+    
+    # Test with valid plotly code
+    valid_code = """import plotly.express as px
+fig = px.scatter(df, x='x', y='y')"""
+    
+    result = agent.execute_plotly_code(valid_code)
+    assert "Code executed successfully" in result
+    assert agent.execution_env.fig is not None
+
+def test_execute_plotly_code_with_error():
+    """Test that execute_plotly_code handles errors properly."""
+    df = pd.DataFrame({
+        'x': [1, 2, 3, 4, 5],
+        'y': [10, 20, 30, 40, 50]
+    })
+    
+    agent = PlotlyAgent()
+    agent.set_df(df)
+    
+    # Test with invalid code
+    invalid_code = """import plotly.express as px
+fig = px.scatter(df, x='non_existent_column', y='y')"""
+    
+    result = agent.execute_plotly_code(invalid_code)
+    assert "Error" in result
+    assert agent.execution_env.fig is None
+
+def test_does_fig_exist():
+    """Test that does_fig_exist correctly reports figure existence."""
+    df = pd.DataFrame({
+        'x': [1, 2, 3, 4, 5],
+        'y': [10, 20, 30, 40, 50]
+    })
+    
+    agent = PlotlyAgent()
+    agent.set_df(df)
+    
+    # Initially no figure should exist
+    assert "No figure has been created yet" in agent.does_fig_exist()
+    
+    # Create a figure
+    valid_code = """import plotly.express as px
+fig = px.scatter(df, x='x', y='y')"""
+    agent.execute_plotly_code(valid_code)
+    
+    # Now a figure should exist
+    assert "A figure is available for display" in agent.does_fig_exist()
+
+def test_reset_conversation():
+    """Test that reset_conversation clears the chat history."""
+    agent = PlotlyAgent()
+    agent.chat_history = ["message1", "message2"]
+    agent.reset_conversation()
+    assert agent.chat_history == []
+
+def test_view_generated_code():
+    """Test that view_generated_code returns the last generated code."""
+    agent = PlotlyAgent()
+    test_code = "test code"
+    agent.generated_code = test_code
+    assert agent.view_generated_code() == test_code
+
+def test_get_figure():
+    """Test that get_figure returns the current figure if it exists."""
+    df = pd.DataFrame({
+        'x': [1, 2, 3, 4, 5],
+        'y': [10, 20, 30, 40, 50]
+    })
+    
+    agent = PlotlyAgent()
+    agent.set_df(df)
+    
+    # Initially no figure should exist
+    assert agent.get_figure() is None
+    
+    # Create a figure
+    valid_code = """import plotly.express as px
+fig = px.scatter(df, x='x', y='y')"""
+    agent.execute_plotly_code(valid_code)
+    
+    # Now a figure should exist
+    assert agent.get_figure() is not None
+
+def test_process_message():
+    """Test that process_message updates chat history and handles responses."""
+    df = pd.DataFrame({
+        'x': [1, 2, 3, 4, 5],
+        'y': [10, 20, 30, 40, 50]
+    })
+    
+    agent = PlotlyAgent()
+    agent.set_df(df)
+    
+    # Test processing a message
+    response = agent.process_message("Create a scatter plot")
+    
+    # Check that chat history was updated
+    assert len(agent.chat_history) == 2  # One human message and one AI message
+    assert isinstance(agent.chat_history[0], HumanMessage)
+    assert isinstance(agent.chat_history[1], AIMessage)
+    assert agent.chat_history[0].content == "Create a scatter plot"
+
+def test_execute_plotly_code_without_df():
+    """Test that execute_plotly_code handles the case when no dataframe is set."""
+    agent = PlotlyAgent()
+    result = agent.execute_plotly_code("some code")
+    assert "Error" in result and "No dataframe has been set" in result 

plot_agent-0.1.1/README.md

@@ -1,52 +0,0 @@
-# Plot Agent
-
-An AI-powered data visualization assistant that helps users create Plotly visualizations in Python.
-
-## Installation
-
-You can install the package using pip:
-
-```bash
-pip install plot-agent
-```
-
-## Usage
-
-Here's a simple example of how to use Plot Agent:
-
-```python
-import pandas as pd
-from plot_agent import PlotlyAgent
-
-# Create a sample dataframe
-df = pd.DataFrame({
-    'x': [1, 2, 3, 4, 5],
-    'y': [10, 20, 30, 40, 50]
-})
-
-# Initialize the agent
-agent = PlotlyAgent()
-
-# Set the dataframe
-agent.set_df(df)
-
-# Process a visualization request
-response = agent.process_message("Create a line plot of x vs y")
-```
-
-## Features
-
-- AI-powered visualization generation
-- Support for various Plotly chart types
-- Automatic data preprocessing
-- Interactive visualization capabilities
-- Integration with LangChain for advanced AI capabilities
-
-## Requirements
-
-- Python 3.8 or higher
-- Dependencies are automatically installed with the package
-
-## License
-
-This project is licensed under the MIT License - see the LICENSE file for details.