diff --git a/vizro-ai/changelog.d/20240725_132140_nadija_ratkusic_graca_vizroai_update_docs.md b/vizro-ai/changelog.d/20240725_132140_nadija_ratkusic_graca_vizroai_update_docs.md
new file mode 100644
index 000000000..f1f65e73c
--- /dev/null
+++ b/vizro-ai/changelog.d/20240725_132140_nadija_ratkusic_graca_vizroai_update_docs.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/run-vizro-ai.md b/vizro-ai/docs/pages/user-guides/run-vizro-ai.md
index e26209774..f940641d8 100644
--- a/vizro-ai/docs/pages/user-guides/run-vizro-ai.md
+++ b/vizro-ai/docs/pages/user-guides/run-vizro-ai.md
@@ -81,32 +81,32 @@ There are two ways to integrate Vizro-AI into an application, directly and by ac
             ```
 
 
-2. Vizro-AI's `_get_chart_code` method returns a string of Python code that manipulates the data and creates the visualization. Vizro-AI validates the code to ensure that it is executable and can be integrated.
+2. When the `return_elements` argument of VizroAI's `plot` method is set to `True`, the method returns a `PlotOutputs` data class which contains all possible `VizroAI.plot()` outputs.
+By setting `return_elements=True` you can access code or the figure object. Vizro-AI validates the code to ensure that it is executable and can be integrated.
 
-    !!! example "Application integration via chart code"
+    !!! example "Accessing outputs from PlotOutputs data class"
 
         === "app.py"
-            ```py
-            import vizro.plotly.express as px
-            from vizro_ai import VizroAI
-
-            vizro_ai = VizroAI()
+           ```py
+           import vizro_ai
+           from vizro_ai import VizroAI
+           import vizro.plotly.express as px
+           from dotenv import load_dotenv
 
-            df = px.data.gapminder()
-            code_string = vizro_ai._get_chart_code(df, "describe life expectancy per continent over time")
-            ```
-        === "code_string"
-            [![ResultCode]][ResultCode]
+           load_dotenv()
 
-        [ResultCode]: ../../assets/user_guides/code_string_app_integration.png
+           df = px.data.gapminder()
+           vizro_ai = VizroAI()
 
-    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.
+           plot_outputs = vizro_ai.plot(df, "describe life expectancy per continent over time", explain=True, return_elements=True)
+           fig = plot_outputs.figure
+           code_string = plot_outputs.code
+           ```
 
-    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.
 
 ### How to use `max_debug_retry` parameter in plot function
 - Default Value: 3
-- Type: int
+- Type: `int`
 - Brief: By default, the `max_debug_retry` is set to 3, the function will try to debug errors up to three times.
 If the errors are not resolved after the maximum number of retries, the function will stop further debugging retries.
 For example, if you would like adjust to 5 retries, you can set `max_debug_retry = 5` in the plot function:
@@ -114,3 +114,17 @@ For example, if you would like adjust to 5 retries, you can set `max_debug_retry
 ```py
 vizro_ai.plot(df = df, user_input = "your user input", max_debug_retry= 5)
 ```
+
+### How to use `return_elements` parameter in plot function
+- Default Value: False
+- Type: `bool`
+- Brief: By default, the `return_elements` is set to `False`, and the `VizroAI.plot()` returns a `plotly.graph_objects` object. If `return_elements` is set to `True` `VizroAI.plot()` returns the `PlotOutputs` data class.
+The `PlotOutputs` data class is designed to encapsulate all possible outputs generated by the `VizroAI.plot()` method.
+
+Attributes of `PlotOutputs`:
+
+- **`code`**: A string representing of the Python code that manipulates the data and creates the visualization.
+- **`figure`**: A [`CapturedCallable`](https://vizro.readthedocs.io/en/stable/pages/API-reference/models/#vizro.models.types.CapturedCallable) object from [`Vizro`](https://vizro.readthedocs.io/en/stable/), representing the visual plot generated by `VizroAI.plot()`. This object is designed to run immediately in the Vizro dashboard, but otherwise, it behaves like a plotly `go.Figure`.
+- **`business_insights`**: A string containing high-level business insights derived from the plot. `business_insights` is only available if `explain=True`.
+- **`code_explanation`**: A string offering a detailed explanation of the code used to produce the plot. `code_explanation` is only available if `explain=True`.
+-