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

plot_agent/agent.py

@@ -1,13 +1,6 @@
 import pandas as pd
-import numpy as np
-import matplotlib.pyplot as plt
-import plotly.express as px
-import plotly.graph_objects as go
-from plotly.subplots import make_subplots
 from io import StringIO
-import traceback
-import sys
-from typing import Dict, Optional, Any
+from typing import Optional
 
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain_core.messages import AIMessage, HumanMessage
@@ -75,6 +68,14 @@ class PlotlyAgent:
         Returns:
             None
         """
+
+        # Check df
+        assert isinstance(df, pd.DataFrame), "The dataframe must be a pandas dataframe."
+        assert not df.empty, "The dataframe must not be empty."
+
+        if sql_query:
+            assert isinstance(sql_query, str), "The SQL query must be a string."
+
         self.df = df
 
         # Capture df.info() output
@@ -104,18 +105,27 @@ class PlotlyAgent:
         Returns:
             str: The result of the execution.
         """
+        assert isinstance(generated_code, str), "The generated code must be a string."
+
         if not self.execution_env:
             return "Error: No dataframe has been set. Please set a dataframe first."
 
         # Store this as the last generated code
         self.generated_code = generated_code
 
-        result = self.execution_env.execute_code(generated_code)
+        # Execute the generated code
+        code_execution_result = self.execution_env.execute_code(generated_code)
 
-        if result["success"]:
-            return f"Code executed successfully! A figure object was created.\n{result.get('output', '')}"
+        # Extract the results from the code execution
+        code_execution_success = code_execution_result.get("success", False)
+        code_execution_output = code_execution_result.get("output", "")
+        code_execution_error = code_execution_result.get("error", "")
+
+        # Check if the code executed successfully
+        if code_execution_success:
+            return f"Code executed successfully! A figure object was created.\n{code_execution_output}"
         else:
-            return f"Error: {result.get('error', 'Unknown error')}\n{result.get('output', '')}"
+            return f"Error: {code_execution_error}\n{code_execution_output}"
 
     def does_fig_exist(self, *args, **kwargs) -> str:
         """
@@ -144,23 +154,25 @@ class PlotlyAgent:
 
     def _initialize_agent(self):
         """Initialize the LangChain agent with the necessary tools and prompt."""
+
+        # Initialize the tools
         tools = [
             Tool.from_function(
                 func=self.execute_plotly_code,
                 name="execute_plotly_code",
-                description="Execute the provided Plotly code and return the result",
+                description="Execute the provided Plotly code and return a result indicating if the code executed successfully and if a figure object was created.",
                 args_schema=GeneratedCodeInput,
             ),
             StructuredTool.from_function(
                 func=self.does_fig_exist,
                 name="does_fig_exist",
-                description="Check if a figure exists and is available for display. This tool takes no arguments.",
+                description="Check if a figure exists and is available for display. This tool takes no arguments and returns a string indicating if a figure is available for display or not.",
                 args_schema=DoesFigExistInput,
             ),
             StructuredTool.from_function(
                 func=self.view_generated_code,
                 name="view_generated_code",
-                description="View the generated code.",
+                description="View the generated code. This tool takes no arguments and returns the generated code as a string.",
                 args_schema=ViewGeneratedCodeInput,
             ),
         ]
@@ -198,6 +210,8 @@ class PlotlyAgent:
 
     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."
+
         if not self.agent_executor:
             return "Please set a dataframe first using set_df() method."
 
@@ -238,3 +252,4 @@ class PlotlyAgent:
     def reset_conversation(self):
         """Reset the conversation history."""
         self.chat_history = []
+        self.generated_code = None

plot_agent/execution.py

@@ -11,9 +11,20 @@ from typing import Dict, Any
 
 
 class PlotlyAgentExecutionEnvironment:
-    """Environment to safely execute plotly code and capture the fig object."""
+    """
+    Environment to safely execute plotly code and capture the fig object.
+
+    Args:
+        df (pd.DataFrame): The dataframe to use for the execution environment.
+    """
 
     def __init__(self, df: pd.DataFrame):
+        """
+        Initialize the execution environment with the given dataframe.
+
+        Args:
+            df (pd.DataFrame): The dataframe to use for the execution environment.
+        """
         self.df = df
         self.locals_dict = {
             "df": df,

{plot_agent-0.2.1.dist-info → plot_agent-0.2.2.dist-info}/METADATA

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

@@ -0,0 +1,10 @@
+plot_agent/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+plot_agent/agent.py,sha256=mflUj-vE_x9_W7XTJ3-GgYFfvg3uXV3wkrUw7wRPlVM,9592
+plot_agent/execution.py,sha256=BBKBVGQDrg7BPWvPbLWOjkAAFRlfyu28QxClXuspd8o,2772
+plot_agent/models.py,sha256=ZOWWeYaqmnKJarXYyXnBQQ4nwUe71ae_gMul3bZXaWU,644
+plot_agent/prompt.py,sha256=HjRgbsAe8HHs8arQogvzOGQdThEWKRqQhtQyaUplxhQ,3064
+plot_agent-0.2.2.dist-info/licenses/LICENSE,sha256=A4DPih7wHrh4VMEG3p1PhorqdhjmGIo8nQdYNQL7daA,1062
+plot_agent-0.2.2.dist-info/METADATA,sha256=dLTSBMjvxg0U63r-EbN2tfI7-eZGYexq854L0kN8M6M,2841
+plot_agent-0.2.2.dist-info/WHEEL,sha256=lTU6B6eIfYoiQJTZNc-fyaR6BpL6ehTzU3xGYxn2n8k,91
+plot_agent-0.2.2.dist-info/top_level.txt,sha256=KyOjpihUssx26Ra-37vKUQ71pI2qgJsHaRwXHJUhjzQ,11
+plot_agent-0.2.2.dist-info/RECORD,,

plot_agent-0.2.1.dist-info/RECORD

@@ -1,10 +0,0 @@
-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,,