plotly images

vanna-ai · Jul 28, 2023 · 7e0460b · 7e0460b
1 parent 9091ec8
commit 7e0460b
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 2,839 deletions.
diff --git a/notebooks/vn-ask.ipynb b/notebooks/vn-ask.ipynb
diff --git a/pyproject.toml b/pyproject.toml
@@ -13,7 +13,7 @@ classifiers = [
     "Operating System :: OS Independent",
 ]
 dependencies = [
-    "requests", "tabulate", "plotly", "pandas"
+    "requests", "tabulate", "plotly", "kaleido", "pandas"
 ]
 
 [project.urls]

diff --git a/src/vanna/__init__.py b/src/vanna/__init__.py
@@ -982,14 +982,21 @@ def ask(question: Union[str, None] = None, print_results: bool = True, auto_trai
             plotly_code = generate_plotly_code(question=question, sql=sql, df=df)
             fig = get_plotly_figure(plotly_code=plotly_code, df=df)
             if print_results:
-                fig.show()
+                try:
+                    display = __import__('IPython.display', fromlist=['display']).display
+                    Image = __import__('IPython.display', fromlist=['Image']).Image
+                    img_bytes = fig.to_image(format="png", scale=2)
+                    display(Image(img_bytes))
+                except Exception as e:
+                    fig.show()
 
             if generate_followups:
                 followup_questions = generate_followup_questions(question=question, df=df)
                 if print_results and followup_questions is not None and len(followup_questions) > 0:
                     print("AI-generated follow-up questions:")
                     for followup_question in followup_questions:
                         print(followup_question)
+                        pass
 
                 if print_results:
                     return None