From 674ef582ffaf8761174654076db21faed86de6ca Mon Sep 17 00:00:00 2001
From: Anna Xiong <anna.xiong@quantumblack.com>
Date: Thu, 11 Jan 2024 10:31:56 -0500
Subject: [PATCH] Feat/add max debug retry parameter (#261)

---
 ...3842_anna_xiong_add_max_retry_parameter.md | 47 +++++++++++++++++++
 .../docs/pages/user_guides/run_vizro_ai.md    | 12 ++++-
 vizro-ai/src/vizro_ai/_vizro_ai.py            | 13 +++--
 3 files changed, 66 insertions(+), 6 deletions(-)
 create mode 100644 vizro-ai/changelog.d/20240110_093842_anna_xiong_add_max_retry_parameter.md

diff --git a/vizro-ai/changelog.d/20240110_093842_anna_xiong_add_max_retry_parameter.md b/vizro-ai/changelog.d/20240110_093842_anna_xiong_add_max_retry_parameter.md
new file mode 100644
index 000000000..7a571f19b
--- /dev/null
+++ b/vizro-ai/changelog.d/20240110_093842_anna_xiong_add_max_retry_parameter.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
+
+- - Added `max_debug_retry` parameter to `VizroAI.plot` to allow users to determine the maximum number of debugging attempts desired ([#261](https://github.com/mckinsey/vizro/pull/261))
+
+<!--
+### Changed
+
+- A bullet item for the Changed 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))
+
+-->
+<!--
+### 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/user_guides/run_vizro_ai.md b/vizro-ai/docs/pages/user_guides/run_vizro_ai.md
index d384ef562..99cffb8f7 100644
--- a/vizro-ai/docs/pages/user_guides/run_vizro_ai.md
+++ b/vizro-ai/docs/pages/user_guides/run_vizro_ai.md
@@ -2,7 +2,7 @@
 
 This guide offers insights into different options of running Vizro-AI, including Jupyter notebook, Python scripts, and integration into your applications.
 
-### 1. Jupyter Notebook
+### 1. Jupyter notebook
 To run Vizro-AI in jupyter, create a new cell and execute the code below to render the described visualization. You should see the chart as an output.
 
 ??? note "Note: API key"
@@ -80,3 +80,13 @@ Vizro-AI's `_get_chart_code` method returns the Python code string that can be u
 
 The returned `code_string` can be used to dynamically render charts within your application. You may have the option to encapsulate the chart within a `fig` object or convert the figure into a JSON string for further integration.
 In case you would like to use the insights or code explanation, you can use `vizro_ai._run_plot_tasks(df, ..., explain=True)`, which returns a dictionary containing the code explanation and chart insights alongside the code.
+
+### 4. `max_debug_retry` parameter in plot function
+- Default Value: 3
+- Type: int
+- Brief: By default, the `max_debug_retry` is set to 3, the function will attempt to debug errors up to three times.
+If the errors are not resolved after the maximum number of retries, the function will cease further debugging attempts.
+E.g. if you would like adjust to 5 debugging attempts, you can set `max_debug_retry = 5` in the plot function:
+    ```py
+     vizro_ai.plot(df = df, user_input = "your user input", max_debug_retry= 5)
+    ```
diff --git a/vizro-ai/src/vizro_ai/_vizro_ai.py b/vizro-ai/src/vizro_ai/_vizro_ai.py
index df1130223..7dbdc62a6 100644
--- a/vizro-ai/src/vizro_ai/_vizro_ai.py
+++ b/vizro-ai/src/vizro_ai/_vizro_ai.py
@@ -103,15 +103,18 @@ def _get_chart_code(self, df: pd.DataFrame, user_input: str) -> str:
         # TODO refine and update error handling
         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) -> Union[None, Dict[str, Any]]:
+    def plot(
+        self, df: pd.DataFrame, user_input: str, explain: bool = False, max_debug_retry: int = 3
+    ) -> Union[None, Dict[str, Any]]:
         """Plot visuals using vizro via english descriptions, english to chart translation.
 
         Args:
-            df: The dataframe to be analyzed
-            user_input: User questions or descriptions of the desired visual
-            explain: Flag to include explanation in response
+            df: The dataframe to be analyzed.
+            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`.
         """
-        output_dict = self._run_plot_tasks(df, user_input, explain=explain)
+        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")