From 538d176af8fb56fcbba856458a9dea3caeb25dea Mon Sep 17 00:00:00 2001
From: Maximilian Schulz <maximilian_schulz@mckinsey.com>
Date: Fri, 10 Nov 2023 22:26:37 +0100
Subject: [PATCH 1/3] Introduce custom exception

---
 vizro-ai/src/vizro_ai/_vizro_ai.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/vizro-ai/src/vizro_ai/_vizro_ai.py b/vizro-ai/src/vizro_ai/_vizro_ai.py
index 622f4421f..572ffd2d2 100644
--- a/vizro-ai/src/vizro_ai/_vizro_ai.py
+++ b/vizro-ai/src/vizro_ai/_vizro_ai.py
@@ -19,6 +19,10 @@
 logger = logging.getLogger(__name__)
 
 
+class DebugFailure(Exception):
+    pass
+
+
 class VizroAI:
     """Vizro-AI main class."""
 
@@ -113,12 +117,11 @@ def plot(self, df: pd.DataFrame, user_input: str, explain: bool = False) -> Unio
         code_explanation = output_dict.get("code_explanation")
 
         if code_string.startswith("Failed to debug code"):
-            logger.warning(
+            raise DebugFailure(
                 "Chart creation failed. Retry debugging has reached maximum limit, fallout response is \
-                            provided. "
+                            provided. Try to rephrase the prompt, or try to select a different model. "
                 + code_string
             )
-            return
         if not explain:
             _exec_code(code=code_string, local_args={"df": df}, show_fig=True, is_notebook_env=_is_jupyter())
         if explain:

From afb3630a1e502f94a4a7d054261ccc7851182ab5 Mon Sep 17 00:00:00 2001
From: Maximilian Schulz <maximilian_schulz@mckinsey.com>
Date: Fri, 10 Nov 2023 22:39:32 +0100
Subject: [PATCH 2/3] Provide better failure message

---
 vizro-ai/src/vizro_ai/_vizro_ai.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/vizro-ai/src/vizro_ai/_vizro_ai.py b/vizro-ai/src/vizro_ai/_vizro_ai.py
index 572ffd2d2..df1130223 100644
--- a/vizro-ai/src/vizro_ai/_vizro_ai.py
+++ b/vizro-ai/src/vizro_ai/_vizro_ai.py
@@ -118,9 +118,8 @@ def plot(self, df: pd.DataFrame, user_input: str, explain: bool = False) -> Unio
 
         if code_string.startswith("Failed to debug code"):
             raise DebugFailure(
-                "Chart creation failed. Retry debugging has reached maximum limit, fallout response is \
-                            provided. Try to rephrase the prompt, or try to select a different model. "
-                + code_string
+                "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 not explain:
             _exec_code(code=code_string, local_args={"df": df}, show_fig=True, is_notebook_env=_is_jupyter())

From 07d332de2d1d6701e48c5d394469850d7aa3bd10 Mon Sep 17 00:00:00 2001
From: Maximilian Schulz <maximilian_schulz@mckinsey.com>
Date: Fri, 10 Nov 2023 22:40:01 +0100
Subject: [PATCH 3/3] Changelog

---
 ...imilian_schulz_improve_AI_error_message.md | 48 +++++++++++++++++++
 1 file changed, 48 insertions(+)
 create mode 100644 vizro-ai/changelog.d/20231110_223949_maximilian_schulz_improve_AI_error_message.md

diff --git a/vizro-ai/changelog.d/20231110_223949_maximilian_schulz_improve_AI_error_message.md b/vizro-ai/changelog.d/20231110_223949_maximilian_schulz_improve_AI_error_message.md
new file mode 100644
index 000000000..f1f65e73c
--- /dev/null
+++ b/vizro-ai/changelog.d/20231110_223949_maximilian_schulz_improve_AI_error_message.md
@@ -0,0 +1,48 @@
+<!--
+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
+
+- 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))
+
+-->