From 64d8602b801e550f8fb9bd0c1cc76863b95704e6 Mon Sep 17 00:00:00 2001
From: Anna Xiong <anna.xiong@quantumblack.com>
Date: Tue, 14 May 2024 12:13:12 -0400
Subject: [PATCH] [Docs]fix vizro ai chart example (#472)

Co-authored-by: Lingyi Zhang <lingyi_zhang@mckinsey.com>
---
 ...9_anna_xiong_fix_vizro_ai_chart_example.md | 48 +++++++++++++++++++
 .../add-generated-chart-usecase.md            |  6 +--
 .../pages/user-guides/customize-vizro-ai.md   |  2 +-
 vizro-ai/examples/example.ipynb               |  2 +-
 vizro-ai/examples/example.py                  |  2 +-
 .../vizro_ai/components/chart_selection.py    |  2 +-
 .../vizro_ai/components/code_validation.py    |  2 +-
 .../vizro_ai/components/custom_chart_wrap.py  |  2 +-
 .../vizro_ai/components/dataframe_craft.py    |  2 +-
 .../src/vizro_ai/components/visual_code.py    |  2 +-
 vizro-ai/tests/integration/test_example.py    |  2 +-
 11 files changed, 60 insertions(+), 12 deletions(-)
 create mode 100644 vizro-ai/changelog.d/20240510_140409_anna_xiong_fix_vizro_ai_chart_example.md

diff --git a/vizro-ai/changelog.d/20240510_140409_anna_xiong_fix_vizro_ai_chart_example.md b/vizro-ai/changelog.d/20240510_140409_anna_xiong_fix_vizro_ai_chart_example.md
new file mode 100644
index 000000000..f1f65e73c
--- /dev/null
+++ b/vizro-ai/changelog.d/20240510_140409_anna_xiong_fix_vizro_ai_chart_example.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/docs/pages/user-guides/add-generated-chart-usecase.md b/vizro-ai/docs/pages/user-guides/add-generated-chart-usecase.md
index 2241de7ea..286357b4a 100644
--- a/vizro-ai/docs/pages/user-guides/add-generated-chart-usecase.md
+++ b/vizro-ai/docs/pages/user-guides/add-generated-chart-usecase.md
@@ -14,7 +14,7 @@ This guide explains the different ways in which you can add a chart generated by
             ```py
             import vizro_ai
             from vizro_ai import VizroAI
-            import plotly.express as px
+            import vizro.plotly.express as px
             from dotenv import load_dotenv
 
             load_dotenv()
@@ -64,7 +64,7 @@ This object must then be passed to the `figure` argument of the Vizro [Graph](ht
         import vizro.models as vm
         from vizro.models.types import capture
         import pandas as pd
-        import plotly.express as px
+        import vizro.plotly.express as px
 
 
         @capture('graph')
@@ -120,7 +120,7 @@ Executing the code below yields the identical dashboard as the example above.
         from vizro import Vizro
         import vizro.models as vm
         import pandas as pd
-        import plotly.express as px
+        import vizro.plotly.express as px
         import vizro_ai
         from vizro_ai import VizroAI
 
diff --git a/vizro-ai/docs/pages/user-guides/customize-vizro-ai.md b/vizro-ai/docs/pages/user-guides/customize-vizro-ai.md
index c7105a699..0ac4e4679 100644
--- a/vizro-ai/docs/pages/user-guides/customize-vizro-ai.md
+++ b/vizro-ai/docs/pages/user-guides/customize-vizro-ai.md
@@ -51,7 +51,7 @@ To ensure a deterministic answer to our queries, we've set the temperature to 0.
         from vizro_ai import VizroAI
         from langchain_openai import ChatOpenAI
 
-        import plotly.express as px
+        import vizro.plotly.express as px
         df = px.data.gapminder()
 
         llm = ChatOpenAI(
diff --git a/vizro-ai/examples/example.ipynb b/vizro-ai/examples/example.ipynb
index a257a2ea8..372a958c3 100644
--- a/vizro-ai/examples/example.ipynb
+++ b/vizro-ai/examples/example.ipynb
@@ -39,7 +39,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import plotly.express as px\n",
+    "import vizro.plotly.express as px\n",
     "df = px.data.gapminder()\n",
     "df"
    ]
diff --git a/vizro-ai/examples/example.py b/vizro-ai/examples/example.py
index b62e66ea1..1789f0256 100644
--- a/vizro-ai/examples/example.py
+++ b/vizro-ai/examples/example.py
@@ -1,6 +1,6 @@
 """AI plot example."""
 
-import plotly.express as px
+import vizro.plotly.express as px
 from vizro_ai import VizroAI
 
 vizro_ai = VizroAI()
diff --git a/vizro-ai/src/vizro_ai/components/chart_selection.py b/vizro-ai/src/vizro_ai/components/chart_selection.py
index 7743880ad..a14353639 100644
--- a/vizro-ai/src/vizro_ai/components/chart_selection.py
+++ b/vizro-ai/src/vizro_ai/components/chart_selection.py
@@ -106,7 +106,7 @@ def _chart_to_use(load_args) -> str:
 
 
 if __name__ == "__main__":
-    import plotly.express as px
+    import vizro.plotly.express as px
 
     from vizro_ai.chains._llm_models import _get_llm_model
 
diff --git a/vizro-ai/src/vizro_ai/components/code_validation.py b/vizro-ai/src/vizro_ai/components/code_validation.py
index 384b6a178..ab2cbd650 100644
--- a/vizro-ai/src/vizro_ai/components/code_validation.py
+++ b/vizro-ai/src/vizro_ai/components/code_validation.py
@@ -86,7 +86,7 @@ def run(self, code_snippet: str, chain_input: str = "") -> str:
 
 
 if __name__ == "__main__":
-    import plotly.express as px
+    import vizro.plotly.express as px
 
     from vizro_ai.chains._llm_models import _get_llm_model
 
diff --git a/vizro-ai/src/vizro_ai/components/custom_chart_wrap.py b/vizro-ai/src/vizro_ai/components/custom_chart_wrap.py
index befb1aa04..5ff938368 100644
--- a/vizro-ai/src/vizro_ai/components/custom_chart_wrap.py
+++ b/vizro-ai/src/vizro_ai/components/custom_chart_wrap.py
@@ -103,7 +103,7 @@ def _add_capture_code(code_string: str) -> str:
 
 
 if __name__ == "__main__":
-    import plotly.express as px
+    import vizro.plotly.express as px
 
     from vizro_ai.chains._llm_models import _get_llm_model
 
diff --git a/vizro-ai/src/vizro_ai/components/dataframe_craft.py b/vizro-ai/src/vizro_ai/components/dataframe_craft.py
index 5f63a19f5..37741a4f6 100755
--- a/vizro-ai/src/vizro_ai/components/dataframe_craft.py
+++ b/vizro-ai/src/vizro_ai/components/dataframe_craft.py
@@ -123,7 +123,7 @@ def _format_dataframe_string(s: str) -> str:
 
 
 if __name__ == "__main__":
-    import plotly.express as px
+    import vizro.plotly.express as px
 
     from vizro_ai.chains._llm_models import _get_llm_model
 
diff --git a/vizro-ai/src/vizro_ai/components/visual_code.py b/vizro-ai/src/vizro_ai/components/visual_code.py
index 7314648f4..137f15bc7 100644
--- a/vizro-ai/src/vizro_ai/components/visual_code.py
+++ b/vizro-ai/src/vizro_ai/components/visual_code.py
@@ -101,7 +101,7 @@ def _clean_visual_code(raw_code: str) -> str:
 
 
 if __name__ == "__main__":
-    import plotly.express as px
+    import vizro.plotly.express as px
 
     from vizro_ai.chains._llm_models import _get_llm_model
 
diff --git a/vizro-ai/tests/integration/test_example.py b/vizro-ai/tests/integration/test_example.py
index 770d56ba6..a81ddb79b 100644
--- a/vizro-ai/tests/integration/test_example.py
+++ b/vizro-ai/tests/integration/test_example.py
@@ -1,4 +1,4 @@
-import plotly.express as px
+import vizro.plotly.express as px
 from hamcrest import all_of, any_of, assert_that, contains_string, equal_to
 from vizro_ai import VizroAI