Merge branch 'main' into HenriqueAJNB/main

.gitignore

@@ -4,7 +4,10 @@
 # __pycache__
 __pycache__
 
+# macOS
+.DS_Store
+
 # build
 build
 dist
-pandasai.egg-info
+pandasai.egg-info

README.md

@@ -1,5 +1,6 @@
 # PandasAI 🐼
 
+[![release](https://img.shields.io/pypi/v/pandasai?label=Release&style=flat-square)](https://pypi.org/project/pandasai/)
 [![lint](https://github.com/gventuri/pandas-ai/actions/workflows/ci.yml/badge.svg)](https://github.com/gventuri/pandas-ai/actions/workflows/ci.yml/badge.svg)
 [![](https://dcbadge.vercel.app/api/server/kF7FqH2FwS?style=flat&compact=true)](https://discord.gg/kF7FqH2FwS)
 [![Downloads](https://static.pepy.tech/badge/pandasai/month)](https://pepy.tech/project/pandasai) [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
@@ -25,6 +26,8 @@ pip install pandasai
 
 ## Usage
 
+> Disclaimer: GDP data was collected from [this source](https://ourworldindata.org/grapher/gross-domestic-product?tab=table), published by World Development Indicators - World Bank (2022.05.26) and collected at National accounts data - World Bank / OECD. It relates to the year of 2020. Happiness indexes were extracted from [the World Happiness Report](https://ftnnews.com/images/stories/documents/2020/WHR20.pdf). Another useful [link](https://data.world/makeovermonday/2020w19-world-happiness-report-2020).
+
 PandasAI is designed to be used in conjunction with Pandas. It makes Pandas conversational, allowing you to ask questions about your data and get answers back, in the form of Pandas DataFrames. For example, you can ask PandasAI to find all the rows in a DataFrame where the value of a column is greater than 5, and it will return a DataFrame containing only those rows:
 
 ```python
@@ -34,8 +37,8 @@ from pandasai import PandasAI
 # Sample DataFrame
 df = pd.DataFrame({
     "country": ["United States", "United Kingdom", "France", "Germany", "Italy", "Spain", "Canada", "Australia", "Japan", "China"],
-    "gdp": [21400000, 2940000, 2830000, 3870000, 2160000, 1350000, 1780000, 1320000, 516000, 14000000],
-    "happiness_index": [7.3, 7.2, 6.5, 7.0, 6.0, 6.3, 7.3, 7.3, 5.9, 5.0]
+    "gdp": [19294482071552, 2891615567872, 2411255037952, 3435817336832, 1745433788416, 1181205135360, 1607402389504, 1490967855104, 4380756541440, 14631844184064],
+    "happiness_index": [6.94, 7.16, 6.66, 7.07, 6.38, 6.4, 7.23, 7.22, 5.87, 5.12]
 })
 
 # Instantiate a LLM
@@ -49,11 +52,11 @@ pandas_ai.run(df, prompt='Which are the 5 happiest countries?')
 The above code will return the following:
 
 ```
-0     United States
 6            Canada
 7         Australia
 1    United Kingdom
 3           Germany
+0     United States
 Name: country, dtype: object
 ```
 
@@ -66,7 +69,7 @@ pandas_ai.run(df, prompt='What is the sum of the GDPs of the 2 unhappiest countr
 The above code will return the following:
 
 ```
-14516000
+19012600725504
 ```
 
 You can also ask PandasAI to draw a graph:

examples/data/sample_dataframe.py

@@ -14,16 +14,16 @@
         "China",
     ],
     "gdp": [
-        21400000,
-        2940000,
-        2830000,
-        3870000,
-        2160000,
-        1350000,
-        1780000,
-        1320000,
-        516000,
-        14000000,
+        19294482071552,
+        2891615567872,
+        2411255037952,
+        3435817336832,
+        1745433788416,
+        1181205135360,
+        1607402389504,
+        1490967855104,
+        4380756541440,
+        14631844184064,
     ],
-    "happiness_index": [7.3, 7.2, 6.5, 7.0, 6.0, 6.3, 7.3, 7.3, 5.9, 5.0],
+    "happiness_index": [6.94, 7.16, 6.66, 7.07, 6.38, 6.4, 7.23, 7.22, 5.87, 5.12],
 }

examples/from_dataframe.py

@@ -12,4 +12,4 @@
 pandas_ai = PandasAI(llm, verbose=True, conversational=False)
 response = pandas_ai.run(df, "Calculate the sum of the gdp of north american countries")
 print(response)
-# Output: 26200000
+# Output: 20901884461056

examples/with_privacy_enforced.py

@@ -15,4 +15,4 @@
     "Calculate the sum of the gdp of north american countries",
 )
 print(response)
-# Output: 26200000
+# Output: 20901884461056

pandasai/__init__.py

@@ -11,7 +11,7 @@
 
 
 class PandasAI:
-    """PandasAI is a wrapper around a LLM to make dataframes convesational"""
+    """PandasAI is a wrapper around a LLM to make dataframes conversational"""
 
     _task_instruction: str = """
 There is a dataframe in pandas (python).
@@ -172,9 +172,7 @@ def run_code(
         lines = code.strip().split("\n")
         last_line = lines[-1].strip()
         if last_line.startswith("print(") and last_line.endswith(")"):
-            # Last line is already printing
-            return eval(last_line[6:-1])
-        # Evaluate last line and return its value or the captured output
+            last_line = last_line[6:-1]
         try:
             return eval(last_line)
         except Exception:  # pylint: disable=W0718

pandasai/llm/openai.py

@@ -62,7 +62,7 @@ def _set_params(self, **kwargs):
 
     def completion(self, prompt: str) -> str:
         """
-        Query the completation API
+        Query the completion API
 
         Args:
             prompt (str): Prompt

tests/test_pandasai.py

@@ -167,3 +167,10 @@ def test_run_without_privacy_enforcement(self):
 How many countries are in the dataframe?"""
         self.pandasai.run(df, "How many countries are in the dataframe?")
         assert self.pandasai._llm.last_prompt == expected_prompt
+
+    def test_run_with_print_at_the_end(self):
+        code = """
+result = {"happiness": 0.5, "gdp": 0.8}
+print(result)"""
+        self.setup(code)
+        self.pandasai.run_code(code, pd.DataFrame())