From 492ee80591e39369f29c9d3d334eaf3fe25eb034 Mon Sep 17 00:00:00 2001
From: nadijagraca <nadija_ratkusic_graca@external.mckinsey.com>
Date: Thu, 5 Dec 2024 11:00:52 +0100
Subject: [PATCH] update safeguards

---
 ...tkusic_graca_update_safeguard_whitelist.md | 48 +++++++++++++++++++
 .../src/vizro_ai/plot/_response_models.py     |  3 +-
 .../src/vizro_ai/plot/_utils/_constants.py    | 10 ++--
 3 files changed, 55 insertions(+), 6 deletions(-)
 create mode 100644 vizro-ai/changelog.d/20241205_104116_nadija_ratkusic_graca_update_safeguard_whitelist.md

diff --git a/vizro-ai/changelog.d/20241205_104116_nadija_ratkusic_graca_update_safeguard_whitelist.md b/vizro-ai/changelog.d/20241205_104116_nadija_ratkusic_graca_update_safeguard_whitelist.md
new file mode 100644
index 000000000..7c0d58d4f
--- /dev/null
+++ b/vizro-ai/changelog.d/20241205_104116_nadija_ratkusic_graca_update_safeguard_whitelist.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))
+
+-->
diff --git a/vizro-ai/src/vizro_ai/plot/_response_models.py b/vizro-ai/src/vizro_ai/plot/_response_models.py
index 17333373a..8daa0f4de 100644
--- a/vizro-ai/src/vizro_ai/plot/_response_models.py
+++ b/vizro-ai/src/vizro_ai/plot/_response_models.py
@@ -196,7 +196,8 @@ def _test_execute_chart_code(v, values):
                 fig = custom_chart(data_frame.sample(10, replace=True))
             except Exception as e:
                 raise ValueError(
-                    f"Produced code execution failed the following error: <{e}>. Please check the code and try again."
+                    f"Produced code execution failed the following error: <{e}>. Please check the code and try again, "
+                    f"alternatively try with a more powerful model."
                 )
             assert isinstance(
                 fig, go.Figure
diff --git a/vizro-ai/src/vizro_ai/plot/_utils/_constants.py b/vizro-ai/src/vizro_ai/plot/_utils/_constants.py
index 86935f633..30db3eaf8 100644
--- a/vizro-ai/src/vizro_ai/plot/_utils/_constants.py
+++ b/vizro-ai/src/vizro_ai/plot/_utils/_constants.py
@@ -22,7 +22,7 @@
     # '__build_class__',
     # '__import__',
     "abs",
-    # 'all',
+    "all",
     "any",
     # 'ascii',
     # 'bin',
@@ -45,7 +45,7 @@
     # 'input',
     # 'isinstance',
     # 'issubclass',
-    # 'iter',
+    "iter",
     # 'aiter',
     "len",
     # 'locals',
@@ -60,7 +60,7 @@
     # 'repr',
     "round",
     # 'setattr',
-    # 'sorted',
+    "sorted",
     "sum",
     # 'vars',
     "None",
@@ -76,7 +76,7 @@
     # 'complex',
     "dict",
     "enumerate",
-    # 'filter',
+    "filter",
     "float",
     # 'frozenset',
     # 'property',
@@ -85,7 +85,7 @@
     "map",
     # 'object',
     "range",
-    # 'reversed',
+    "reversed",
     "set",
     # 'slice',
     # 'staticmethod',