fix: removed extra "today" from prompts + liniting (#386)

Sinaptik-AI · Jul 22, 2023 · 8ff0a9b · 8ff0a9b
1 parent 2f9cf17
commit 8ff0a9b
Show file tree

Hide file tree

Showing 14 changed files with 73 additions and 93 deletions.
diff --git a/examples/from_csv_vertexai.py b/examples/from_csv_vertexai.py
@@ -7,9 +7,9 @@
 
 df = pd.read_csv("examples/data/Loan payments data.csv")
 
-llm = GoogleVertexai(project_id="generative-ai-training",
-                     location="us-central1",
-                     model="text-bison@001")
+llm = GoogleVertexai(
+    project_id="generative-ai-training", location="us-central1", model="text-bison@001"
+)
 pandas_ai = PandasAI(llm, verbose=True, conversational=True)
 response = pandas_ai.run(df, "How many loans are from men and have been paid off?")
 print(response)

diff --git a/examples/from_dataframe_azure.py b/examples/from_dataframe_azure.py
@@ -14,7 +14,7 @@
 # The name of your deployed model
 # This will correspond to the custom name you chose for your
 # deployment when you deployed a model.
-deployment_name = 'YOUR-MODEL-DEPLOYMENT-NAME'
+deployment_name = "YOUR-MODEL-DEPLOYMENT-NAME"
 
 llm = AzureOpenAI(
     deployment_name=deployment_name,

diff --git a/examples/save_chart.py b/examples/save_chart.py
@@ -12,9 +12,9 @@
 llm = OpenAI()
 
 user_defined_path = os.getcwd()
-pandas_ai = PandasAI(llm, save_charts=True,
-                     save_charts_path=user_defined_path,
-                     verbose=True)
+pandas_ai = PandasAI(
+    llm, save_charts=True, save_charts_path=user_defined_path, verbose=True
+)
 response = pandas_ai(
     df,
     "Plot the histogram of countries showing for each the gpd,"

diff --git a/pandasai/prompts/correct_error_prompt.py b/pandasai/prompts/correct_error_prompt.py
@@ -1,6 +1,5 @@
 """ Prompt to correct Python Code on Error
 ```
-Today is {today_date}.
 You are provided with a pandas dataframe (df) with {num_rows} rows and {num_columns} columns.
 This is the metadata of the dataframe:
 {df_head}.
@@ -19,7 +18,6 @@
 code with {START_CODE_TAG} exactly and suffix the code with {END_CODE_TAG} exactly.
 ```
 """  # noqa: E501
-from datetime import date
 
 from pandasai.constants import END_CODE_TAG, START_CODE_TAG
 from .base import Prompt
@@ -29,7 +27,6 @@ class CorrectErrorPrompt(Prompt):
     """Prompt to Correct Python code on Error"""
 
     text: str = """
-Today is {today_date}.
 You are provided with a pandas dataframe (df) with {num_rows} rows and {num_columns} columns.
 This is the metadata of the dataframe:
 {df_head}.
@@ -52,5 +49,4 @@ def __init__(self, **kwargs):
             **kwargs,
             START_CODE_TAG=START_CODE_TAG,
             END_CODE_TAG=END_CODE_TAG,
-            today_date=date.today()
         )
diff --git a/pandasai/prompts/generate_python_code.py b/pandasai/prompts/generate_python_code.py
@@ -1,6 +1,5 @@
 """ Prompt to generate Python code
 ```
-Today is {today_date}.
 You are provided with a pandas dataframe (df) with {num_rows} rows and {num_columns} columns.
 This is the metadata of the dataframe:
 {df_head}.
@@ -12,7 +11,6 @@
 ```
 """  # noqa: E501
 
-from datetime import date
 
 from pandasai.constants import END_CODE_TAG, START_CODE_TAG
 from .base import Prompt
@@ -22,7 +20,6 @@ class GeneratePythonCodePrompt(Prompt):
     """Prompt to generate Python code"""
 
     text: str = """
-Today is {today_date}.
 You are provided with a pandas dataframe (df) with {num_rows} rows and {num_columns} columns.
 This is the metadata of the dataframe:
 {df_head}.
@@ -33,8 +30,5 @@ class GeneratePythonCodePrompt(Prompt):
 
     def __init__(self, **kwargs):
         super().__init__(
-            **kwargs,
-            START_CODE_TAG=START_CODE_TAG,
-            END_CODE_TAG=END_CODE_TAG,
-            today_date=date.today()
+            **kwargs, START_CODE_TAG=START_CODE_TAG, END_CODE_TAG=END_CODE_TAG
         )
diff --git a/pandasai/prompts/multiple_dataframes.py b/pandasai/prompts/multiple_dataframes.py
@@ -1,6 +1,5 @@
 """ Prompt to generate Python code for multiple dataframes """
 
-from datetime import date
 
 import pandas as pd
 
@@ -12,7 +11,6 @@ class MultipleDataframesPrompt(Prompt):
     """Prompt to generate Python code"""
 
     text: str = """
-Today is {today_date}.
 You are provided with the following pandas dataframes:"""
     instruction: str = """
 When asked about the data, your response should include a python code that describes the dataframes provided.
@@ -30,7 +28,6 @@ def __init__(self, dataframes: list[pd.DataFrame]):
 
         self.text += self.instruction
         self.text = self.text.format(
-            today_date=date.today(),
             START_CODE_TAG=START_CODE_TAG,
             END_CODE_TAG=END_CODE_TAG,
         )

diff --git a/setup.py b/setup.py
@@ -1,12 +1,8 @@
 from setuptools import setup
 
 setup(
-    name = 'pandasai-cli',
-    version = '0.1.2',
-    packages = ['pai'],
-    entry_points = {
-        'console_scripts': [
-            'pai = pai.__main__:main'
-        ]
-    }
-)
+    name="pandasai-cli",
+    version="0.1.2",
+    packages=["pai"],
+    entry_points={"console_scripts": ["pai = pai.__main__:main"]},
+)
diff --git a/tests/helpers/test_openai_info.py b/tests/helpers/test_openai_info.py
@@ -21,14 +21,16 @@ class TestOpenAIInfo:
     """Unit tests for OpenAI Info Callback"""
 
     def test_handler(self, handler: OpenAICallbackHandler) -> None:
-        response = OpenAIObject.construct_from({
-            "usage": {
-                "prompt_tokens": 2,
-                "completion_tokens": 1,
-                "total_tokens": 3,
-            },
-            "model": "gpt-35-turbo",
-        })
+        response = OpenAIObject.construct_from(
+            {
+                "usage": {
+                    "prompt_tokens": 2,
+                    "completion_tokens": 1,
+                    "total_tokens": 3,
+                },
+                "model": "gpt-35-turbo",
+            }
+        )
 
         handler(response)
         assert handler.total_tokens == 3
@@ -37,14 +39,16 @@ def test_handler(self, handler: OpenAICallbackHandler) -> None:
         assert handler.total_cost > 0
 
     def test_handler_unknown_model(self, handler: OpenAICallbackHandler) -> None:
-        response = OpenAIObject.construct_from({
-            "usage": {
-                "prompt_tokens": 2,
-                "completion_tokens": 1,
-                "total_tokens": 3,
-            },
-            "model": "foo-bar",
-        })
+        response = OpenAIObject.construct_from(
+            {
+                "usage": {
+                    "prompt_tokens": 2,
+                    "completion_tokens": 1,
+                    "total_tokens": 3,
+                },
+                "model": "foo-bar",
+            }
+        )
 
         handler(response)
         assert handler.total_tokens == 3
@@ -56,23 +60,25 @@ def test_handler_unknown_model(self, handler: OpenAICallbackHandler) -> None:
     def test_openai_callback(self, mocker):
         df = pd.DataFrame([1, 2, 3])
         llm = OpenAI(api_token="test")
-        llm_response = OpenAIObject.construct_from({
-            "choices": [
-                {
-                    "text": "```df.sum()```",
-                    "index": 0,
-                    "logprobs": None,
-                    "finish_reason": "stop",
-                    "start_text": "",
-                }
-            ],
-            "model": llm.model,
-            "usage": {
-                "prompt_tokens": 2,
-                "completion_tokens": 1,
-                "total_tokens": 3,
-            },
-        })
+        llm_response = OpenAIObject.construct_from(
+            {
+                "choices": [
+                    {
+                        "text": "```df.sum()```",
+                        "index": 0,
+                        "logprobs": None,
+                        "finish_reason": "stop",
+                        "start_text": "",
+                    }
+                ],
+                "model": llm.model,
+                "usage": {
+                    "prompt_tokens": 2,
+                    "completion_tokens": 1,
+                    "total_tokens": 3,
+                },
+            }
+        )
         mocker.patch.object(openai.ChatCompletion, "create", return_value=llm_response)
 
         pandas_ai = PandasAI(llm, enable_cache=False)

diff --git a/tests/llms/test_azure_openai.py b/tests/llms/test_azure_openai.py
@@ -26,12 +26,15 @@ def test_type_without_deployment(self):
             AzureOpenAI(api_token="test", api_base="test", api_version="test")
 
     def test_type_with_token(self):
-        assert AzureOpenAI(
-            api_token="test",
-            api_base="test",
-            api_version="test",
-            deployment_name="test"
-        ).type == "azure-openai"
+        assert (
+            AzureOpenAI(
+                api_token="test",
+                api_base="test",
+                api_version="test",
+                deployment_name="test",
+            ).type
+            == "azure-openai"
+        )
 
     def test_proxy(self):
         proxy = "http://proxy.mycompany.com:8080"
@@ -40,7 +43,7 @@ def test_proxy(self):
             api_base="test",
             api_version="test",
             deployment_name="test",
-            openai_proxy=proxy
+            openai_proxy=proxy,
         )
         assert client.openai_proxy == proxy
         assert openai.proxy["http"] == proxy
@@ -79,7 +82,7 @@ def test_completion(self, mocker):
             api_token="test",
             api_base="test",
             api_version="test",
-            deployment_name="test"
+            deployment_name="test",
         )
         result = openai.completion("Some prompt.")
 
@@ -101,7 +104,7 @@ def test_chat_completion(self, mocker):
             api_base="test",
             api_version="test",
             deployment_name="test",
-            is_chat_model=True
+            is_chat_model=True,
         )
         expected_response = {
             "choices": [

diff --git a/tests/llms/test_openai.py b/tests/llms/test_openai.py
@@ -26,10 +26,7 @@ def test_type_with_token(self):
 
     def test_proxy(self):
         proxy = "http://proxy.mycompany.com:8080"
-        client = OpenAI(
-            api_token="test",
-            openai_proxy=proxy
-        )
+        client = OpenAI(api_token="test", openai_proxy=proxy)
         assert client.openai_proxy == proxy
         assert openai.proxy["http"] == proxy
         assert openai.proxy["https"] == proxy

diff --git a/tests/prompts/test_correct_error_prompt.py b/tests/prompts/test_correct_error_prompt.py
@@ -1,6 +1,5 @@
 """Unit tests for the correct error prompt class"""
 
-from datetime import date
 
 from pandasai.prompts.correct_error_prompt import CorrectErrorPrompt
 
@@ -23,8 +22,7 @@ def test_str_with_args(self):
                     error_returned="error",
                 )
             )
-            == f"""
-Today is {date.today()}.
+            == """
 You are provided with a pandas dataframe (df) with 5 rows and 5 columns.
 This is the metadata of the dataframe:
 df.head().

diff --git a/tests/prompts/test_generate_python_code_prompt.py b/tests/prompts/test_generate_python_code_prompt.py
@@ -1,6 +1,5 @@
 """Unit tests for the generate python code prompt class"""
 
-from datetime import date
 
 from pandasai.prompts.generate_python_code import GeneratePythonCodePrompt
 
@@ -18,8 +17,7 @@ def test_str_with_args(self):
                     num_columns=5,
                 )
             )
-            == f"""
-Today is {date.today()}.
+            == """
 You are provided with a pandas dataframe (df) with 10 rows and 5 columns.
 This is the metadata of the dataframe:
 df.head().

diff --git a/tests/prompts/test_multiple_dataframes_prompt.py b/tests/prompts/test_multiple_dataframes_prompt.py
@@ -1,6 +1,5 @@
 """Unit tests for the multiple dataframes prompt class"""
 
-from datetime import date
 
 import pandas as pd
 
@@ -22,8 +21,7 @@ def test_str_with_args(self):
 
         assert (
             str(MultipleDataframesPrompt(dataframes=[df1, df2]))
-            == f"""
-Today is {date.today()}.
+            == """
 You are provided with the following pandas dataframes:
 Dataframe df1, with 5 rows and 3 columns.
 This is the metadata of the dataframe df1:

diff --git a/tests/test_pandasai.py b/tests/test_pandasai.py
@@ -1,7 +1,6 @@
 """Unit tests for the PandasAI class"""
 import logging
 import sys
-from datetime import date
 from typing import Optional
 from unittest.mock import Mock, patch
 from uuid import UUID
@@ -18,6 +17,7 @@
 from langchain.llms import OpenAI
 from pandasai.callbacks.base import StdoutCallBack
 
+
 class TestPandasAI:
     """Unit tests for the PandasAI class"""
 
@@ -131,7 +131,7 @@ def test_callback(self, pandasai):
         callback = StdoutCallBack()
         pandasai.callback = callback
 
-        #mock on_code function
+        # mock on_code function
         with patch.object(callback, "on_code") as mock_on_code:
             pandasai.run(df, "Give me sum of all gdps?")
             mock_on_code.assert_called()
@@ -191,8 +191,7 @@ def test_run_with_privacy_enforcement(self, pandasai):
         pandasai._enforce_privacy = True
         pandasai._is_conversational_answer = True
 
-        expected_prompt = f"""
-Today is {date.today()}.
+        expected_prompt = """
 You are provided with a pandas dataframe (df) with 3 rows and 1 columns.
 This is the metadata of the dataframe:
 country
@@ -255,8 +254,7 @@ def test_run_without_privacy_enforcement(self, pandasai):
         pandasai._enforce_privacy = False
         pandasai._is_conversational_answer = False
 
-        expected_prompt = f"""
-Today is {date.today()}.
+        expected_prompt = """
 You are provided with a pandas dataframe (df) with 3 rows and 1 columns.
 This is the metadata of the dataframe:
 country
@@ -569,8 +567,7 @@ def test_retry_on_error_with_single_df(self, pandasai, sample_df):
         pandasai._retry_run_code(code, e=Exception("Test error"), multiple=False)
         assert (
             pandasai.last_prompt
-            == f"""
-Today is {date.today()}.
+            == """
 You are provided with a pandas dataframe (df) with 10 rows and 3 columns.
 This is the metadata of the dataframe:
           country             gdp  happiness_index