Add total_cost estimates based on token count for openai (#2243)

We have completion and prompt tokens, model names, so if we can, let's keep a running total of the cost.
1 year ago · 7a8f1d2854
parent 632c2b49da
commit 7a8f1d2854
2 changed files with 84 additions and 26 deletions
--- a/docs/modules/models/llms/examples/token_usage_tracking.ipynb
+++ b/docs/modules/models/llms/examples/token_usage_tracking.ipynb
@ -14,7 +14,7 @@
  },
  {
   "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 2,
   "id": "9455db35",
   "metadata": {},
   "outputs": [],
@ -35,7 +35,7 @@
  },
  {
   "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 4,
   "id": "31667d54",
   "metadata": {},
   "outputs": [
@ -43,10 +43,11 @@
     "name": "stdout",
     "output_type": "stream",
     "text": [
-      "Total Tokens: 42\n",
+      "Total Tokens: 39\n",
      "Prompt Tokens: 4\n",
-      "Completion Tokens: 38\n",
-      "Successful Requests: 1\n"
+      "Completion Tokens: 35\n",
+      "Successful Requests: 1\n",
+      "Total Cost (USD): $0.0007800000000000001\n"
     ]
    }
   ],
@ -56,7 +57,8 @@
    "    print(f\"Total Tokens: {cb.total_tokens}\")\n",
    "    print(f\"Prompt Tokens: {cb.prompt_tokens}\")\n",
    "    print(f\"Completion Tokens: {cb.completion_tokens}\")\n",
-    "    print(f\"Successful Requests: {cb.successful_requests}\")"
+    "    print(f\"Successful Requests: {cb.successful_requests}\")\n",
+    "    print(f\"Total Cost (USD): ${cb.total_cost}\")"
   ]
  },
  {
@ -69,7 +71,7 @@
  },
  {
   "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 5,
   "id": "e09420f4",
   "metadata": {},
   "outputs": [
@ -77,7 +79,7 @@
     "name": "stdout",
     "output_type": "stream",
     "text": [
-      "83\n"
+      "91\n"
     ]
    }
   ],
@ -98,7 +100,7 @@
  },
  {
   "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 6,
   "id": "5d1125c6",
   "metadata": {},
   "outputs": [],
@ -124,37 +126,43 @@
     "text": [
      "\n",
      "\n",
-      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
-      "\u001b[32;1m\u001b[1;3m I need to find out who Olivia Wilde's boyfriend is and then calculate his age raised to the 0.23 power.\n",
+      "\u001B[1m> Entering new AgentExecutor chain...\u001B[0m\n",
+      "\u001B[32;1m\u001B[1;3m I need to find out who Olivia Wilde's boyfriend is and then calculate his age raised to the 0.23 power.\n",
      "Action: Search\n",
-      "Action Input: \"Olivia Wilde boyfriend\"\u001b[0m\n",
-      "Observation: \u001b[36;1m\u001b[1;3mJason Sudeikis\u001b[0m\n",
-      "Thought:\u001b[32;1m\u001b[1;3m I need to find out Jason Sudeikis' age\n",
+      "Action Input: \"Olivia Wilde boyfriend\"\u001B[0m\n",
+      "Observation: \u001B[36;1m\u001B[1;3mSudeikis and Wilde's relationship ended in November 2020. Wilde was publicly served with court documents regarding child custody while she was presenting Don't Worry Darling at CinemaCon 2022. In January 2021, Wilde began dating singer Harry Styles after meeting during the filming of Don't Worry Darling.\u001B[0m\n",
+      "Thought:\u001B[32;1m\u001B[1;3m I need to find out Harry Styles' age.\n",
      "Action: Search\n",
-      "Action Input: \"Jason Sudeikis age\"\u001b[0m\n",
-      "Observation: \u001b[36;1m\u001b[1;3m47 years\u001b[0m\n",
-      "Thought:\u001b[32;1m\u001b[1;3m I need to calculate 47 raised to the 0.23 power\n",
+      "Action Input: \"Harry Styles age\"\u001B[0m\n",
+      "Observation: \u001B[36;1m\u001B[1;3m29 years\u001B[0m\n",
+      "Thought:\u001B[32;1m\u001B[1;3m I need to calculate 29 raised to the 0.23 power.\n",
      "Action: Calculator\n",
-      "Action Input: 47^0.23\u001b[0m\n",
-      "Observation: \u001b[33;1m\u001b[1;3mAnswer: 2.4242784855673896\n",
-      "\u001b[0m\n",
-      "Thought:\u001b[32;1m\u001b[1;3m I now know the final answer\n",
-      "Final Answer: Jason Sudeikis, Olivia Wilde's boyfriend, is 47 years old and his age raised to the 0.23 power is 2.4242784855673896.\u001b[0m\n",
+      "Action Input: 29^0.23\u001B[0m\n",
+      "Observation: \u001B[33;1m\u001B[1;3mAnswer: 2.169459462491557\n",
+      "\u001B[0m\n",
+      "Thought:\u001B[32;1m\u001B[1;3m I now know the final answer.\n",
+      "Final Answer: Harry Styles, Olivia Wilde's boyfriend, is 29 years old and his age raised to the 0.23 power is 2.169459462491557.\u001B[0m\n",
      "\n",
-      "\u001b[1m> Finished chain.\u001b[0m\n",
-      "1465\n"
+      "\u001B[1m> Finished chain.\u001B[0m\n",
+      "Total Tokens: 1506\n",
+      "Prompt Tokens: 1350\n",
+      "Completion Tokens: 156\n",
+      "Total Cost (USD): $0.03012\n"
     ]
    }
   ],
   "source": [
    "with get_openai_callback() as cb:\n",
    "    response = agent.run(\"Who is Olivia Wilde's boyfriend? What is his current age raised to the 0.23 power?\")\n",
-    "    print(cb.total_tokens)"
+    "    print(f\"Total Tokens: {cb.total_tokens}\")\n",
+    "    print(f\"Prompt Tokens: {cb.prompt_tokens}\")\n",
+    "    print(f\"Completion Tokens: {cb.completion_tokens}\")\n",
+    "    print(f\"Total Cost (USD): ${cb.total_cost}\")"
   ]
  },
  {
   "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 7,
   "id": "80ca77a3",
   "metadata": {},
   "outputs": [],
--- a/langchain/callbacks/openai_info.py
+++ b/langchain/callbacks/openai_info.py
@ -5,6 +5,45 @@ from langchain.callbacks.base import BaseCallbackHandler
 from langchain.schema import AgentAction, AgentFinish, LLMResult


+def get_openai_model_cost_per_1k_tokens(
+    model_name: str, is_completion: bool = False
+) -> float:
+    model_cost_mapping = {
+        "gpt-4": 0.03,
+        "gpt-4-0314": 0.03,
+        "gpt-4-completion": 0.06,
+        "gpt-4-0314-completion": 0.06,
+        "gpt-4-32k": 0.06,
+        "gpt-4-32k-0314": 0.06,
+        "gpt-4-32k-completion": 0.12,
+        "gpt-4-32k-0314-completion": 0.12,
+        "gpt-3.5-turbo": 0.002,
+        "gpt-3.5-turbo-0301": 0.002,
+        "text-ada-001": 0.0004,
+        "ada": 0.0004,
+        "text-babbage-001": 0.0005,
+        "babbage": 0.0005,
+        "text-curie-001": 0.002,
+        "curie": 0.002,
+        "text-davinci-003": 0.02,
+        "text-davinci-002": 0.02,
+        "code-davinci-002": 0.02,
+    }
+
+    cost = model_cost_mapping.get(
+        model_name.lower()
+        + ("-completion" if is_completion and model_name.startswith("gpt-4") else ""),
+        None,
+    )
+    if cost is None:
+        raise ValueError(
+            f"Unknown model: {model_name}. Please provide a valid OpenAI model name."
+            "Known models are: " + ", ".join(model_cost_mapping.keys())
+        )
+
+    return cost
+
+
 class OpenAICallbackHandler(BaseCallbackHandler):
    """Callback Handler that tracks OpenAI info."""

@ -12,6 +51,7 @@ class OpenAICallbackHandler(BaseCallbackHandler):
    prompt_tokens: int = 0
    completion_tokens: int = 0
    successful_requests: int = 0
+    total_cost: float = 0.0

    @property
    def always_verbose(self) -> bool:
@ -34,6 +74,16 @@ class OpenAICallbackHandler(BaseCallbackHandler):
            self.successful_requests += 1
            if "token_usage" in response.llm_output:
                token_usage = response.llm_output["token_usage"]
+                if "model_name" in response.llm_output:
+                    completion_cost = get_openai_model_cost_per_1k_tokens(
+                        response.llm_output["model_name"], is_completion=True
+                    ) * (token_usage.get("completion_tokens", 0) / 1000)
+                    prompt_cost = get_openai_model_cost_per_1k_tokens(
+                        response.llm_output["model_name"]
+                    ) * (token_usage.get("prompt_tokens", 0) / 1000)
+
+                    self.total_cost += prompt_cost + completion_cost
+
                if "total_tokens" in token_usage:
                    self.total_tokens += token_usage["total_tokens"]
                if "prompt_tokens" in token_usage: