#16 Update Data Visualization Example

business-science · Jan 10, 2025 · 159daab · 159daab
1 parent 0c0f5fd
commit 159daab
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 5 deletions.
diff --git a/ai_data_science_team/agents/__init__.py b/ai_data_science_team/agents/__init__.py
@@ -1,6 +1,6 @@
 from ai_data_science_team.agents.data_cleaning_agent import make_data_cleaning_agent, DataCleaningAgent
-from ai_data_science_team.agents.feature_engineering_agent import make_feature_engineering_agent
-from ai_data_science_team.agents.data_wrangling_agent import make_data_wrangling_agent
-from ai_data_science_team.agents.sql_database_agent import make_sql_database_agent
-from ai_data_science_team.agents.data_visualization_agent import make_data_visualization_agent
+from ai_data_science_team.agents.feature_engineering_agent import make_feature_engineering_agent, FeatureEngineeringAgent
+from ai_data_science_team.agents.data_wrangling_agent import make_data_wrangling_agent, DataWranglingAgent
+from ai_data_science_team.agents.sql_database_agent import make_sql_database_agent, SQLDatabaseAgent
+from ai_data_science_team.agents.data_visualization_agent import make_data_visualization_agent, DataVisualizationAgent
 
diff --git a/ai_data_science_team/agents/data_visualization_agent.py b/ai_data_science_team/agents/data_visualization_agent.py
@@ -33,6 +33,7 @@
 from ai_data_science_team.tools.regex import relocate_imports_inside_function, add_comments_to_top, format_agent_name
 from ai_data_science_team.tools.metadata import get_dataframe_summary
 from ai_data_science_team.tools.logging import log_ai_function
+from ai_data_science_team.utils.plotly import plotly_from_dict
 
 # Setup
 AGENT_NAME = "data_visualization_agent"
@@ -288,7 +289,7 @@ def get_plotly_graph(self):
             The Plotly graph dictionary if available, otherwise None.
         """
         if self.response:
-            return self.response.get("plotly_graph", None)
+            return plotly_from_dict(self.response.get("plotly_graph", None))
         return None
 
     def get_data_raw(self):

diff --git a/ai_data_science_team/utils/plotly.py b/ai_data_science_team/utils/plotly.py
@@ -18,4 +18,7 @@ def plotly_from_dict(plotly_graph_dict: dict):
         A Plotly graph object.
     """
 
+    if plotly_from_dict is None:
+        return None
+
     return pio.from_json(json.dumps(plotly_graph_dict))