github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-07-14 17:07:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add prompt and completion token tracking (#2080)

Browse Source 
Tracking the breakdown of token usage is useful when using GPT-4, where
prompt and completion tokens are priced differently.
This commit is contained in:
iocuydi 2023-03-28 01:41:25 -05:00 committed by GitHub
parent b5020c7d9c
commit e8d9cbca3f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 6 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
langchain/callbacks/openai_info.py
View File

@ -9,6 +9,8 @@ class OpenAICallbackHandler(BaseCallbackHandler):
"""Callback Handler that tracks OpenAI info.""" """Callback Handler that tracks OpenAI info."""
total_tokens: int = 0 total_tokens: int = 0
prompt_tokens: int = 0
completion_tokens: int = 0
@property @property
def always_verbose(self) -> bool: def always_verbose(self) -> bool:
@ -32,6 +34,10 @@ class OpenAICallbackHandler(BaseCallbackHandler):
token_usage = response.llm_output["token_usage"] token_usage = response.llm_output["token_usage"]
if "total_tokens" in token_usage: if "total_tokens" in token_usage:
self.total_tokens += token_usage["total_tokens"] 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( def on_llm_error(
self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any