feat: reduce samples from 5 to 3 to save context

Sinaptik-AI · Nov 6, 2023 · 9386751 · 9386751
1 parent fa64a5a
commit 9386751
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/pandasai/helpers/data_sampler.py b/pandasai/helpers/data_sampler.py
@@ -25,7 +25,7 @@ def __init__(self, df: DataFrameType):
             df = df.to_pandas()
         self.df = df
 
-    def sample(self, n: int = 5) -> pd.DataFrame:
+    def sample(self, n: int = 3) -> pd.DataFrame:
         """Sample the dataframe.
 
         Args:

diff --git a/pandasai/smart_dataframe/__init__.py b/pandasai/smart_dataframe/__init__.py
@@ -428,7 +428,7 @@ def _truncate_head_columns(self, df: DataFrameType, max_size=25) -> DataFrameTyp
 
     def _get_sample_head(self) -> DataFrameType:
         head = None
-        rows_to_display = 0 if self.lake.config.enforce_privacy else 5
+        rows_to_display = 0 if self.lake.config.enforce_privacy else 3
         if self._sample_head is not None:
             head = self.sample_head
         elif not self._core._df_loaded and self.connector:

diff --git a/tests/test_query_tracker.py b/tests/test_query_tracker.py
@@ -92,7 +92,7 @@ def test_add_dataframes(
         # Check if the dataframe was added correctly
         assert len(tracker._dataframes) == 1
         assert len(tracker._dataframes[0]["headers"]) == 3
-        assert len(tracker._dataframes[0]["rows"]) == 5
+        assert len(tracker._dataframes[0]["rows"]) == 3
 
     def test_add_step(self, tracker: QueryExecTracker):
         # Create a sample step