Refactor ai utils (#297)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Alexey Snigir <[email protected]>

vizro-ai/changelog.d/20240131_142042_anna_xiong_refactor_ai_utils.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))
+
+-->

vizro-ai/src/vizro_ai/_vizro_ai.py

@@ -1,22 +1,17 @@
 import logging
-import traceback
-from typing import Any, Callable, Dict, Optional, Union
+from typing import Any, Dict, Union
 
 import pandas as pd
 
 from vizro_ai.chains import ModelConstructor
 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"<h4>Insights:</h4>\n\n{biz_insights}\n<br><br><h4>Code:</h4>\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())

vizro-ai/src/vizro_ai/utils/__init__.py

@@ -1,3 +0,0 @@
-from .safeguard import _safeguard_check
-
-__all__ = ["_safeguard_check"]

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"<h4>Insights:</h4>\n\n{biz_insights}\n<br><br><h4>Code:</h4>\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

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: