From ebb91e3c79823e0295ff3b45e5dcaf822dec01c7 Mon Sep 17 00:00:00 2001
From: mspronesti <massimiliano.pronesti@gmail.com>
Date: Wed, 4 Oct 2023 18:09:42 +0200
Subject: [PATCH 1/2] fix(helpers): separate cost of prompt and completion

---
 pandasai/helpers/openai_info.py | 52 +++++++++++++++++++++++++++++----
 1 file changed, 46 insertions(+), 6 deletions(-)

diff --git a/pandasai/helpers/openai_info.py b/pandasai/helpers/openai_info.py
index b1de2a784..66cb1cfed 100644
--- a/pandasai/helpers/openai_info.py
+++ b/pandasai/helpers/openai_info.py
@@ -10,30 +10,66 @@
     "gpt-4-0613": 0.03,
     "gpt-4-32k": 0.06,
     "gpt-4-32k-0613": 0.06,
+    # GPT-4 output
+    "gpt-4-completion": 0.06,
+    "gpt-4-0613-completion": 0.06,
+    "gpt-4-32k-completion": 0.12,
+    "gpt-4-32k-0613-completion": 0.12,
     # GPT-3.5 input
     "gpt-3.5-turbo": 0.0015,
-    "gpt-3.5-turbo-instruct": 0.0015,
     "gpt-3.5-turbo-0613": 0.0015,
+    "gpt-3.5-turbo-instruct": 0.0015,
     "gpt-3.5-turbo-16k": 0.003,
     "gpt-3.5-turbo-16k-0613": 0.003,
+    # GPT-3.5 output
+    "gpt-3.5-turbo-completion": 0.002,
+    "gpt-3.5-turbo-0613-completion": 0.002,
+    "gpt-3.5-turbo-instruct-completion": 0.002,
+    "gpt-3.5-turbo-16k-completion": 0.004,
+    "gpt-3.5-turbo-16k-0613-completion": 0.004,
+    # Azure GPT-35 input
+    "gpt-35-turbo": 0.0015,  # Azure OpenAI version of ChatGPT
+    "gpt-35-turbo-0613": 0.0015,
+    "gpt-35-turbo-instruct": 0.0015,
+    "gpt-35-turbo-16k": 0.003,
+    "gpt-35-turbo-16k-0613": 0.003,
+    # Azure GPT-35 output
+    "gpt-35-turbo-completion": 0.002,  # Azure OpenAI version of ChatGPT
+    "gpt-35-turbo-0613-completion": 0.002,
+    "gpt-35-turbo-instruct-completion": 0.002,
+    "gpt-35-turbo-16k-completion": 0.004,
+    "gpt-35-turbo-16k-0613-completion": 0.004,
     # Others
-    "gpt-35-turbo": 0.002,  # Azure OpenAI version of ChatGPT
+    "text-davinci-003": 0.02,
 }
 
 
 def get_openai_token_cost_for_model(
-    model_name: str,
-    num_tokens: int,
+        model_name: str,
+        num_tokens: int,
+        is_completion: bool = False,
 ) -> float:
     """
     Get the cost in USD for a given model and number of tokens.
+
     Args:
         model_name (str): Name of the model
         num_tokens (int): Number of tokens.
+        is_completion: Whether `num_tokens` refers to completion tokens or not.
+            Defaults to False.
+
     Returns:
         float: Cost in USD.
     """
     model_name = model_name.lower()
+    if is_completion and (
+            model_name.startswith("gpt-4")
+            or model_name.startswith("gpt-3.5")
+            or model_name.startswith("gpt-35")
+    ):
+        # The cost of completion token is different from
+        # the cost of prompt tokens.
+        model_name = model_name + "-completion"
     if model_name not in MODEL_COST_PER_1K_TOKENS:
         raise ValueError(
             f"Unknown model: {model_name}. Please provide a valid OpenAI model name."
@@ -63,10 +99,14 @@ def __call__(self, response: OpenAIObject) -> None:
         usage = response.usage
         if "total_tokens" not in usage:
             return None
+
         model_name = response.model
         if model_name in MODEL_COST_PER_1K_TOKENS:
-            total_cost = get_openai_token_cost_for_model(model_name, usage.total_tokens)
-            self.total_cost += total_cost
+            prompt_cost = get_openai_token_cost_for_model(model_name, usage.prompt_tokens)
+            completion_cost = get_openai_token_cost_for_model(
+                model_name, usage.completion_tokens, is_completion=True
+            )
+            self.total_cost += prompt_cost + completion_cost
 
         self.total_tokens += usage.total_tokens
         self.prompt_tokens += usage.prompt_tokens

From f1bf92dc80ca517cf0837a53409303b6a8473616 Mon Sep 17 00:00:00 2001
From: mspronesti <massimiliano.pronesti@gmail.com>
Date: Thu, 5 Oct 2023 00:06:24 +0200
Subject: [PATCH 2/2] chore: fix ruff lint

---
 pandasai/helpers/openai_info.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/pandasai/helpers/openai_info.py b/pandasai/helpers/openai_info.py
index 66cb1cfed..581f6a97f 100644
--- a/pandasai/helpers/openai_info.py
+++ b/pandasai/helpers/openai_info.py
@@ -102,7 +102,9 @@ def __call__(self, response: OpenAIObject) -> None:
 
         model_name = response.model
         if model_name in MODEL_COST_PER_1K_TOKENS:
-            prompt_cost = get_openai_token_cost_for_model(model_name, usage.prompt_tokens)
+            prompt_cost = get_openai_token_cost_for_model(
+                model_name, usage.prompt_tokens
+            )
             completion_cost = get_openai_token_cost_for_model(
                 model_name, usage.completion_tokens, is_completion=True
             )