fix: removed deprecated df._data

sinaptik-ai · Jan 28, 2025 · 71279d5 · 71279d5
1 parent a564920
commit 71279d5
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 27 deletions.
diff --git a/pandasai/data_loader/loader.py b/pandasai/data_loader/loader.py
@@ -74,10 +74,10 @@ def load(
             df = self._apply_transformations(df)
 
             # Convert to pandas DataFrame while preserving internal data
-            df = pd.DataFrame(df._data)
+            df = pd.DataFrame(df)
 
             return DataFrame(
-                df._data,
+                df,
                 schema=self.schema,
                 name=self.schema.name,
                 description=self.schema.description,

diff --git a/tests/unit_tests/agent/test_agent_llm_judge.py b/tests/unit_tests/agent/test_agent_llm_judge.py
@@ -201,7 +201,6 @@ def test_combined_questions_with_type(self, question):
 
         assert evaluation_response.score > 5, evaluation_response.justification
 
-    @pytest.mark.final
     def test_average_score(self):
         if self.evaluation_scores:
             average_score = sum(self.evaluation_scores) / len(self.evaluation_scores)

diff --git a/tests/unit_tests/prompts/test_sql_prompt.py b/tests/unit_tests/prompts/test_sql_prompt.py
@@ -19,28 +19,26 @@ class TestGeneratePythonCodeWithSQLPrompt:
     @pytest.mark.parametrize(
         "output_type,output_type_template",
         [
-            *[
-                (
-                    "",
-                    """type (possible values "string", "number", "dataframe", "plot"). Examples: { "type": "string", "value": f"The highest salary is {highest_salary}." } or { "type": "number", "value": 125 } or { "type": "dataframe", "value": pd.DataFrame({...}) } or { "type": "plot", "value": "temp_chart.png" }""",
-                ),
-                (
-                    "number",
-                    """type (must be "number"), value must int. Example: { "type": "number", "value": 125 }""",
-                ),
-                (
-                    "dataframe",
-                    """type (must be "dataframe"), value must be pd.DataFrame or pd.Series. Example: { "type": "dataframe", "value": pd.DataFrame({...}) }""",
-                ),
-                (
-                    "plot",
-                    """type (must be "plot"), value must be string. Example: { "type": "plot", "value": "temp_chart.png" }""",
-                ),
-                (
-                    "string",
-                    """type (must be "string"), value must be string. Example: { "type": "string", "value": f"The highest salary is {highest_salary}." }""",
-                ),
-            ]
+            (
+                "",
+                """type (possible values "string", "number", "dataframe", "plot"). Examples: { "type": "string", "value": f"The highest salary is {highest_salary}." } or { "type": "number", "value": 125 } or { "type": "dataframe", "value": pd.DataFrame({...}) } or { "type": "plot", "value": "temp_chart.png" }""",
+            ),
+            (
+                "number",
+                """type (must be "number"), value must int. Example: { "type": "number", "value": 125 }""",
+            ),
+            (
+                "dataframe",
+                """type (must be "dataframe"), value must be pd.DataFrame or pd.Series. Example: { "type": "dataframe", "value": pd.DataFrame({...}) }""",
+            ),
+            (
+                "plot",
+                """type (must be "plot"), value must be string. Example: { "type": "plot", "value": "temp_chart.png" }""",
+            ),
+            (
+                "string",
+                """type (must be "string"), value must be string. Example: { "type": "string", "value": f"The highest salary is {highest_salary}." }""",
+            ),
         ],
     )
     def test_str_with_args(self, output_type, output_type_template):
@@ -62,8 +60,6 @@ def test_str_with_args(self, output_type, output_type_template):
         if sys.platform.startswith("win"):
             prompt_content = prompt_content.replace("\r\n", "\n")
 
-        print(prompt_content)
-
         assert (
             prompt_content
             == f'''<tables>