fix

mckinsey · Jun 4, 2024 · 477ad82 · 477ad82
1 parent 4c69844
commit 477ad82
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 49 deletions.
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
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
@@ -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.
 
 <!--
 ### Changed

diff --git a/vizro-ai/src/vizro_ai/_vizro_ai.py b/vizro-ai/src/vizro_ai/_vizro_ai.py
@@ -155,6 +155,8 @@ def plot(  # pylint: disable=too-many-arguments  # noqa: PLR0913
 
         # TODO Tentative for integration test, will be updated/removed for new tests
         if self._return_all_text:
-            return asdict(vizro_plot)
+            output_dict = asdict(vizro_plot)
+            output_dict["code_string"] = vizro_plot.code
+            return output_dict
 
         return vizro_plot if return_elements else vizro_plot.figure