From 993ad1b68749d1848f24922d48a4852e51a76004 Mon Sep 17 00:00:00 2001 From: AnnaXiong Date: Thu, 16 May 2024 09:32:32 -0400 Subject: [PATCH 01/21] update helper --- vizro-ai/src/vizro_ai/utils/helper.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/vizro-ai/src/vizro_ai/utils/helper.py b/vizro-ai/src/vizro_ai/utils/helper.py index c7b6744f7..5ded6af90 100644 --- a/vizro-ai/src/vizro_ai/utils/helper.py +++ b/vizro-ai/src/vizro_ai/utils/helper.py @@ -75,14 +75,10 @@ def _exec_code_and_retrieve_fig( return dashboard_ready_fig -def _exec_fig_code_display_markdown( - df: pd.DataFrame, code_snippet: str, biz_insights: str, code_explain: str -) -> go.Figure: - # TODO change default test str to other +def _display_markdown(code_snippet: str, biz_insights: str, code_explain: str) -> go.Figure: """Display chart and Markdown format description in jupyter and returns fig object. Args: - df: The dataframe to be analyzed. code_snippet: code string to be executed biz_insights: business insights to be displayed in markdown cell code_explain: code explanation to be displayed in markdown cell @@ -100,7 +96,6 @@ def _exec_fig_code_display_markdown( markdown_code = f"```\n{code_snippet}\n```" output_text = f"

Insights:

\n\n{biz_insights}\n

Code:

\n\n{code_explain}\n{markdown_code}" display(Markdown(output_text)) - return _exec_code_and_retrieve_fig(code_snippet, local_args={"df": df}, is_notebook_env=_is_jupyter()) class DebugFailure(Exception): From dee20fbb791d2410081c94ed5d56b302badc0b70 Mon Sep 17 00:00:00 2001 From: AnnaXiong Date: Thu, 16 May 2024 09:32:51 -0400 Subject: [PATCH 02/21] add get all output method --- vizro-ai/src/vizro_ai/_vizro_ai.py | 60 ++++++++++++++++++++++-------- 1 file changed, 45 insertions(+), 15 deletions(-) diff --git a/vizro-ai/src/vizro_ai/_vizro_ai.py b/vizro-ai/src/vizro_ai/_vizro_ai.py index 8417ca583..c90d6eb55 100644 --- a/vizro-ai/src/vizro_ai/_vizro_ai.py +++ b/vizro-ai/src/vizro_ai/_vizro_ai.py @@ -11,8 +11,8 @@ from vizro_ai.utils.helper import ( DebugFailure, _debug_helper, + _display_markdown, _exec_code_and_retrieve_fig, - _exec_fig_code_display_markdown, _is_jupyter, ) @@ -78,6 +78,12 @@ def _run_plot_tasks( code_string = validated_code_dict.get("code_string") business_insights, code_explanation = None, None + if pass_validation is False: + raise DebugFailure( + "Chart creation failed. Retry debugging has reached maximum limit. Try to rephrase the prompt, " + "or try to select a different model. Fallout response is provided: \n\n" + code_string + ) + if explain and pass_validation: business_insights, code_explanation = self._lazy_get_component(GetCodeExplanation).run( chain_input=user_input, code_snippet=code_string @@ -103,7 +109,11 @@ def _get_chart_code(self, df: pd.DataFrame, user_input: str) -> str: return self._run_plot_tasks(df, user_input, explain=False).get("code_string") def plot( - self, df: pd.DataFrame, user_input: str, explain: bool = False, max_debug_retry: int = 3 + self, + df: pd.DataFrame, + user_input: str, + explain: bool = False, + max_debug_retry: int = 3, ) -> Union[go.Figure, Dict[str, Any]]: """Plot visuals using vizro via english descriptions, english to chart translation. @@ -114,7 +124,7 @@ def plot( max_debug_retry: Maximum number of retries to debug errors. Defaults to `3`. Returns: - Plotly Figure object or a dictionary containing data + go.Figure """ output_dict = self._run_plot_tasks(df, user_input, explain=explain, max_debug_retry=max_debug_retry) @@ -122,18 +132,38 @@ def plot( business_insights = output_dict.get("business_insights") code_explanation = output_dict.get("code_explanation") - if code_string.startswith("Failed to debug code"): - raise DebugFailure( - "Chart creation failed. Retry debugging has reached maximum limit. Try to rephrase the prompt, " - "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}, is_notebook_env=_is_jupyter()) - # TODO Tentative for integration test - if self._return_all_text: - return output_dict if not explain: - return _exec_code_and_retrieve_fig(code=code_string, local_args={"df": df}, is_notebook_env=_is_jupyter()) + return fig_object + if explain: - return _exec_fig_code_display_markdown( - df=df, code_snippet=code_string, biz_insights=business_insights, code_explain=code_explanation - ) + _display_markdown(code_snippet=code_string, biz_insights=business_insights, code_explain=code_explanation) + return fig_object + + def get_all_output( + self, + df: pd.DataFrame, + user_input: str, + explain: bool = False, + max_debug_retry: int = 3, + ): + """Executes plotting tasks based on user input and retrieves output dictionary. + + Args: + df: The dataframe to be analyzed. + user_input: User input that dictates what operations to perform on the DataFrame. + explain: A flag to indicate whether to include explanations. Defaults to False. + max_debug_retry: Maximum number of retries for debugging the generated code. Defaults to 3. + + Returns: + dict: A dictionary containing the generated code string, the figure object, and other text explanation + depending on whether explain flag is true. + + """ + output_dict = self._run_plot_tasks(df, user_input, explain=explain, max_debug_retry=max_debug_retry) + code_string = output_dict.get("code_string") + fig_object = _exec_code_and_retrieve_fig(code=code_string, local_args={"df": df}, is_notebook_env=_is_jupyter()) + output_dict.update({"fig": fig_object}) + + return output_dict From 2dd6ddc83c2d601148e89a32154917837640e975 Mon Sep 17 00:00:00 2001 From: AnnaXiong Date: Thu, 16 May 2024 13:47:42 -0400 Subject: [PATCH 03/21] update get_outputs --- vizro-ai/src/vizro_ai/_vizro_ai.py | 31 +++++++++++++++++++----------- 1 file changed, 20 insertions(+), 11 deletions(-) diff --git a/vizro-ai/src/vizro_ai/_vizro_ai.py b/vizro-ai/src/vizro_ai/_vizro_ai.py index c90d6eb55..6f96f6495 100644 --- a/vizro-ai/src/vizro_ai/_vizro_ai.py +++ b/vizro-ai/src/vizro_ai/_vizro_ai.py @@ -76,24 +76,27 @@ def _run_plot_tasks( pass_validation = validated_code_dict.get("debug_status") code_string = validated_code_dict.get("code_string") - business_insights, code_explanation = None, None - if pass_validation is False: + if not pass_validation: raise DebugFailure( "Chart creation failed. Retry debugging has reached maximum limit. Try to rephrase the prompt, " "or try to select a different model. Fallout response is provided: \n\n" + code_string ) - if explain and pass_validation: + elif explain: business_insights, code_explanation = self._lazy_get_component(GetCodeExplanation).run( chain_input=user_input, code_snippet=code_string ) - return { - "business_insights": business_insights, - "code_explanation": code_explanation, - "code_string": code_string, - } + return { + "business_insights": business_insights, + "code_explanation": code_explanation, + "code_string": code_string, + } + else: + return { + "code_string": code_string, + } def _get_chart_code(self, df: pd.DataFrame, user_input: str) -> str: """Get Chart code of vizro via english descriptions, English to chart translation. @@ -129,8 +132,6 @@ def plot( """ output_dict = self._run_plot_tasks(df, user_input, explain=explain, max_debug_retry=max_debug_retry) code_string = output_dict.get("code_string") - business_insights = output_dict.get("business_insights") - code_explanation = output_dict.get("code_explanation") fig_object = _exec_code_and_retrieve_fig(code=code_string, local_args={"df": df}, is_notebook_env=_is_jupyter()) @@ -138,10 +139,12 @@ def plot( return fig_object if explain: + business_insights = output_dict.get("business_insights", None) + code_explanation = output_dict.get("code_explanation", None) _display_markdown(code_snippet=code_string, biz_insights=business_insights, code_explain=code_explanation) return fig_object - def get_all_output( + def get_outputs( self, df: pd.DataFrame, user_input: str, @@ -166,4 +169,10 @@ def get_all_output( fig_object = _exec_code_and_retrieve_fig(code=code_string, local_args={"df": df}, is_notebook_env=_is_jupyter()) output_dict.update({"fig": fig_object}) + if explain is False: + logger.info( + "Flag explain is set to False. business_insights and code_explanation will not be included in " + "output dictionary" + ) + return output_dict From b1c0929f099ebb01dac14fd96d33aaa1b87c226a Mon Sep 17 00:00:00 2001 From: AnnaXiong Date: Thu, 16 May 2024 13:52:14 -0400 Subject: [PATCH 04/21] update get_outputs --- vizro-ai/src/vizro_ai/_vizro_ai.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vizro-ai/src/vizro_ai/_vizro_ai.py b/vizro-ai/src/vizro_ai/_vizro_ai.py index 6f96f6495..cbb29ea31 100644 --- a/vizro-ai/src/vizro_ai/_vizro_ai.py +++ b/vizro-ai/src/vizro_ai/_vizro_ai.py @@ -172,7 +172,7 @@ def get_outputs( if explain is False: logger.info( "Flag explain is set to False. business_insights and code_explanation will not be included in " - "output dictionary" + "the output dictionary" ) return output_dict From c7289b8e23fbacabf479ec11ef7842a20be7c8b2 Mon Sep 17 00:00:00 2001 From: AnnaXiong Date: Thu, 16 May 2024 14:03:33 -0400 Subject: [PATCH 05/21] update get_plot_outputs --- vizro-ai/src/vizro_ai/_vizro_ai.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vizro-ai/src/vizro_ai/_vizro_ai.py b/vizro-ai/src/vizro_ai/_vizro_ai.py index cbb29ea31..d96589ce1 100644 --- a/vizro-ai/src/vizro_ai/_vizro_ai.py +++ b/vizro-ai/src/vizro_ai/_vizro_ai.py @@ -144,7 +144,7 @@ def plot( _display_markdown(code_snippet=code_string, biz_insights=business_insights, code_explain=code_explanation) return fig_object - def get_outputs( + def get_plot_outputs( self, df: pd.DataFrame, user_input: str, From e6debe987846edb6a68815e0bb704b140043d290 Mon Sep 17 00:00:00 2001 From: AnnaXiong Date: Thu, 16 May 2024 14:03:42 -0400 Subject: [PATCH 06/21] changelog --- ...na_xiong_get_vizro_ai_customized_output.md | 47 +++++++++++++++++++ 1 file changed, 47 insertions(+) create mode 100644 vizro-ai/changelog.d/20240516_135901_anna_xiong_get_vizro_ai_customized_output.md diff --git a/vizro-ai/changelog.d/20240516_135901_anna_xiong_get_vizro_ai_customized_output.md b/vizro-ai/changelog.d/20240516_135901_anna_xiong_get_vizro_ai_customized_output.md new file mode 100644 index 000000000..7cadc52f0 --- /dev/null +++ b/vizro-ai/changelog.d/20240516_135901_anna_xiong_get_vizro_ai_customized_output.md @@ -0,0 +1,47 @@ + + + + + +### Added + +- Enable feature to get all possible outputs from VizroAI.plot() using VizroAI.get_plot_outputs() ([#488](https://github.com/mckinsey/vizro/pull/488)). It returns a dictionary that contains code string, fig object, business insights, code explanation. + + + + + From 4b40f38e94c0817f893a8578bf64ef7825c7d815 Mon Sep 17 00:00:00 2001 From: AnnaXiong Date: Thu, 23 May 2024 11:30:10 -0400 Subject: [PATCH 07/21] fix --- vizro-ai/src/vizro_ai/_vizro_ai.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/vizro-ai/src/vizro_ai/_vizro_ai.py b/vizro-ai/src/vizro_ai/_vizro_ai.py index d96589ce1..8e9816c0a 100644 --- a/vizro-ai/src/vizro_ai/_vizro_ai.py +++ b/vizro-ai/src/vizro_ai/_vizro_ai.py @@ -127,7 +127,7 @@ def plot( max_debug_retry: Maximum number of retries to debug errors. Defaults to `3`. Returns: - go.Figure + Plotly Figure object or a dictionary containing data """ output_dict = self._run_plot_tasks(df, user_input, explain=explain, max_debug_retry=max_debug_retry) @@ -135,6 +135,10 @@ def plot( fig_object = _exec_code_and_retrieve_fig(code=code_string, local_args={"df": df}, is_notebook_env=_is_jupyter()) + # TODO Tentative for integration test, will be updated/removed for new tests + if self._return_all_text: + return output_dict + if not explain: return fig_object From a6c811dd38b2a36742d198bbe5ddab46967a1ec9 Mon Sep 17 00:00:00 2001 From: AnnaXiong Date: Thu, 23 May 2024 11:32:14 -0400 Subject: [PATCH 08/21] changelog --- ...na_xiong_get_vizro_ai_customized_output.md | 47 +++++++++++++++++++ 1 file changed, 47 insertions(+) create mode 100644 vizro-ai/changelog.d/20240523_113018_anna_xiong_get_vizro_ai_customized_output.md diff --git a/vizro-ai/changelog.d/20240523_113018_anna_xiong_get_vizro_ai_customized_output.md b/vizro-ai/changelog.d/20240523_113018_anna_xiong_get_vizro_ai_customized_output.md new file mode 100644 index 000000000..1d256e27b --- /dev/null +++ b/vizro-ai/changelog.d/20240523_113018_anna_xiong_get_vizro_ai_customized_output.md @@ -0,0 +1,47 @@ + + + + + +### Added + +- Enable feature to get customized text outputs from `VizroAI.plot` ([#488](https://github.com/mckinsey/vizro/pull/488)) + + + + + From db450e4973aed305dc3ec00e6645edbfb48f8ea7 Mon Sep 17 00:00:00 2001 From: AnnaXiong Date: Thu, 23 May 2024 12:54:25 -0400 Subject: [PATCH 09/21] fix --- vizro-ai/src/vizro_ai/_vizro_ai.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vizro-ai/src/vizro_ai/_vizro_ai.py b/vizro-ai/src/vizro_ai/_vizro_ai.py index 8e9816c0a..cc1503ba5 100644 --- a/vizro-ai/src/vizro_ai/_vizro_ai.py +++ b/vizro-ai/src/vizro_ai/_vizro_ai.py @@ -164,7 +164,7 @@ def get_plot_outputs( max_debug_retry: Maximum number of retries for debugging the generated code. Defaults to 3. Returns: - dict: A dictionary containing the generated code string, the figure object, and other text explanation + A dictionary containing the generated code string, the figure object, and other text explanation depending on whether explain flag is true. """ From cc6fea63bebdd32bf1f0d6de82a7edc600955749 Mon Sep 17 00:00:00 2001 From: AnnaXiong Date: Thu, 23 May 2024 13:00:31 -0400 Subject: [PATCH 10/21] docstring --- vizro-ai/src/vizro_ai/_vizro_ai.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vizro-ai/src/vizro_ai/_vizro_ai.py b/vizro-ai/src/vizro_ai/_vizro_ai.py index cc1503ba5..1e96d800f 100644 --- a/vizro-ai/src/vizro_ai/_vizro_ai.py +++ b/vizro-ai/src/vizro_ai/_vizro_ai.py @@ -117,7 +117,7 @@ def plot( user_input: str, explain: bool = False, max_debug_retry: int = 3, - ) -> Union[go.Figure, Dict[str, Any]]: + ) -> go.Figure: """Plot visuals using vizro via english descriptions, english to chart translation. Args: @@ -154,7 +154,7 @@ def get_plot_outputs( user_input: str, explain: bool = False, max_debug_retry: int = 3, - ): + ) -> Dict[str, Any]: """Executes plotting tasks based on user input and retrieves output dictionary. Args: From 20ba6a5a8f65b265cbcd5d52b6951d180165cd23 Mon Sep 17 00:00:00 2001 From: AnnaXiong Date: Mon, 3 Jun 2024 17:22:10 -0400 Subject: [PATCH 11/21] add dataclass and update show fig utils --- vizro-ai/src/vizro_ai/utils/helper.py | 27 +++++++++++++++++++++------ 1 file changed, 21 insertions(+), 6 deletions(-) diff --git a/vizro-ai/src/vizro_ai/utils/helper.py b/vizro-ai/src/vizro_ai/utils/helper.py index 5ded6af90..e7c2b75bf 100644 --- a/vizro-ai/src/vizro_ai/utils/helper.py +++ b/vizro-ai/src/vizro_ai/utils/helper.py @@ -1,6 +1,7 @@ """Helper Functions For Vizro AI.""" import traceback +from dataclasses import dataclass, field from typing import Callable, Dict, Optional import pandas as pd @@ -9,6 +10,16 @@ from .safeguard import _safeguard_check +@dataclass +class Plot: + """Data class about a vizro ai plot.""" + + code_string: str + fig: go.Figure = field(default=None) + business_insights: str = field(default=None) + code_explanation: str = field(default=None) + + # Taken from rich.console. See https://github.com/Textualize/rich. def _is_jupyter() -> bool: # pragma: no cover """Checks if we're running in a Jupyter notebook.""" @@ -49,22 +60,29 @@ def _debug_helper( def _exec_code_and_retrieve_fig( - code: str, local_args: Optional[Dict] = None, is_notebook_env: bool = True + code: str, local_args: Optional[Dict] = None, show_fig: bool = False, 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 to flag whether to add show fig code line is_notebook_env: boolean flag indicating if code is run in Jupyter notebook Returns: - go.Figure + Plotly go figure """ 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: namespace.update(local_args) _safeguard_check(code) @@ -75,7 +93,7 @@ def _exec_code_and_retrieve_fig( return dashboard_ready_fig -def _display_markdown(code_snippet: str, biz_insights: str, code_explain: str) -> go.Figure: +def _display_markdown(code_snippet: str, biz_insights: str, code_explain: str) -> None: """Display chart and Markdown format description in jupyter and returns fig object. Args: @@ -83,9 +101,6 @@ def _display_markdown(code_snippet: str, biz_insights: str, code_explain: str) - biz_insights: business insights to be displayed in markdown cell code_explain: code explanation to be displayed in markdown cell - Returns: - go.Figure - """ try: # pylint: disable=import-outside-toplevel From cbb2a64a0d7d3d4f7bd2e28ac6a0603d2d6284b6 Mon Sep 17 00:00:00 2001 From: AnnaXiong Date: Mon, 3 Jun 2024 17:22:27 -0400 Subject: [PATCH 12/21] refactor return to Plot dataclass --- vizro-ai/src/vizro_ai/_vizro_ai.py | 108 ++++++++++++----------------- 1 file changed, 45 insertions(+), 63 deletions(-) diff --git a/vizro-ai/src/vizro_ai/_vizro_ai.py b/vizro-ai/src/vizro_ai/_vizro_ai.py index 1e96d800f..b7a16d66c 100644 --- a/vizro-ai/src/vizro_ai/_vizro_ai.py +++ b/vizro-ai/src/vizro_ai/_vizro_ai.py @@ -1,5 +1,6 @@ import logging -from typing import Any, Dict, Optional, Union +from dataclasses import asdict +from typing import Any, Optional, Union import pandas as pd import plotly.graph_objects as go @@ -10,6 +11,7 @@ from vizro_ai.task_pipeline._pipeline_manager import PipelineManager from vizro_ai.utils.helper import ( DebugFailure, + Plot, _debug_helper, _display_markdown, _exec_code_and_retrieve_fig, @@ -23,7 +25,7 @@ class VizroAI: """Vizro-AI main class.""" pipeline_manager: PipelineManager = PipelineManager() - _return_all_text: bool = False + _return_all_text: bool = False # TODO deleted after adding new integration test def __init__(self, model: Optional[Union[ChatOpenAI, str]] = None): """Initialization of VizroAI. @@ -57,7 +59,7 @@ def _lazy_get_component(self, component_class: Any) -> Any: # TODO configure co def _run_plot_tasks( self, df: pd.DataFrame, user_input: str, max_debug_retry: int = 3, explain: bool = False - ) -> Dict[str, Any]: + ) -> Plot: """Task execution.""" chart_type_pipeline = self.pipeline_manager.chart_type_pipeline chart_types = chart_type_pipeline.run(initial_args={"chain_input": user_input, "df": df}) @@ -82,21 +84,23 @@ def _run_plot_tasks( "Chart creation failed. Retry debugging has reached maximum limit. Try to rephrase the prompt, " "or try to select a different model. Fallout response is provided: \n\n" + code_string ) - - elif explain: - business_insights, code_explanation = self._lazy_get_component(GetCodeExplanation).run( - chain_input=user_input, code_snippet=code_string - ) - - return { - "business_insights": business_insights, - "code_explanation": code_explanation, - "code_string": code_string, - } else: - return { - "code_string": code_string, - } + fig_object = _exec_code_and_retrieve_fig( + code=code_string, local_args={"df": df}, show_fig=True, 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 + ) + + return Plot( + code_string=code_string, + fig=fig_object, + business_insights=business_insights, + code_explanation=code_explanation, + ) + else: + return Plot(code_string=code_string, fig=fig_object) def _get_chart_code(self, df: pd.DataFrame, user_input: str) -> str: """Get Chart code of vizro via english descriptions, English to chart translation. @@ -108,7 +112,7 @@ def _get_chart_code(self, df: pd.DataFrame, user_input: str) -> str: user_input: User questions or descriptions of the desired visual """ - # TODO refine and update error handling + # TODO retained for some chat application integration, need deprecation handling return self._run_plot_tasks(df, user_input, explain=False).get("code_string") def plot( @@ -117,7 +121,8 @@ def plot( user_input: str, explain: bool = False, max_debug_retry: int = 3, - ) -> go.Figure: + return_plot_components: bool = False, + ) -> Union[go.Figure, Plot]: """Plot visuals using vizro via english descriptions, english to chart translation. Args: @@ -125,53 +130,15 @@ def plot( user_input: User questions or descriptions of the desired visual. explain: Flag to include explanation in response. max_debug_retry: Maximum number of retries to debug errors. Defaults to `3`. + return_plot_components: Flag to return plot dataclass that includes all components. Returns: - Plotly Figure object or a dictionary containing data - - """ - output_dict = self._run_plot_tasks(df, user_input, explain=explain, max_debug_retry=max_debug_retry) - code_string = output_dict.get("code_string") - - fig_object = _exec_code_and_retrieve_fig(code=code_string, local_args={"df": df}, is_notebook_env=_is_jupyter()) - - # TODO Tentative for integration test, will be updated/removed for new tests - if self._return_all_text: - return output_dict - - if not explain: - return fig_object - - if explain: - business_insights = output_dict.get("business_insights", None) - code_explanation = output_dict.get("code_explanation", None) - _display_markdown(code_snippet=code_string, biz_insights=business_insights, code_explain=code_explanation) - return fig_object - - def get_plot_outputs( - self, - df: pd.DataFrame, - user_input: str, - explain: bool = False, - max_debug_retry: int = 3, - ) -> Dict[str, Any]: - """Executes plotting tasks based on user input and retrieves output dictionary. - - Args: - df: The dataframe to be analyzed. - user_input: User input that dictates what operations to perform on the DataFrame. - explain: A flag to indicate whether to include explanations. Defaults to False. - max_debug_retry: Maximum number of retries for debugging the generated code. Defaults to 3. - - Returns: - A dictionary containing the generated code string, the figure object, and other text explanation - depending on whether explain flag is true. + go.Figure or Plot dataclass """ - output_dict = self._run_plot_tasks(df, user_input, explain=explain, max_debug_retry=max_debug_retry) - code_string = output_dict.get("code_string") - fig_object = _exec_code_and_retrieve_fig(code=code_string, local_args={"df": df}, is_notebook_env=_is_jupyter()) - output_dict.update({"fig": fig_object}) + vizro_plot = self._run_plot_tasks( + df=df, user_input=user_input, explain=explain, max_debug_retry=max_debug_retry + ) if explain is False: logger.info( @@ -179,4 +146,19 @@ def get_plot_outputs( "the output dictionary" ) - return output_dict + else: + _display_markdown( + code_snippet=vizro_plot.code_string, + biz_insights=vizro_plot.business_insights, + code_explain=vizro_plot.code_explanation, + ) + + if return_plot_components: + return vizro_plot + + else: + return vizro_plot.fig + + # TODO Tentative for integration test, will be updated/removed for new tests + if self._return_all_text: + return asdict(vizro_plot) From 1c0fc50e06496ed975ac8c2d7a92f8f4a5b1ff25 Mon Sep 17 00:00:00 2001 From: AnnaXiong Date: Mon, 3 Jun 2024 17:34:55 -0400 Subject: [PATCH 13/21] fix --- vizro-ai/src/vizro_ai/_vizro_ai.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/vizro-ai/src/vizro_ai/_vizro_ai.py b/vizro-ai/src/vizro_ai/_vizro_ai.py index b7a16d66c..d14a92c09 100644 --- a/vizro-ai/src/vizro_ai/_vizro_ai.py +++ b/vizro-ai/src/vizro_ai/_vizro_ai.py @@ -115,7 +115,7 @@ def _get_chart_code(self, df: pd.DataFrame, user_input: str) -> str: # TODO retained for some chat application integration, need deprecation handling return self._run_plot_tasks(df, user_input, explain=False).get("code_string") - def plot( + def plot( # pylint: disable=too-many-arguments # noqa: PLR0913 self, df: pd.DataFrame, user_input: str, @@ -153,12 +153,12 @@ def plot( code_explain=vizro_plot.code_explanation, ) + # TODO Tentative for integration test, will be updated/removed for new tests + if self._return_all_text: + return asdict(vizro_plot) + if return_plot_components: return vizro_plot else: return vizro_plot.fig - - # TODO Tentative for integration test, will be updated/removed for new tests - if self._return_all_text: - return asdict(vizro_plot) From 3ec92b9788e10c06e36f4a879de233278c6621c5 Mon Sep 17 00:00:00 2001 From: AnnaXiong Date: Mon, 3 Jun 2024 17:38:30 -0400 Subject: [PATCH 14/21] changelog --- ...0240603_173503_anna_xiong_get_vizro_ai_customized_output.md} | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) rename vizro-ai/changelog.d/{20240516_135901_anna_xiong_get_vizro_ai_customized_output.md => 20240603_173503_anna_xiong_get_vizro_ai_customized_output.md} (86%) diff --git a/vizro-ai/changelog.d/20240516_135901_anna_xiong_get_vizro_ai_customized_output.md b/vizro-ai/changelog.d/20240603_173503_anna_xiong_get_vizro_ai_customized_output.md similarity index 86% rename from vizro-ai/changelog.d/20240516_135901_anna_xiong_get_vizro_ai_customized_output.md rename to vizro-ai/changelog.d/20240603_173503_anna_xiong_get_vizro_ai_customized_output.md index 7cadc52f0..dfed0e321 100644 --- a/vizro-ai/changelog.d/20240516_135901_anna_xiong_get_vizro_ai_customized_output.md +++ b/vizro-ai/changelog.d/20240603_173503_anna_xiong_get_vizro_ai_customized_output.md @@ -19,7 +19,7 @@ Uncomment the section that is right (remove the HTML comment wrapper). ### Added -- Enable feature to get all possible outputs from VizroAI.plot() using VizroAI.get_plot_outputs() ([#488](https://github.com/mckinsey/vizro/pull/488)). It returns a dictionary that contains code string, fig object, business insights, code explanation. +- Enable feature to get all possible outputs from VizroAI.plot() using by specifying 'return_plot_components = True' ([#488](https://github.com/mckinsey/vizro/pull/488)). It returns a datalcass that contains code string, fig object, business insights, and code explanation. - - - - -### Added - -- Enable feature to get customized text outputs from `VizroAI.plot` ([#488](https://github.com/mckinsey/vizro/pull/488)) - - - - - diff --git a/vizro-ai/changelog.d/20240603_173503_anna_xiong_get_vizro_ai_customized_output.md b/vizro-ai/changelog.d/20240603_173503_anna_xiong_get_vizro_ai_customized_output.md index dfed0e321..59c82e7da 100644 --- a/vizro-ai/changelog.d/20240603_173503_anna_xiong_get_vizro_ai_customized_output.md +++ b/vizro-ai/changelog.d/20240603_173503_anna_xiong_get_vizro_ai_customized_output.md @@ -19,7 +19,7 @@ Uncomment the section that is right (remove the HTML comment wrapper). ### Added -- Enable feature to get all possible outputs from VizroAI.plot() using by specifying 'return_plot_components = True' ([#488](https://github.com/mckinsey/vizro/pull/488)). It returns a datalcass that contains code string, fig object, business insights, and code explanation. +- Enable feature to get all possible outputs from VizroAI.plot() using by specifying 'return_elements = True' ([#488](https://github.com/mckinsey/vizro/pull/488)). It returns a datalcass that contains code string, figure object, business insights, and code explanation.