[Bug] Fix duplicate figure display in jupyter notebook (#527)

Signed-off-by: dependabot[bot] <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Petar Pejovic <[email protected]> Co-authored-by: Anna Xiong <[email protected]> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: Li Nguyen <[email protected]> Co-authored-by: Maximilian Schulz <[email protected]>
mckinsey · Jul 15, 2024 · 6ac7fce · 6ac7fce
1 parent b8b15d4
commit 6ac7fce
Show file tree

Hide file tree

Showing 10 changed files with 80 additions and 27 deletions.
diff --git a/...-ai/changelog.d/20240619_092410_nadija_ratkusic_graca_duplicate_fig_diplayed.md b/...-ai/changelog.d/20240619_092410_nadija_ratkusic_graca_duplicate_fig_diplayed.md
@@ -0,0 +1,47 @@
+<!--
+A new scriv changelog fragment.
+
+Uncomment the section that is right (remove the HTML comment wrapper).
+-->
+
+<!--
+### Highlights ✨
+
+- A bullet item for the Highlights ✨ category with a link to the relevant PR at the end of your entry, e.g. Enable feature XXX ([#1](https://github.com/mckinsey/vizro/pull/1))
+
+-->
+<!--
+### Removed
+
+- A bullet item for the Removed category with a link to the relevant PR at the end of your entry, e.g. Enable feature XXX ([#1](https://github.com/mckinsey/vizro/pull/1))
+
+-->
+<!--
+### Added
+
+- A bullet item for the Added category with a link to the relevant PR at the end of your entry, e.g. Enable feature XXX ([#1](https://github.com/mckinsey/vizro/pull/1))
+
+-->
+
+### Changed
+
+- Disabled figure display upon variable assignment, to display outcome of `VizroAI.plot()` add `.show()`. Enable feature XXX ([#527](https://github.com/mckinsey/vizro/pull/527))
+
+<!--
+### Deprecated
+
+- A bullet item for the Deprecated category with a link to the relevant PR at the end of your entry, e.g. Enable feature XXX ([#1](https://github.com/mckinsey/vizro/pull/1))
+
+-->
+<!--
+### Fixed
+
+- A bullet item for the Fixed category with a link to the relevant PR at the end of your entry, e.g. Enable feature XXX ([#1](https://github.com/mckinsey/vizro/pull/1))
+
+-->
+<!--
+### Security
+
+- A bullet item for the Security category with a link to the relevant PR at the end of your entry, e.g. Enable feature XXX ([#1](https://github.com/mckinsey/vizro/pull/1))
+
+-->
diff --git a/vizro-ai/docs/pages/tutorials/quickstart.md b/vizro-ai/docs/pages/tutorials/quickstart.md
@@ -92,7 +92,8 @@ And that's it! By passing the prepared data and written visualization request, V
         df = px.data.gapminder()
 
         vizro_ai = VizroAI()
-        vizro_ai.plot(df, "create a line graph for GDP per capita since 1950 for each continent. Mark the x axis as Year, y axis as GDP Per Cap and don't include a title", explain=True)
+        fig = vizro_ai.plot(df, "create a line graph for GDP per capita since 1950 for each continent. Mark the x axis as Year, y axis as GDP Per Cap and don't include a title", explain=True)
+        fig.show()
         ```
     === "Result"
         [![LineGraph]][LineGraph]
@@ -128,7 +129,8 @@ Let's create another example to illustrate the code and insights returned when p
 
     === "Code for the cell"
         ```py
-        vizro_ai.plot(df, "show me the geo distribution of life expectancy", explain=True)
+        fig = vizro_ai.plot(df, "show me the geo distribution of life expectancy", explain=True)
+        fig.show()
         ```
     === "Result"
         [![GeoDistribution]][GeoDistribution]

diff --git a/vizro-ai/docs/pages/user-guides/add-generated-chart-usecase.md b/vizro-ai/docs/pages/user-guides/add-generated-chart-usecase.md
@@ -22,12 +22,13 @@ This guide explains the different ways in which you can add a chart generated by
             df = px.data.gapminder()
             vizro_ai = VizroAI(model="gpt-4-0613")
 
-            vizro_ai.plot(df,
+            fig = vizro_ai.plot(df,
                           """Plot a bubble chart to show the changes in life expectancy
                            and GDP per capita for each country over time.
                            Color the bubbles by continent.
                            Add animation on yearly basis, and do not use facets.
                            Put the legend on top""", explain=True)
+            fig.show()
             ```
         === "Result"
             [![VizroAIChart]][VizroAIChart]

diff --git a/vizro-ai/docs/pages/user-guides/chart-examples.md b/vizro-ai/docs/pages/user-guides/chart-examples.md
@@ -25,11 +25,11 @@ It's suitable for comparing multiple variables across different categories or di
         df = px.data.wind()
 
         vizro_ai = VizroAI(model="gpt-4-0613")
-        vizro_ai.plot(df,
+        fig = vizro_ai.plot(df,
                       """Describe wind frequency and direction using bar_polar chart.
                          Increase the width and height of the figure.
                          Improve layout by placing title to the left. Show legend""", explain=True)
-
+        fig.show()
         ```
 
     [VizroAIChart1]: ../../assets/user_guides/polar_bar_chart.png
@@ -56,12 +56,12 @@ The next chart we'll look at is a geographical map chart to visualize spatial pa
         df = px.data.wind()
 
         vizro_ai = VizroAI(model="gpt-4-0613")
-        vizro_ai.plot(df,
+        fig = vizro_ai.plot(df,
                       """Visualize life expectancy over the years using map chart. Use life expectancy as the color dimension.
                          Improve layout by using Arial font. Increase the width and height of the map area. Outline continents on the map.
                          Show countries on the map.
                          Increase the width and height of the figure.""", explain=True)
-
+        fig.show()
         ```
 
     [VizroAIChart2]: ../../assets/user_guides/map_chart.gif
@@ -88,8 +88,8 @@ Let's explore how to generate a 3-dimensional surface plot with VizroAI.
         df = px.data.gapminder()
 
         vizro_ai = VizroAI(model="gpt-4-0613")
-        vizro_ai.plot(df, "create a surface plot")
-
+        fig = vizro_ai.plot(df, "create a surface plot")
+        fig.show()
         ```
 
     [VizroAIChart3]: ../../assets/user_guides/surface_plot.gif
diff --git a/vizro-ai/docs/pages/user-guides/create-advanced-charts.md b/vizro-ai/docs/pages/user-guides/create-advanced-charts.md
@@ -13,7 +13,8 @@ We'll create an animated bar chart illustrating the population development of ea
         df = px.data.gapminder()
 
         vizro_ai = VizroAI()
-        vizro_ai.plot(df, "The chart should be an animated stacked bar chart with population on the y axis and continent on the x axis with all respective countries, allowing you to observe changes in the population over consecutive years.")
+        fig = vizro_ai.plot(df, "The chart should be an animated stacked bar chart with population on the y axis and continent on the x axis with all respective countries, allowing you to observe changes in the population over consecutive years.")
+        fig.show()
         ```
     === "Result"
         [![AnimatedChart1]][AnimatedChart1]
@@ -32,7 +33,8 @@ Having unveiled our animated bar chart showcasing population development per cou
         df = px.data.gapminder()
 
         vizro_ai = VizroAI()
-        vizro_ai.plot(df, "The chart should be an animated stacked bar chart with population on the y axis and continent on the x axis with all respective countries, allowing you to observe changes in the population over consecutive years. Please improve layout.")
+        fig = vizro_ai.plot(df, "The chart should be an animated stacked bar chart with population on the y axis and continent on the x axis with all respective countries, allowing you to observe changes in the population over consecutive years. Please improve layout.")
+        fig.show()
         ```
     === "Result"
         [![AnimatedChart2]][AnimatedChart2]
@@ -54,7 +56,8 @@ Upon closer inspection, two challenges emerge: the legend overlaps the x-axis an
         df = px.data.gapminder()
 
         vizro_ai = VizroAI()
-        vizro_ai.plot(df, "The chart should be an animated stacked bar chart with population on the y axis and continent on the x axis with all respective countries, allowing you to observe changes in the population over consecutive years. Make sure that y axis range fits entire data. Please improve layout and optimize layout of legend.")
+        fig = vizro_ai.plot(df, "The chart should be an animated stacked bar chart with population on the y axis and continent on the x axis with all respective countries, allowing you to observe changes in the population over consecutive years. Make sure that y axis range fits entire data. Please improve layout and optimize layout of legend.")
+        fig.show()
         ```
     === "Result"
         [![AnimatedChart3]][AnimatedChart3]

diff --git a/vizro-ai/docs/pages/user-guides/customize-vizro-ai.md b/vizro-ai/docs/pages/user-guides/customize-vizro-ai.md
@@ -77,7 +77,8 @@ To ensure a deterministic answer to our queries, we've set the temperature to 0.
             max_retries=5,
         )
         vizro_ai = VizroAI(model=llm)
-        vizro_ai.plot(df, "describe the composition of gdp in continent")
+        fig = vizro_ai.plot(df, "describe the composition of gdp in continent")
+        fig.show()
         ```
 
 Passing an instantiated model to `VizroAI` lets you customize it, and additionally, it enables you to use an OpenAI model that is not included in the above list of [supported models](#supported-models).

diff --git a/vizro-ai/docs/pages/user-guides/run-vizro-ai.md b/vizro-ai/docs/pages/user-guides/run-vizro-ai.md
@@ -19,7 +19,8 @@ To run Vizro-AI code in a Jupyter Notebook, create a new cell and execute the co
         vizro_ai = VizroAI()
 
         df = px.data.gapminder()
-        vizro_ai.plot(df, "visualize the life expectancy per continent and color each continent")
+        fig = vizro_ai.plot(df, "visualize the life expectancy per continent and color each continent")
+        fig.show()
         ```
     === "Result"
         [![BarChart]][BarChart]

diff --git a/vizro-ai/docs/pages/user-guides/use-different-languages.md b/vizro-ai/docs/pages/user-guides/use-different-languages.md
@@ -11,7 +11,8 @@ Vizro-AI is versatile, supporting prompts and chart visualizations in languages
         df = px.data.gapminder()
 
         vizro_ai = VizroAI()
-        vizro_ai.plot(df, "请画一个世界年均GDP的趋势图")
+        fig = vizro_ai.plot(df, "请画一个世界年均GDP的趋势图")
+        fig.show()
         ```
     === "Result"
         [![ChineseChart]][ChineseChart]
@@ -30,7 +31,8 @@ Subsequently, we'll switch to German and prompt the visualization of life expect
         df = px.data.gapminder()
 
         vizro_ai = VizroAI()
-        vizro_ai.plot(df, "Visualiere den Trend von der Lebenserwartung in USA über die Jahre im Vergleich zur Veränderung der weltweiten Lebenserwartung über die Jahre und kreiere eine deutsche Visualisierung", explain=True)
+        fig = vizro_ai.plot(df, "Visualiere den Trend von der Lebenserwartung in USA über die Jahre im Vergleich zur Veränderung der weltweiten Lebenserwartung über die Jahre und kreiere eine deutsche Visualisierung", explain=True)
+        fig.show()
         ```
     === "Result"
         [![GermanChart]][GermanChart]

diff --git a/vizro-ai/src/vizro_ai/_vizro_ai.py b/vizro-ai/src/vizro_ai/_vizro_ai.py
@@ -56,7 +56,11 @@ def _lazy_get_component(self, component_class: Any) -> Any:  # TODO configure co
         return self.components_instances[component_class]
 
     def _run_plot_tasks(
-        self, df: pd.DataFrame, user_input: str, max_debug_retry: int = 3, explain: bool = False
+        self,
+        df: pd.DataFrame,
+        user_input: str,
+        max_debug_retry: int = 3,
+        explain: bool = False,
     ) -> PlotOutputs:
         """Task execution."""
         chart_type_pipeline = self.pipeline_manager.chart_type_pipeline
@@ -83,9 +87,7 @@ def _run_plot_tasks(
                 "or try to select a different model. Fallout response is provided: \n\n" + code_string
             )
 
-        fig_object = _exec_code_and_retrieve_fig(
-            code=code_string, local_args={"df": df}, show_fig=_is_jupyter(), is_notebook_env=_is_jupyter()
-        )
+        fig_object = _exec_code_and_retrieve_fig(code=code_string, local_args={"df": df}, is_notebook_env=_is_jupyter())
         if explain:
             business_insights, code_explanation = self._lazy_get_component(GetCodeExplanation).run(
                 chain_input=user_input, code_snippet=code_string

diff --git a/vizro-ai/src/vizro_ai/utils/helper.py b/vizro-ai/src/vizro_ai/utils/helper.py
@@ -60,14 +60,13 @@ def _debug_helper(
 
 
 def _exec_code_and_retrieve_fig(
-    code: str, local_args: Optional[Dict] = None, show_fig: bool = False, is_notebook_env: bool = True
+    code: str, local_args: Optional[Dict] = None, is_notebook_env: bool = True
 ) -> go.Figure:
     """Execute code in notebook with correct namespace and return fig object.
 
     Args:
         code: code string to be executed
         local_args: additional local arguments
-        show_fig: boolean flag indicating if fig will be rendered automatically
         is_notebook_env: boolean flag indicating if code is run in Jupyter notebook
 
     Returns:
@@ -76,11 +75,6 @@ def _exec_code_and_retrieve_fig(
     """
     from IPython import get_ipython
 
-    if show_fig and "\nfig.show()" not in code:
-        code += "\nfig.show()"
-    elif not show_fig:
-        code = code.replace("fig.show()", "")
-
     namespace = get_ipython().user_ns if is_notebook_env else globals()
 
     if local_args: