feat: reduce query token consumption (#270)

* Update __init__.py Converted the dataframe to json for lower token usage. * test: fix tests --------- Co-authored-by: Gabriele Venturi <[email protected]>
Sinaptik-AI · Jun 25, 2023 · 09a41f0 · 09a41f0
1 parent a09101f
commit 09a41f0
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 8 deletions.
diff --git a/pandasai/__init__.py b/pandasai/__init__.py
@@ -334,6 +334,7 @@ def run(
                     df_head = data_frame.head(rows_to_display)
                     if anonymize_df:
                         df_head = anonymize_dataframe_head(df_head)
+                    df_head=df_head.to_csv(index=False)
 
                     generate_code_instruction = self._non_default_prompts.get(
                         "generate_python_code", GeneratePythonCodePrompt

diff --git a/tests/test_pandasai.py b/tests/test_pandasai.py
@@ -184,9 +184,8 @@ def test_run_with_privacy_enforcement(self, pandasai):
 Today is {date.today()}.
 You are provided with a pandas dataframe (df) with 3 rows and 1 columns.
 This is the metadata of the dataframe:
-Empty DataFrame
-Columns: [country]
-Index: [].
+country
+.
 
 When asked about the data, your response should include a python code that describes the dataframe `df`.
 Using the provided dataframe, df, return the python code and make sure to prefix the requested python code with <startCode> exactly and suffix the code with <endCode> exactly to get the answer to the following question:
@@ -246,10 +245,11 @@ def test_run_without_privacy_enforcement(self, pandasai):
 Today is {date.today()}.
 You are provided with a pandas dataframe (df) with 3 rows and 1 columns.
 This is the metadata of the dataframe:
-          country
-0   United States
-1  United Kingdom
-2          France.
+country
+United States
+United Kingdom
+France
+.
 
 When asked about the data, your response should include a python code that describes the dataframe `df`.
 Using the provided dataframe, df, return the python code and make sure to prefix the requested python code with <startCode> exactly and suffix the code with <endCode> exactly to get the answer to the following question:
@@ -663,7 +663,9 @@ def test_replace_generate_code_prompt(self, llm):
         expected_last_prompt = (
             str(
                 replacement_prompt(
-                    num_rows=df.shape[0], num_columns=df.shape[1], df_head=df.head()
+                    num_rows=df.shape[0],
+                    num_columns=df.shape[1],
+                    df_head=df.head().to_csv(index=False),
                 )
             )
             + question