plot-agent 0.4.0__tar.gz → 0.5.1__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {plot_agent-0.4.0 → plot_agent-0.5.1}/PKG-INFO +13 -1
- plot_agent-0.5.1/plot_agent/__init__.py +5 -0
- {plot_agent-0.4.0 → plot_agent-0.5.1}/plot_agent/agent.py +251 -5
- {plot_agent-0.4.0 → plot_agent-0.5.1}/plot_agent/execution.py +90 -6
- {plot_agent-0.4.0 → plot_agent-0.5.1}/plot_agent/models.py +6 -0
- {plot_agent-0.4.0 → plot_agent-0.5.1}/plot_agent/prompt.py +10 -5
- {plot_agent-0.4.0 → plot_agent-0.5.1}/plot_agent.egg-info/PKG-INFO +13 -1
- {plot_agent-0.4.0 → plot_agent-0.5.1}/plot_agent.egg-info/SOURCES.txt +1 -0
- plot_agent-0.5.1/plot_agent.egg-info/requires.txt +13 -0
- {plot_agent-0.4.0 → plot_agent-0.5.1}/pyproject.toml +16 -1
- plot_agent-0.4.0/plot_agent/__init__.py +0 -0
- {plot_agent-0.4.0 → plot_agent-0.5.1}/LICENSE +0 -0
- {plot_agent-0.4.0 → plot_agent-0.5.1}/README.md +0 -0
- {plot_agent-0.4.0 → plot_agent-0.5.1}/plot_agent.egg-info/dependency_links.txt +0 -0
- {plot_agent-0.4.0 → plot_agent-0.5.1}/plot_agent.egg-info/top_level.txt +0 -0
- {plot_agent-0.4.0 → plot_agent-0.5.1}/setup.cfg +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: plot-agent
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.5.1
|
|
4
4
|
Summary: An AI-powered data visualization assistant using Plotly
|
|
5
5
|
Author-email: andrewm4894 <andrewm4894@gmail.com>
|
|
6
6
|
Project-URL: Homepage, https://github.com/andrewm4894/plot-agent
|
|
@@ -10,6 +10,18 @@ Classifier: Operating System :: OS Independent
|
|
|
10
10
|
Requires-Python: >=3.8
|
|
11
11
|
Description-Content-Type: text/markdown
|
|
12
12
|
License-File: LICENSE
|
|
13
|
+
Requires-Dist: pandas
|
|
14
|
+
Requires-Dist: numpy
|
|
15
|
+
Requires-Dist: matplotlib
|
|
16
|
+
Requires-Dist: plotly
|
|
17
|
+
Requires-Dist: kaleido
|
|
18
|
+
Requires-Dist: langchain-core
|
|
19
|
+
Requires-Dist: langchain-openai
|
|
20
|
+
Requires-Dist: langgraph
|
|
21
|
+
Requires-Dist: pydantic
|
|
22
|
+
Requires-Dist: python-dotenv
|
|
23
|
+
Provides-Extra: posthog
|
|
24
|
+
Requires-Dist: posthog; extra == "posthog"
|
|
13
25
|
Dynamic: license-file
|
|
14
26
|
|
|
15
27
|
# Plot Agent
|
|
@@ -7,10 +7,10 @@ from io import StringIO
|
|
|
7
7
|
import os
|
|
8
8
|
import re
|
|
9
9
|
import logging
|
|
10
|
-
from typing import Optional
|
|
10
|
+
from typing import List, Optional, Union
|
|
11
11
|
from dotenv import load_dotenv
|
|
12
12
|
|
|
13
|
-
from langchain_core.messages import AIMessage, HumanMessage
|
|
13
|
+
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
|
|
14
14
|
from langchain_core.tools import Tool, StructuredTool
|
|
15
15
|
from langgraph.prebuilt import create_react_agent
|
|
16
16
|
from langchain_openai import ChatOpenAI
|
|
@@ -20,9 +20,18 @@ from plot_agent.models import (
|
|
|
20
20
|
GeneratedCodeInput,
|
|
21
21
|
DoesFigExistInput,
|
|
22
22
|
ViewGeneratedCodeInput,
|
|
23
|
+
CheckPlotOutputsInput,
|
|
23
24
|
)
|
|
24
25
|
from plot_agent.execution import PlotAgentExecutionEnvironment
|
|
25
26
|
|
|
27
|
+
# Optional PostHog integration
|
|
28
|
+
try:
|
|
29
|
+
from posthog import Posthog
|
|
30
|
+
from posthog.ai.langchain import CallbackHandler as PostHogCallbackHandler
|
|
31
|
+
POSTHOG_AVAILABLE = True
|
|
32
|
+
except ImportError:
|
|
33
|
+
POSTHOG_AVAILABLE = False
|
|
34
|
+
|
|
26
35
|
|
|
27
36
|
class PlotAgent:
|
|
28
37
|
"""
|
|
@@ -41,6 +50,7 @@ class PlotAgent:
|
|
|
41
50
|
llm_timeout: int = 60,
|
|
42
51
|
llm_max_retries: int = 1,
|
|
43
52
|
debug: bool = False,
|
|
53
|
+
include_plot_image: bool = False,
|
|
44
54
|
):
|
|
45
55
|
"""
|
|
46
56
|
Initialize the PlotAgent.
|
|
@@ -52,6 +62,11 @@ class PlotAgent:
|
|
|
52
62
|
max_iterations (int): Maximum number of iterations for the agent to take.
|
|
53
63
|
early_stopping_method (str): Method to use for early stopping.
|
|
54
64
|
handle_parsing_errors (bool): Whether to handle parsing errors gracefully.
|
|
65
|
+
llm_temperature (float): Temperature for LLM sampling.
|
|
66
|
+
llm_timeout (int): Timeout in seconds for LLM calls.
|
|
67
|
+
llm_max_retries (int): Maximum retries for LLM calls.
|
|
68
|
+
debug (bool): Enable debug logging.
|
|
69
|
+
include_plot_image (bool): Generate PNG image of plots for PostHog analytics.
|
|
55
70
|
"""
|
|
56
71
|
# Load .env if present, then require a valid API key
|
|
57
72
|
load_dotenv()
|
|
@@ -76,6 +91,73 @@ class PlotAgent:
|
|
|
76
91
|
)
|
|
77
92
|
self._logger.addHandler(handler)
|
|
78
93
|
|
|
94
|
+
# Initialize PostHog for LLM analytics (optional)
|
|
95
|
+
self.posthog_client = None
|
|
96
|
+
self.posthog_callback_handler = None
|
|
97
|
+
posthog_enabled = os.getenv("POSTHOG_ENABLED", "false").lower() == "true"
|
|
98
|
+
|
|
99
|
+
# Enable PostHog multimodal capture if include_plot_image is True
|
|
100
|
+
if include_plot_image and posthog_enabled:
|
|
101
|
+
os.environ["_INTERNAL_LLMA_MULTIMODAL"] = "true"
|
|
102
|
+
|
|
103
|
+
if posthog_enabled:
|
|
104
|
+
if not POSTHOG_AVAILABLE:
|
|
105
|
+
self._logger.warning(
|
|
106
|
+
"PostHog is enabled but the posthog package is not installed. "
|
|
107
|
+
"Install it with: pip install posthog"
|
|
108
|
+
)
|
|
109
|
+
else:
|
|
110
|
+
posthog_api_key = os.getenv("POSTHOG_API_KEY")
|
|
111
|
+
posthog_host = os.getenv("POSTHOG_HOST", "https://app.posthog.com")
|
|
112
|
+
|
|
113
|
+
if not posthog_api_key:
|
|
114
|
+
self._logger.warning(
|
|
115
|
+
"POSTHOG_ENABLED is true but POSTHOG_API_KEY is not set. "
|
|
116
|
+
"PostHog tracking will be disabled."
|
|
117
|
+
)
|
|
118
|
+
else:
|
|
119
|
+
try:
|
|
120
|
+
# Build super_properties for session tracking
|
|
121
|
+
super_properties = {}
|
|
122
|
+
|
|
123
|
+
# Add session ID from environment if provided
|
|
124
|
+
ai_session_id = os.getenv("POSTHOG_AI_SESSION_ID")
|
|
125
|
+
if ai_session_id:
|
|
126
|
+
super_properties["$ai_session_id"] = ai_session_id
|
|
127
|
+
|
|
128
|
+
# Initialize PostHog client with super_properties
|
|
129
|
+
self.posthog_client = Posthog(
|
|
130
|
+
posthog_api_key,
|
|
131
|
+
host=posthog_host,
|
|
132
|
+
super_properties=super_properties
|
|
133
|
+
)
|
|
134
|
+
|
|
135
|
+
# Build callback handler config
|
|
136
|
+
callback_config = {"client": self.posthog_client}
|
|
137
|
+
|
|
138
|
+
# Add optional distinct_id
|
|
139
|
+
distinct_id = os.getenv("POSTHOG_DISTINCT_ID")
|
|
140
|
+
if distinct_id:
|
|
141
|
+
callback_config["distinct_id"] = distinct_id
|
|
142
|
+
|
|
143
|
+
# Add privacy mode setting
|
|
144
|
+
privacy_mode = os.getenv("POSTHOG_PRIVACY_MODE", "false").lower() == "true"
|
|
145
|
+
callback_config["privacy_mode"] = privacy_mode
|
|
146
|
+
|
|
147
|
+
self.posthog_callback_handler = PostHogCallbackHandler(**callback_config)
|
|
148
|
+
|
|
149
|
+
if self.debug:
|
|
150
|
+
session_info = f"session_id={ai_session_id}" if ai_session_id else "no session"
|
|
151
|
+
self._logger.debug(
|
|
152
|
+
f"PostHog LLM analytics initialized (host={posthog_host}, "
|
|
153
|
+
f"distinct_id={distinct_id or 'anonymous'}, "
|
|
154
|
+
f"privacy_mode={privacy_mode}, {session_info})"
|
|
155
|
+
)
|
|
156
|
+
except Exception as e:
|
|
157
|
+
self._logger.error(f"Failed to initialize PostHog: {e}")
|
|
158
|
+
self.posthog_client = None
|
|
159
|
+
self.posthog_callback_handler = None
|
|
160
|
+
|
|
79
161
|
self.llm = ChatOpenAI(
|
|
80
162
|
model=model,
|
|
81
163
|
temperature=llm_temperature,
|
|
@@ -97,6 +179,7 @@ class PlotAgent:
|
|
|
97
179
|
self.max_iterations = max_iterations
|
|
98
180
|
self.early_stopping_method = early_stopping_method
|
|
99
181
|
self.handle_parsing_errors = handle_parsing_errors
|
|
182
|
+
self.include_plot_image = include_plot_image
|
|
100
183
|
|
|
101
184
|
def set_df(self, df: pd.DataFrame, sql_query: Optional[str] = None):
|
|
102
185
|
"""
|
|
@@ -131,7 +214,9 @@ class PlotAgent:
|
|
|
131
214
|
self.sql_query = sql_query
|
|
132
215
|
|
|
133
216
|
# Initialize execution environment
|
|
134
|
-
self.execution_env = PlotAgentExecutionEnvironment(
|
|
217
|
+
self.execution_env = PlotAgentExecutionEnvironment(
|
|
218
|
+
df, include_plot_image=self.include_plot_image
|
|
219
|
+
)
|
|
135
220
|
|
|
136
221
|
# Initialize the agent with tools
|
|
137
222
|
self._initialize_agent()
|
|
@@ -191,6 +276,44 @@ class PlotAgent:
|
|
|
191
276
|
else:
|
|
192
277
|
return "No figure has been created yet."
|
|
193
278
|
|
|
279
|
+
def check_plot_outputs(self, *args, **kwargs) -> str:
|
|
280
|
+
"""
|
|
281
|
+
Check if all required plot outputs (fig, plot_title, plot_summary) are available.
|
|
282
|
+
|
|
283
|
+
Args:
|
|
284
|
+
*args: Any positional arguments (ignored)
|
|
285
|
+
**kwargs: Any keyword arguments (ignored)
|
|
286
|
+
|
|
287
|
+
Returns:
|
|
288
|
+
str: A message indicating which plot outputs are available.
|
|
289
|
+
"""
|
|
290
|
+
if not self.execution_env:
|
|
291
|
+
return "No execution environment has been initialized. Please set a dataframe first."
|
|
292
|
+
|
|
293
|
+
available = []
|
|
294
|
+
missing = []
|
|
295
|
+
|
|
296
|
+
if self.execution_env.fig is not None:
|
|
297
|
+
available.append("fig")
|
|
298
|
+
else:
|
|
299
|
+
missing.append("fig")
|
|
300
|
+
|
|
301
|
+
if self.execution_env.plot_title is not None:
|
|
302
|
+
available.append("plot_title")
|
|
303
|
+
else:
|
|
304
|
+
missing.append("plot_title")
|
|
305
|
+
|
|
306
|
+
if self.execution_env.plot_summary is not None:
|
|
307
|
+
available.append("plot_summary")
|
|
308
|
+
else:
|
|
309
|
+
missing.append("plot_summary")
|
|
310
|
+
|
|
311
|
+
if not missing:
|
|
312
|
+
return "All required plot outputs are available: fig, plot_title, and plot_summary."
|
|
313
|
+
else:
|
|
314
|
+
status = f"Available: {', '.join(available) if available else 'none'}. Missing: {', '.join(missing)}."
|
|
315
|
+
return status
|
|
316
|
+
|
|
194
317
|
def view_generated_code(self, *args, **kwargs) -> str:
|
|
195
318
|
"""
|
|
196
319
|
View the generated code.
|
|
@@ -230,6 +353,15 @@ class PlotAgent:
|
|
|
230
353
|
),
|
|
231
354
|
args_schema=ViewGeneratedCodeInput,
|
|
232
355
|
),
|
|
356
|
+
StructuredTool.from_function(
|
|
357
|
+
func=self.check_plot_outputs,
|
|
358
|
+
name="check_plot_outputs",
|
|
359
|
+
description=(
|
|
360
|
+
"Check if all required plot outputs (fig, plot_title, plot_summary) are available. "
|
|
361
|
+
"This tool takes no arguments and returns the status of all plot outputs."
|
|
362
|
+
),
|
|
363
|
+
args_schema=CheckPlotOutputsInput,
|
|
364
|
+
),
|
|
233
365
|
]
|
|
234
366
|
|
|
235
367
|
# Prepare system prompt with dataframe information
|
|
@@ -300,10 +432,16 @@ class PlotAgent:
|
|
|
300
432
|
if self.debug:
|
|
301
433
|
self._logger.debug(f"process_message() user: {user_message}")
|
|
302
434
|
self._logger.debug(f"graph message count before invoke: {len(graph_messages)}")
|
|
435
|
+
|
|
436
|
+
# Build config with optional PostHog callback
|
|
437
|
+
invoke_config = {"recursion_limit": self.max_iterations}
|
|
438
|
+
if self.posthog_callback_handler:
|
|
439
|
+
invoke_config["callbacks"] = [self.posthog_callback_handler]
|
|
440
|
+
|
|
303
441
|
# Invoke the LangGraph agent
|
|
304
442
|
result = self.agent_executor.invoke(
|
|
305
443
|
{"messages": graph_messages},
|
|
306
|
-
config=
|
|
444
|
+
config=invoke_config,
|
|
307
445
|
)
|
|
308
446
|
|
|
309
447
|
# Extract the latest AI message from the returned messages
|
|
@@ -362,9 +500,14 @@ class PlotAgent:
|
|
|
362
500
|
)
|
|
363
501
|
),
|
|
364
502
|
]
|
|
503
|
+
# Build config with optional PostHog callback for retry
|
|
504
|
+
retry_config = {"recursion_limit": max(3, self.max_iterations // 2)}
|
|
505
|
+
if self.posthog_callback_handler:
|
|
506
|
+
retry_config["callbacks"] = [self.posthog_callback_handler]
|
|
507
|
+
|
|
365
508
|
retry_result = self.agent_executor.invoke(
|
|
366
509
|
{"messages": guided_messages},
|
|
367
|
-
config=
|
|
510
|
+
config=retry_config,
|
|
368
511
|
)
|
|
369
512
|
self._graph_messages = retry_result.get("messages", [])
|
|
370
513
|
retry_ai_messages = [
|
|
@@ -393,6 +536,16 @@ class PlotAgent:
|
|
|
393
536
|
if self.debug:
|
|
394
537
|
self._logger.debug(f"execution result success={exec_result.get('success')} error={exec_result.get('error')!r}")
|
|
395
538
|
|
|
539
|
+
# Run verification step with image if enabled and figure exists
|
|
540
|
+
# This sends the plot image to the LLM for verification, which gets captured by PostHog
|
|
541
|
+
if (
|
|
542
|
+
self.include_plot_image
|
|
543
|
+
and self.posthog_callback_handler
|
|
544
|
+
and self.execution_env
|
|
545
|
+
and self.execution_env.fig is not None
|
|
546
|
+
):
|
|
547
|
+
self._verify_plot_with_image(user_message)
|
|
548
|
+
|
|
396
549
|
return ai_content if isinstance(ai_content, str) else str(ai_content)
|
|
397
550
|
|
|
398
551
|
def get_figure(self):
|
|
@@ -401,7 +554,100 @@ class PlotAgent:
|
|
|
401
554
|
return self.execution_env.fig
|
|
402
555
|
return None
|
|
403
556
|
|
|
557
|
+
def get_plot_title(self):
|
|
558
|
+
"""Return the current plot title if one exists."""
|
|
559
|
+
if self.execution_env and self.execution_env.plot_title:
|
|
560
|
+
return self.execution_env.plot_title
|
|
561
|
+
return None
|
|
562
|
+
|
|
563
|
+
def get_plot_summary(self):
|
|
564
|
+
"""Return the current plot summary if one exists."""
|
|
565
|
+
if self.execution_env and self.execution_env.plot_summary:
|
|
566
|
+
return self.execution_env.plot_summary
|
|
567
|
+
return None
|
|
568
|
+
|
|
569
|
+
def get_plot_image_base64(self) -> Optional[str]:
|
|
570
|
+
"""Return the current plot image as base64-encoded data URI."""
|
|
571
|
+
if self.execution_env and self.execution_env.plot_image_base64:
|
|
572
|
+
return self.execution_env.plot_image_base64
|
|
573
|
+
return None
|
|
574
|
+
|
|
575
|
+
def _verify_plot_with_image(self, user_request: str) -> Optional[str]:
|
|
576
|
+
"""
|
|
577
|
+
Send the generated plot image to the LLM for verification.
|
|
578
|
+
|
|
579
|
+
This step serves two purposes:
|
|
580
|
+
1. Verifies the plot matches the user's request
|
|
581
|
+
2. Captures the plot image in PostHog LLM traces (via multimodal message)
|
|
582
|
+
|
|
583
|
+
Args:
|
|
584
|
+
user_request: The original user request for context.
|
|
585
|
+
|
|
586
|
+
Returns:
|
|
587
|
+
The LLM's verification response, or None if verification fails.
|
|
588
|
+
"""
|
|
589
|
+
plot_image = self.get_plot_image_base64()
|
|
590
|
+
if not plot_image:
|
|
591
|
+
return None
|
|
592
|
+
|
|
593
|
+
# Build multimodal message with the plot image
|
|
594
|
+
human_content: List[Union[dict, str]] = [
|
|
595
|
+
{
|
|
596
|
+
"type": "text",
|
|
597
|
+
"text": (
|
|
598
|
+
f"Please verify this generated plot matches the user's request.\n\n"
|
|
599
|
+
f"User request: {user_request}\n\n"
|
|
600
|
+
f"Generated code:\n```python\n{self.generated_code or 'N/A'}\n```\n\n"
|
|
601
|
+
f"Plot title: {self.get_plot_title() or 'N/A'}\n"
|
|
602
|
+
f"Plot summary: {self.get_plot_summary() or 'N/A'}\n\n"
|
|
603
|
+
f"Respond with a brief confirmation that the plot is correct, "
|
|
604
|
+
f"or note any issues you see."
|
|
605
|
+
),
|
|
606
|
+
},
|
|
607
|
+
{
|
|
608
|
+
"type": "image_url",
|
|
609
|
+
"image_url": {"url": plot_image},
|
|
610
|
+
},
|
|
611
|
+
]
|
|
612
|
+
|
|
613
|
+
messages = [
|
|
614
|
+
SystemMessage(
|
|
615
|
+
content=(
|
|
616
|
+
"You are a plot verification assistant. "
|
|
617
|
+
"Review the generated plot image and verify it matches the user's request. "
|
|
618
|
+
"Be concise in your response."
|
|
619
|
+
)
|
|
620
|
+
),
|
|
621
|
+
HumanMessage(content=human_content),
|
|
622
|
+
]
|
|
623
|
+
|
|
624
|
+
try:
|
|
625
|
+
# Build config with PostHog callback to capture this verification step
|
|
626
|
+
invoke_config = {}
|
|
627
|
+
if self.posthog_callback_handler:
|
|
628
|
+
invoke_config["callbacks"] = [self.posthog_callback_handler]
|
|
629
|
+
|
|
630
|
+
if self.debug:
|
|
631
|
+
self._logger.debug("Running plot verification with image")
|
|
632
|
+
|
|
633
|
+
# Call LLM directly (not through agent) for verification
|
|
634
|
+
response = self.llm.invoke(messages, config=invoke_config)
|
|
635
|
+
verification_result = response.content if hasattr(response, "content") else str(response)
|
|
636
|
+
|
|
637
|
+
if self.debug:
|
|
638
|
+
self._logger.debug(f"Plot verification result: {verification_result[:200]}...")
|
|
639
|
+
|
|
640
|
+
return verification_result
|
|
641
|
+
except Exception as e:
|
|
642
|
+
self._logger.warning(f"Plot verification failed: {e}")
|
|
643
|
+
return None
|
|
644
|
+
|
|
404
645
|
def reset_conversation(self):
|
|
405
646
|
"""Reset the conversation history."""
|
|
406
647
|
self.chat_history = []
|
|
407
648
|
self.generated_code = None
|
|
649
|
+
if self.execution_env:
|
|
650
|
+
self.execution_env.fig = None
|
|
651
|
+
self.execution_env.plot_title = None
|
|
652
|
+
self.execution_env.plot_summary = None
|
|
653
|
+
self.execution_env.plot_image_base64 = None
|
|
@@ -9,12 +9,15 @@ Security features:
|
|
|
9
9
|
• Enforce a 60 second timeout via signal.alarm
|
|
10
10
|
"""
|
|
11
11
|
import ast
|
|
12
|
+
import base64
|
|
12
13
|
import builtins
|
|
14
|
+
import logging
|
|
13
15
|
import signal
|
|
14
16
|
import threading
|
|
15
17
|
import traceback
|
|
16
18
|
from io import StringIO
|
|
17
19
|
import contextlib
|
|
20
|
+
from typing import Optional
|
|
18
21
|
|
|
19
22
|
import pandas as pd
|
|
20
23
|
import numpy as np
|
|
@@ -111,11 +114,17 @@ class PlotAgentExecutionEnvironment:
|
|
|
111
114
|
"__import__": _safe_import,
|
|
112
115
|
}
|
|
113
116
|
|
|
114
|
-
def __init__(self, df: pd.DataFrame):
|
|
117
|
+
def __init__(self, df: pd.DataFrame, include_plot_image: bool = False):
|
|
115
118
|
"""
|
|
116
119
|
Initialize the execution environment with a dataframe.
|
|
120
|
+
|
|
121
|
+
Args:
|
|
122
|
+
df: The pandas dataframe to use for plotting.
|
|
123
|
+
include_plot_image: If True, generate a PNG image of the plot after execution.
|
|
117
124
|
"""
|
|
118
125
|
self.df = df
|
|
126
|
+
self.include_plot_image = include_plot_image
|
|
127
|
+
self._logger = logging.getLogger("plot_agent.execution")
|
|
119
128
|
# Base namespace for both globals & locals
|
|
120
129
|
self._base_ns = {
|
|
121
130
|
"__builtins__": self._SAFE_BUILTINS,
|
|
@@ -128,6 +137,30 @@ class PlotAgentExecutionEnvironment:
|
|
|
128
137
|
"make_subplots": make_subplots,
|
|
129
138
|
}
|
|
130
139
|
self.fig = None
|
|
140
|
+
self.plot_title = None
|
|
141
|
+
self.plot_summary = None
|
|
142
|
+
self.plot_image_base64 = None
|
|
143
|
+
|
|
144
|
+
def _generate_plot_png(self, fig, width: int = 800, height: int = 600) -> Optional[str]:
|
|
145
|
+
"""
|
|
146
|
+
Generate PNG as base64 data URI from a Plotly figure.
|
|
147
|
+
|
|
148
|
+
Args:
|
|
149
|
+
fig: The Plotly figure to convert.
|
|
150
|
+
width: Image width in pixels.
|
|
151
|
+
height: Image height in pixels.
|
|
152
|
+
|
|
153
|
+
Returns:
|
|
154
|
+
Base64-encoded data URI string, or None if generation fails.
|
|
155
|
+
"""
|
|
156
|
+
try:
|
|
157
|
+
import plotly.io as pio
|
|
158
|
+
img_bytes = pio.to_image(fig, format='png', width=width, height=height)
|
|
159
|
+
b64 = base64.b64encode(img_bytes).decode('utf-8')
|
|
160
|
+
return f"data:image/png;base64,{b64}"
|
|
161
|
+
except Exception as e:
|
|
162
|
+
self._logger.warning(f"Failed to generate plot PNG: {e}")
|
|
163
|
+
return None
|
|
131
164
|
|
|
132
165
|
def _validate_ast(self, node: ast.AST):
|
|
133
166
|
"""
|
|
@@ -167,8 +200,10 @@ class PlotAgentExecutionEnvironment:
|
|
|
167
200
|
|
|
168
201
|
# Copy the base namespace
|
|
169
202
|
ns = self._base_ns.copy()
|
|
170
|
-
# Purge any old
|
|
203
|
+
# Purge any old variables
|
|
171
204
|
ns.pop("fig", None)
|
|
205
|
+
ns.pop("plot_title", None)
|
|
206
|
+
ns.pop("plot_summary", None)
|
|
172
207
|
|
|
173
208
|
try:
|
|
174
209
|
# Parse the generated code
|
|
@@ -179,6 +214,8 @@ class PlotAgentExecutionEnvironment:
|
|
|
179
214
|
# If the code is rejected on safety grounds, return an error
|
|
180
215
|
return {
|
|
181
216
|
"fig": None,
|
|
217
|
+
"plot_title": None,
|
|
218
|
+
"plot_summary": None,
|
|
182
219
|
"output": "",
|
|
183
220
|
"error": f"Code rejected on safety grounds: {e}",
|
|
184
221
|
"success": False,
|
|
@@ -209,6 +246,8 @@ class PlotAgentExecutionEnvironment:
|
|
|
209
246
|
tb = traceback.format_exc()
|
|
210
247
|
return {
|
|
211
248
|
"fig": None,
|
|
249
|
+
"plot_title": None,
|
|
250
|
+
"plot_summary": None,
|
|
212
251
|
"output": out_buf.getvalue(),
|
|
213
252
|
"error": f"Code execution timed out: {te}\n{tb}",
|
|
214
253
|
"success": False,
|
|
@@ -218,6 +257,8 @@ class PlotAgentExecutionEnvironment:
|
|
|
218
257
|
tb = traceback.format_exc()
|
|
219
258
|
return {
|
|
220
259
|
"fig": None,
|
|
260
|
+
"plot_title": None,
|
|
261
|
+
"plot_summary": None,
|
|
221
262
|
"output": out_buf.getvalue(),
|
|
222
263
|
"error": f"Error executing code: {e}\n{tb}",
|
|
223
264
|
"success": False,
|
|
@@ -230,21 +271,64 @@ class PlotAgentExecutionEnvironment:
|
|
|
230
271
|
except Exception:
|
|
231
272
|
pass
|
|
232
273
|
|
|
233
|
-
# Get the
|
|
274
|
+
# Get the variables
|
|
234
275
|
fig = ns.get("fig")
|
|
276
|
+
plot_title = ns.get("plot_title")
|
|
277
|
+
plot_summary = ns.get("plot_summary")
|
|
278
|
+
|
|
279
|
+
# Store the variables
|
|
235
280
|
self.fig = fig
|
|
281
|
+
self.plot_title = plot_title
|
|
282
|
+
self.plot_summary = plot_summary
|
|
283
|
+
|
|
284
|
+
# Generate PNG if enabled and figure exists
|
|
285
|
+
self.plot_image_base64 = None
|
|
286
|
+
if self.include_plot_image and fig is not None:
|
|
287
|
+
self.plot_image_base64 = self._generate_plot_png(fig)
|
|
288
|
+
|
|
289
|
+
# Validate required variables
|
|
290
|
+
missing_vars = []
|
|
236
291
|
if fig is None:
|
|
292
|
+
missing_vars.append("fig")
|
|
293
|
+
if plot_title is None:
|
|
294
|
+
missing_vars.append("plot_title")
|
|
295
|
+
if plot_summary is None:
|
|
296
|
+
missing_vars.append("plot_summary")
|
|
297
|
+
|
|
298
|
+
if missing_vars:
|
|
237
299
|
return {
|
|
238
|
-
"fig":
|
|
300
|
+
"fig": fig,
|
|
301
|
+
"plot_title": plot_title,
|
|
302
|
+
"plot_summary": plot_summary,
|
|
303
|
+
"output": out_buf.getvalue(),
|
|
304
|
+
"error": f"Missing required variables: {', '.join(missing_vars)}. Please create variables named: {', '.join(missing_vars)}.",
|
|
305
|
+
"success": False,
|
|
306
|
+
}
|
|
307
|
+
|
|
308
|
+
# Validate that plot_title and plot_summary are strings
|
|
309
|
+
validation_errors = []
|
|
310
|
+
if not isinstance(plot_title, str):
|
|
311
|
+
validation_errors.append("plot_title must be a string")
|
|
312
|
+
if not isinstance(plot_summary, str):
|
|
313
|
+
validation_errors.append("plot_summary must be a string")
|
|
314
|
+
|
|
315
|
+
if validation_errors:
|
|
316
|
+
return {
|
|
317
|
+
"fig": fig,
|
|
318
|
+
"plot_title": plot_title,
|
|
319
|
+
"plot_summary": plot_summary,
|
|
239
320
|
"output": out_buf.getvalue(),
|
|
240
|
-
"error": "
|
|
321
|
+
"error": f"Validation errors: {'; '.join(validation_errors)}.",
|
|
241
322
|
"success": False,
|
|
242
323
|
}
|
|
243
324
|
|
|
244
325
|
# Return the result
|
|
245
326
|
return {
|
|
246
327
|
"fig": fig,
|
|
247
|
-
"
|
|
328
|
+
"plot_title": plot_title,
|
|
329
|
+
"plot_summary": plot_summary,
|
|
330
|
+
"plot_image_base64": self.plot_image_base64,
|
|
331
|
+
"output": "Code executed successfully. 'fig', 'plot_title', and 'plot_summary' objects were created.",
|
|
248
332
|
"error": "",
|
|
249
333
|
"success": True,
|
|
250
334
|
}
|
|
@@ -31,3 +31,9 @@ class ViewGeneratedCodeInput(BaseModel):
|
|
|
31
31
|
"""Model indicating that the view_generated_code function takes no arguments."""
|
|
32
32
|
|
|
33
33
|
pass
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
class CheckPlotOutputsInput(BaseModel):
|
|
37
|
+
"""Model indicating that the check_plot_outputs function takes no arguments."""
|
|
38
|
+
|
|
39
|
+
pass
|
|
@@ -26,10 +26,12 @@ NOTES:
|
|
|
26
26
|
- You must paste the full code, not just a reference to the code.
|
|
27
27
|
- You must not use fig.show() in your code as it will ultimately be executed elsewhere in a headless environment.
|
|
28
28
|
- 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.
|
|
29
|
+
- Your code MUST create three variables: 'fig' (Plotly figure), 'plot_title' (string), and 'plot_summary' (string).
|
|
29
30
|
|
|
30
31
|
TOOLS:
|
|
31
32
|
- execute_plotly_code(generated_code) to execute the generated code.
|
|
32
33
|
- does_fig_exist() to check that a fig object is available for display. This tool takes no arguments.
|
|
34
|
+
- check_plot_outputs() to check if all required outputs (fig, plot_title, plot_summary) are available. This tool takes no arguments.
|
|
33
35
|
- view_generated_code() to view the generated code if need to fix it. This tool takes no arguments.
|
|
34
36
|
|
|
35
37
|
IMPORTANT CODE FORMATTING INSTRUCTIONS:
|
|
@@ -37,6 +39,9 @@ IMPORTANT CODE FORMATTING INSTRUCTIONS:
|
|
|
37
39
|
2. Use descriptive variable names.
|
|
38
40
|
3. DO NOT include fig.show() in your code - the visualization will be rendered externally.
|
|
39
41
|
4. Ensure your code creates a variable named 'fig' that contains the Plotly figure object.
|
|
42
|
+
5. You MUST also create two string variables:
|
|
43
|
+
- 'plot_title': A concise, descriptive title for the plot (string)
|
|
44
|
+
- 'plot_summary': A brief summary explaining what the plot shows and any key insights (string)
|
|
40
45
|
|
|
41
46
|
When a user asks for a visualization:
|
|
42
47
|
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
|
|
|
48
53
|
YOU MUST CALL execute_plotly_code WITH THE FULL CODE, NOT JUST A REFERENCE TO THE CODE.
|
|
49
54
|
|
|
50
55
|
YOUR WORKFLOW MUST BE:
|
|
51
|
-
1. execute_plotly_code(generated_code) to make sure the code is ran and a figure object
|
|
52
|
-
2.
|
|
53
|
-
3. if there are errors, view the generated code using view_generated_code() to see what went wrong.
|
|
54
|
-
4. fix the code and execute it again with execute_plotly_code(generated_code) to make sure
|
|
55
|
-
5. repeat until
|
|
56
|
+
1. execute_plotly_code(generated_code) to make sure the code is ran and a figure object, plot_title, and plot_summary are created.
|
|
57
|
+
2. use check_plot_outputs() to verify that all required outputs (fig, plot_title, plot_summary) are available.
|
|
58
|
+
3. if there are errors or missing outputs, view the generated code using view_generated_code() to see what went wrong.
|
|
59
|
+
4. fix the code and execute it again with execute_plotly_code(generated_code) to make sure all required outputs are created.
|
|
60
|
+
5. repeat until all outputs (figure object, plot_title, and plot_summary) are available.
|
|
56
61
|
|
|
57
62
|
Always return the final working code (with all the comments) to the user along with an explanation of what the visualization shows.
|
|
58
63
|
Make sure to follow best practices for data visualization, such as appropriate chart types, labels, and colors.
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: plot-agent
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.5.1
|
|
4
4
|
Summary: An AI-powered data visualization assistant using Plotly
|
|
5
5
|
Author-email: andrewm4894 <andrewm4894@gmail.com>
|
|
6
6
|
Project-URL: Homepage, https://github.com/andrewm4894/plot-agent
|
|
@@ -10,6 +10,18 @@ Classifier: Operating System :: OS Independent
|
|
|
10
10
|
Requires-Python: >=3.8
|
|
11
11
|
Description-Content-Type: text/markdown
|
|
12
12
|
License-File: LICENSE
|
|
13
|
+
Requires-Dist: pandas
|
|
14
|
+
Requires-Dist: numpy
|
|
15
|
+
Requires-Dist: matplotlib
|
|
16
|
+
Requires-Dist: plotly
|
|
17
|
+
Requires-Dist: kaleido
|
|
18
|
+
Requires-Dist: langchain-core
|
|
19
|
+
Requires-Dist: langchain-openai
|
|
20
|
+
Requires-Dist: langgraph
|
|
21
|
+
Requires-Dist: pydantic
|
|
22
|
+
Requires-Dist: python-dotenv
|
|
23
|
+
Provides-Extra: posthog
|
|
24
|
+
Requires-Dist: posthog; extra == "posthog"
|
|
13
25
|
Dynamic: license-file
|
|
14
26
|
|
|
15
27
|
# Plot Agent
|
|
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
|
|
4
4
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "plot-agent"
|
|
7
|
-
version = "0.
|
|
7
|
+
version = "0.5.1"
|
|
8
8
|
authors = [
|
|
9
9
|
{ name="andrewm4894", email="andrewm4894@gmail.com" },
|
|
10
10
|
]
|
|
@@ -16,6 +16,21 @@ classifiers = [
|
|
|
16
16
|
"License :: OSI Approved :: MIT License",
|
|
17
17
|
"Operating System :: OS Independent",
|
|
18
18
|
]
|
|
19
|
+
dependencies = [
|
|
20
|
+
"pandas",
|
|
21
|
+
"numpy",
|
|
22
|
+
"matplotlib",
|
|
23
|
+
"plotly",
|
|
24
|
+
"kaleido",
|
|
25
|
+
"langchain-core",
|
|
26
|
+
"langchain-openai",
|
|
27
|
+
"langgraph",
|
|
28
|
+
"pydantic",
|
|
29
|
+
"python-dotenv",
|
|
30
|
+
]
|
|
31
|
+
|
|
32
|
+
[project.optional-dependencies]
|
|
33
|
+
posthog = ["posthog"]
|
|
19
34
|
|
|
20
35
|
[project.urls]
|
|
21
36
|
"Homepage" = "https://github.com/andrewm4894/plot-agent"
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|