Added helper to estimate cost of query with YOLOPandas

tests/integration_tests/test_llm_accessor.py

@@ -3,6 +3,7 @@
 
 from tests import TEST_DIRECTORY
 from yolopandas import pd
+from yolopandas.utils.query_helpers import run_query_with_cost
 
 
 class TestLLMAccessor(unittest.TestCase):
@@ -20,6 +21,11 @@ def test_basic_use(self):
         expected_result = 15
         self.assertEqual(expected_result, result)
 
+        result = run_query_with_cost(
+            self.product_df, "What is the price of the highest-priced book?", yolo=True
+        )
+        self.assertEqual(expected_result, result)
+
         result = self.product_df.llm.query(
             "What is the average price of products grouped by type?",
             yolo=True,

yolopandas/utils/query_helpers.py

@@ -0,0 +1,35 @@
+from typing import Any
+
+from langchain.callbacks import get_openai_callback
+
+from yolopandas import pd
+
+
+def run_query_with_cost(df: pd.DataFrame, query: str, yolo: bool = False) -> Any:
+    """
+    A function to run a YOLOPandas query with cost estimation returned for your query in terms of tokens used.
+    This includes total tokens, prompt tokens, completion tokens, and the total cost in USD.
+
+    Parameters
+    ----------
+    df : pd.DataFrame
+        The Pandas DataFrame with your data
+    query : str
+        The query you want to run against your data
+    yolo : bool
+        Boolean value used to return a prompt to a user or not to accept the code result before
+        running the code (False means to return the prompt)
+
+    Returns
+    -------
+    result : Any
+        The results of the query run against your data. A prompt may be returned as intermediary
+        output to proceed with generating the result or not.
+    """
+    with get_openai_callback() as cb:
+        result = df.llm.query(query, yolo=yolo)
+        print(f"Total Tokens: {cb.total_tokens}")
+        print(f"Prompt Tokens: {cb.prompt_tokens}")
+        print(f"Completion Tokens: {cb.completion_tokens}")
+        print(f"Total Cost (USD): ${cb.total_cost}")
+        return result