fix: cache issue with docker (#269)

Sinaptik-AI · Jun 16, 2023 · 0ee6fd4 · 0ee6fd4
1 parent 8d07fdd
commit 0ee6fd4
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 8 deletions.
diff --git a/pandasai/__init__.py b/pandasai/__init__.py
@@ -116,7 +116,7 @@ class PandasAI:
         "num_rows": None,
         "num_columns": None,
     }
-    _cache: Cache = Cache()
+    _cache: Cache = None
     _enable_cache: bool = True
     _prompt_id: Optional[str] = None
     _middlewares: List[Middleware] = [ChartsMiddleware()]
@@ -182,12 +182,15 @@ def __init__(
         self._verbose = verbose
         self._enforce_privacy = enforce_privacy
         self._save_charts = save_charts
-        self._enable_cache = enable_cache
         self._process_id = str(uuid.uuid4())
 
         self.notebook = Notebook()
         self._in_notebook = self.notebook.in_notebook()
 
+        self._enable_cache = enable_cache
+        if self._enable_cache:
+            self._cache = Cache()
+
         if middlewares is not None:
             self.add_middlewares(*middlewares)
 
@@ -267,7 +270,7 @@ def run(
         self.log(f"Prompt ID: {self._prompt_id}")
 
         try:
-            if self._enable_cache and self._cache.get(prompt):
+            if self._enable_cache and self._cache and self._cache.get(prompt):
                 self.log("Using cached response")
                 code = self._cache.get(prompt)
             else:
@@ -325,7 +328,8 @@ def run(
                     """
                 )
 
-                self._cache.set(prompt, code)
+                if self._enable_cache and self._cache:
+                    self._cache.set(prompt, code)
 
             if show_code and self._in_notebook:
                 self.notebook.create_new_cell(code)
@@ -370,7 +374,8 @@ def clear_cache(self):
         """
         Clears the cache of the PandasAI instance.
         """
-        self._cache.clear()
+        if self._cache:
+            self._cache.clear()
 
     def __call__(
         self,

diff --git a/tests/test_pandasai.py b/tests/test_pandasai.py
@@ -339,9 +339,8 @@ def test_exception_handling(self, pandasai):
         )
         assert pandasai.last_error == "No code found in the answer."
 
-    def test_cache(self, pandasai):
-        pandasai.clear_cache()
-        pandasai._enable_cache = True
+    def test_cache(self, llm):
+        pandasai = PandasAI(llm=llm)
         pandasai._llm.call = Mock(return_value='print("Hello world")')
         assert pandasai._cache.get("How many countries are in the dataframe?") is None
         pandasai(