feat: allow user to pass custom whitelisted libraries

docs/custom-optional-dependencies.md

@@ -0,0 +1,11 @@
+# Custom optional dependencies
+
+By default, `pandasai` only allows to run code that uses some whitelisted modules. This is to prevent malicious code from being executed on the server or locally. However, it is possible to add custom modules to the whitelist. This can be done by passing a list of modules to the `custom_whitelisted_dependencies` parameter when instantiating the `PandasAI` object:
+
+```python
+pandas_ai = PandasAI(llm, custom_whitelisted_dependencies=["my_custom_library"])
+```
+
+The `custom_whitelisted_dependencies` parameter accepts a list of strings, where each string is the name of a module. The module must be installed in the environment where `pandasai` is running.
+
+Please, make sure you have installed the module in the environment where `pandasai` is running. Otherwise, you will get an error when trying to run the code.

mkdocs.yml

@@ -13,6 +13,7 @@ nav:
   - Getting Started: getting-started.md
   - Cache: cache.md
   - Middlewares: middlewares.md
+  - Custom Optional Arguments: custom_optional_arguments.md
   - Command Line Tool: pai_cli.md
   - API:
       - API/pandasai.md

pandasai/__init__.py

@@ -120,6 +120,7 @@ class PandasAI:
     _prompt_id: Optional[str] = None
     _middlewares: List[Middleware] = [ChartsMiddleware()]
     _additional_dependencies: List[dict] = []
+    _custom_whitelisted_dependencies: List[str] = []
     last_code_generated: Optional[str] = None
     last_run_code: Optional[str] = None
     code_output: Optional[str] = None
@@ -134,6 +135,7 @@ def __init__(
         save_charts=False,
         enable_cache=True,
         middlewares=None,
+        custom_whitelisted_dependencies=None,
     ):
         """
 
@@ -149,6 +151,11 @@ def __init__(
             Default to False
             save_charts (bool): Save the charts generated in the notebook.
             Default to False
+            enable_cache (bool): Enable the cache to store the results.
+            Default to True
+            middlewares (list): List of middlewares to be used. Default to None
+            custom_whitelisted_dependencies (list): List of custom dependencies to
+            be used. Default to None
         """
 
         # configure the logging
@@ -183,6 +190,9 @@ def __init__(
         if middlewares is not None:
             self.add_middlewares(*middlewares)
 
+        if custom_whitelisted_dependencies is not None:
+            self._custom_whitelisted_dependencies = custom_whitelisted_dependencies
+
     def conversational_answer(self, question: str, answer: str) -> str:
         """
         Returns the answer in conversational form about the resultant data.
@@ -401,7 +411,7 @@ def _check_imports(self, node: Union[ast.Import, ast.ImportFrom]):
         if library == "pandas":
             return
 
-        if library in WHITELISTED_LIBRARIES:
+        if library in WHITELISTED_LIBRARIES + self._custom_whitelisted_dependencies:
             for alias in node.names:
                 self._additional_dependencies.append(
                     {

tests/test_pandasai.py

@@ -385,3 +385,16 @@ def test_middlewares(self, pandasai, test_middleware):
             == "Overwritten by middleware"
         )
         assert middleware.has_run
+
+    def test_custom_whitelisted_dependencies(self, pandasai):
+        code = """
+import my_custom_library
+my_custom_library.do_something()
+"""
+        pandasai._llm._output = code
+
+        with pytest.raises(BadImportError):
+            pandasai._clean_code(code)
+
+        pandasai._custom_whitelisted_dependencies = ["my_custom_library"]
+        assert pandasai._clean_code(code) == "my_custom_library.do_something()"