AttributeError: 'builtin_function_or_method' object has no attribute 'choice'

[gDanzel]

System Info

OS version: WIN11
Python version: 3.12
The current version of pandasai being used: 2.3.0

🐛 Describe the bug

I'm requesting LLM to generate plotly plot that able to randomly pick a template, and it returns the code using random module, the code snipet like:

import pandas as pd
import plotly.express as px
import random

......
template_style = random.choice(['ggplot2', 'simple_white'])
fig = px.scatter(incidence_all, x='aesoc', y='incidence_rate', size='unique_subjects', color='randarm', title='Incidence Rate', template=template_style)
fig.update_layout(xaxis_title='AESOC', yaxis_title='Incident Rate')
......

But got error at code execution

error message:

2024-11-03 01:31:06 [INFO] Executing Step 7: CodeExecution
2024-11-03 01:31:06 [ERROR] Failed with error: Traceback (most recent call last):
  File "C:\Users\Danzel\PycharmProjects\yun-ai\.venv\Lib\site-packages\pandasai\pipelines\chat\code_execution.py", line 85, in execute
    result = self.execute_code(code_to_run, code_context)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Danzel\PycharmProjects\yun-ai\.venv\Lib\site-packages\pandasai\pipelines\chat\code_execution.py", line 171, in execute_code
    exec(code, environment)
  File "<string>", line 13, in <module>
AttributeError: 'builtin_function_or_method' object has no attribute 'choice'

I did a simple debug and found it's caused by wronly importing the random.random function instead of random module, the random function dose not have choice attribute cuased the error.

in _check_imports function of code_cleaning.py, it parsed the 'import random', to dependency dic as {"module": "random", "name": "random", "alias": "random"},

    def _check_imports(self, node: Union[ast.Import, ast.ImportFrom]):
        """
        Add whitelisted imports to _additional_dependencies.

        Args:
            node (object): ast.Import or ast.ImportFrom

        Raises:
            BadImportError: If the import is not whitelisted

        """
        module = node.names[0].name if isinstance(node, ast.Import) else node.module
        library = module.split(".")[0]

        if library == "pandas":
            return

        if (
            library
            in WHITELISTED_LIBRARIES + self._config.custom_whitelisted_dependencies
        ):
            for alias in node.names:
                self._additional_dependencies.append(
                    {
                        "module": module,
                        "name": alias.name,
                        "alias": alias.asname or alias.name,
                    }
                )
            return

        if library not in WHITELISTED_BUILTINS:
            raise BadImportError(library)

then later in code cleaning, the function get_environment of optional.py, for dependency 'random', it gose getattr(import_dependency(lib["module"]), lib["name"]) instead of import_dependency(lib["module"]), because if hasattr(import_dependency(lib["module"]), lib["name"]) is true for {"module": "random", "name": "random", "alias": "random"}.

......
lib["alias"]: (
getattr(import_dependency(lib["module"]), lib["name"])
if hasattr(import_dependency(lib["module"]), lib["name"])
else import_dependency(lib["module"])
)
......

This is a bug that could potential cause problem handling import libs.

[gdcsinaptik]

@gDanzel have you added plotly to the whitelisted dependencies?
https://docs.pandas-ai.com/custom-whitelisted-dependencies