fix: charts are now displayed again (sinaptik-ai#87)

vritansh · May 8, 2023 · 01f84d2 · 01f84d2
1 parent 5f8851a
commit 01f84d2
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/pandasai/constants.py b/pandasai/constants.py
@@ -4,3 +4,7 @@
 
 START_CODE_TAG = "<startCode>"
 END_CODE_TAG = "<endCode>"
+WHITELISTED_LIBRARIES = [
+    "numpy",
+    "matplotlib",
+]
diff --git a/pandasai/llm/base.py b/pandasai/llm/base.py
@@ -6,7 +6,7 @@
 
 import astor
 
-from ..constants import END_CODE_TAG, START_CODE_TAG
+from ..constants import END_CODE_TAG, START_CODE_TAG, WHITELISTED_LIBRARIES
 from ..exceptions import (
     APIKeyNotFoundError,
     MethodNotImplementedError,
@@ -37,7 +37,10 @@ def _remove_imports(self, code: str) -> str:
         new_body = [
             node
             for node in tree.body
-            if not isinstance(node, (ast.Import, ast.ImportFrom))
+            if not (
+                isinstance(node, (ast.Import, ast.ImportFrom))
+                and any(alias.name in WHITELISTED_LIBRARIES for alias in node.names)
+            )
         ]
         new_tree = ast.Module(body=new_body)
         return astor.to_source(new_tree).strip()