diff --git a/vizro-ai/changelog.d/20240131_142042_anna_xiong_refactor_ai_utils.md b/vizro-ai/changelog.d/20240131_142042_anna_xiong_refactor_ai_utils.md
new file mode 100644
index 000000000..f1f65e73c
--- /dev/null
+++ b/vizro-ai/changelog.d/20240131_142042_anna_xiong_refactor_ai_utils.md
@@ -0,0 +1,48 @@
+
+
+
+
+
+
+
+
+
diff --git a/vizro-ai/src/vizro_ai/_vizro_ai.py b/vizro-ai/src/vizro_ai/_vizro_ai.py
index 4808bff16..c3c378cf8 100644
--- a/vizro-ai/src/vizro_ai/_vizro_ai.py
+++ b/vizro-ai/src/vizro_ai/_vizro_ai.py
@@ -1,6 +1,5 @@
import logging
-import traceback
-from typing import Any, Callable, Dict, Optional, Union
+from typing import Any, Dict, Union
import pandas as pd
@@ -8,15 +7,11 @@
from vizro_ai.chains._llm_models import LLM_MODELS
from vizro_ai.components import GetCodeExplanation, GetDebugger
from vizro_ai.task_pipeline._pipeline_manager import PipelineManager
-from vizro_ai.utils import _safeguard_check
+from vizro_ai.utils.helper import DebugFailure, _debug_helper, _display_markdown_and_chart, _exec_code, _is_jupyter
logger = logging.getLogger(__name__)
-class DebugFailure(Exception):
- pass
-
-
class VizroAI:
"""Vizro-AI main class."""
@@ -136,75 +131,3 @@ def plot(
# TODO Tentative for integration test
if self._return_all_text:
return output_dict
-
-
-def _debug_helper(
- code_string: str, max_debug_retry: int, fix_chain: Callable, df: pd.DataFrame = None
-) -> Dict[bool, str]:
- """Debugging helper."""
- # TODO plug logic back into component
- retry_success = False
- last_exception = None
- for attempt in range(max_debug_retry):
- try:
- _exec_code(code=code_string, local_args={"df": df}, is_notebook_env=_is_jupyter())
- retry_success = True
- break
- except Exception as e:
- error_info = f"{traceback.format_exc()}"
- code_string = fix_chain(chain_input=error_info, code_snippet=code_string)
- last_exception = e
-
- code_string = code_string if retry_success else f"Failed to debug code {code_string}, error: {last_exception}"
-
- return {"debug_status": retry_success, "code_string": code_string}
-
-
-def _exec_code(
- code: str, local_args: Optional[Dict] = None, show_fig: bool = False, is_notebook_env: bool = True
-) -> None:
- """Execute code in notebook with correct namespace."""
- from IPython import get_ipython
-
- if show_fig and "\nfig.show()" not in code:
- code += "\nfig.show()"
- elif not show_fig:
- code = code.replace("fig.show()", "")
- namespace = get_ipython().user_ns if is_notebook_env else globals()
- if local_args:
- namespace.update(local_args)
- _safeguard_check(code)
-
- exec(code, namespace) # nosec
-
-
-# Taken from rich.console. See https://github.com/Textualize/rich.
-def _is_jupyter() -> bool: # pragma: no cover
- """Checks if we're running in a Jupyter notebook."""
- try:
- from IPython import get_ipython
- except NameError:
- return False
- ipython = get_ipython()
- shell = ipython.__class__.__name__
- if "google.colab" in str(ipython.__class__) or shell == "ZMQInteractiveShell":
- return True # Jupyter notebook or qtconsole
- elif shell == "TerminalInteractiveShell":
- return False # Terminal running IPython
- else:
- return False # Other type (?)
-
-
-def _display_markdown_and_chart(df: pd.DataFrame, code_snippet: str, biz_insights: str, code_explain: str) -> None:
- # TODO change default test str to other
- """Display chart and Markdown format description in jupyter."""
- try:
- # pylint: disable=import-outside-toplevel
- from IPython.display import Markdown, display
- except Exception as exc:
- raise ImportError("Please install IPython before proceeding in jupyter environment.") from exc
- # TODO clean up the formatting markdown code
- markdown_code = f"```\n{code_snippet}\n```"
- output_text = f"Insights:
\n\n{biz_insights}\n
Code:
\n\n{code_explain}\n{markdown_code}"
- display(Markdown(output_text))
- _exec_code(code_snippet, local_args={"df": df}, show_fig=True, is_notebook_env=_is_jupyter())
diff --git a/vizro-ai/src/vizro_ai/utils/__init__.py b/vizro-ai/src/vizro_ai/utils/__init__.py
index 3b664f31c..e69de29bb 100644
--- a/vizro-ai/src/vizro_ai/utils/__init__.py
+++ b/vizro-ai/src/vizro_ai/utils/__init__.py
@@ -1,3 +0,0 @@
-from .safeguard import _safeguard_check
-
-__all__ = ["_safeguard_check"]
diff --git a/vizro-ai/src/vizro_ai/utils/helper.py b/vizro-ai/src/vizro_ai/utils/helper.py
new file mode 100644
index 000000000..4bf904707
--- /dev/null
+++ b/vizro-ai/src/vizro_ai/utils/helper.py
@@ -0,0 +1,85 @@
+"""Helper Functions For Vizro AI."""
+import traceback
+from typing import Callable, Dict, Optional
+
+import pandas as pd
+
+from .safeguard import _safeguard_check
+
+
+# Taken from rich.console. See https://github.com/Textualize/rich.
+def _is_jupyter() -> bool: # pragma: no cover
+ """Checks if we're running in a Jupyter notebook."""
+ try:
+ from IPython import get_ipython
+ except NameError:
+ return False
+ ipython = get_ipython()
+ shell = ipython.__class__.__name__
+ if "google.colab" in str(ipython.__class__) or shell == "ZMQInteractiveShell":
+ return True # Jupyter notebook or qtconsole
+ elif shell == "TerminalInteractiveShell":
+ return False # Terminal running IPython
+ else:
+ return False # Other type (?)
+
+
+def _debug_helper(
+ code_string: str, max_debug_retry: int, fix_chain: Callable, df: pd.DataFrame = None
+) -> Dict[bool, str]:
+ """Debugging helper."""
+ retry_success = False
+ last_exception = None
+ is_jupyter = _is_jupyter()
+ for attempt in range(max_debug_retry):
+ try:
+ _exec_code(code=code_string, local_args={"df": df}, is_notebook_env=is_jupyter)
+ retry_success = True
+ break
+ except Exception as e:
+ error_info = f"{traceback.format_exc()}"
+ code_string = fix_chain(chain_input=error_info, code_snippet=code_string)
+ last_exception = e
+
+ code_string = code_string if retry_success else f"Failed to debug code {code_string}, error: {last_exception}"
+
+ return {"debug_status": retry_success, "code_string": code_string}
+
+
+def _exec_code(
+ code: str, local_args: Optional[Dict] = None, show_fig: bool = False, is_notebook_env: bool = True
+) -> None:
+ """Execute code in notebook with correct namespace."""
+ from IPython import get_ipython
+
+ if show_fig and "\nfig.show()" not in code:
+ code += "\nfig.show()"
+ elif not show_fig:
+ code = code.replace("fig.show()", "")
+ namespace = get_ipython().user_ns if is_notebook_env else globals()
+ if local_args:
+ namespace.update(local_args)
+ _safeguard_check(code)
+
+ exec(code, namespace) # nosec
+
+
+def _display_markdown_and_chart(df: pd.DataFrame, code_snippet: str, biz_insights: str, code_explain: str) -> None:
+ # TODO change default test str to other
+ """Display chart and Markdown format description in jupyter."""
+ try:
+ # pylint: disable=import-outside-toplevel
+ from IPython.display import Markdown, display
+ except Exception as exc:
+ raise ImportError("Please install IPython before proceeding in jupyter environment.") from exc
+ # TODO clean up the formatting markdown code to render in jupyter
+ markdown_code = f"```\n{code_snippet}\n```"
+ output_text = f"Insights:
\n\n{biz_insights}\n
Code:
\n\n{code_explain}\n{markdown_code}"
+ display(Markdown(output_text))
+ _exec_code(code_snippet, local_args={"df": df}, show_fig=True, is_notebook_env=_is_jupyter())
+
+
+class DebugFailure(Exception):
+ """Debug Failure."""
+
+ pass
diff --git a/vizro-ai/tests/unit/vizro-ai/utils/test_safeguard_code.py b/vizro-ai/tests/unit/vizro-ai/utils/test_safeguard_code.py
index 533f03801..9bd58121c 100644
--- a/vizro-ai/tests/unit/vizro-ai/utils/test_safeguard_code.py
+++ b/vizro-ai/tests/unit/vizro-ai/utils/test_safeguard_code.py
@@ -1,7 +1,7 @@
import re
import pytest
-from vizro_ai.utils import _safeguard_check
+from vizro_ai.utils.safeguard import _safeguard_check
class TestMaliciousImports: