From e8d9cbca3fad5b2a21ffe4cf13e0941a86c02a48 Mon Sep 17 00:00:00 2001
From: iocuydi <46613640+iocuydi@users.noreply.github.com>
Date: Tue, 28 Mar 2023 01:41:25 -0500
Subject: [PATCH] Add prompt and completion token tracking (#2080)

Tracking the breakdown of token usage is useful when using GPT-4, where
prompt and completion tokens are priced differently.
---
 langchain/callbacks/openai_info.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/langchain/callbacks/openai_info.py b/langchain/callbacks/openai_info.py
index d8cebe585b..64ee90eea9 100644
--- a/langchain/callbacks/openai_info.py
+++ b/langchain/callbacks/openai_info.py
@@ -9,6 +9,8 @@ class OpenAICallbackHandler(BaseCallbackHandler):
     """Callback Handler that tracks OpenAI info."""
 
     total_tokens: int = 0
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
 
     @property
     def always_verbose(self) -> bool:
@@ -32,6 +34,10 @@ class OpenAICallbackHandler(BaseCallbackHandler):
                 token_usage = response.llm_output["token_usage"]
                 if "total_tokens" in token_usage:
                     self.total_tokens += token_usage["total_tokens"]
+                if "prompt_tokens" in token_usage:
+                    self.prompt_tokens += token_usage["prompt_tokens"]
+                if "completion_tokens" in token_usage:
+                    self.completion_tokens += token_usage["completion_tokens"]
 
     def on_llm_error(
         self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any