19 lines
802 B
Python
19 lines
802 B
Python
from langchain_core.callbacks import BaseCallbackHandler
|
|
from langchain_core.outputs import LLMResult
|
|
from typing import Any, Dict
|
|
|
|
class TokenUsageCallbackHandler(BaseCallbackHandler):
|
|
def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
|
|
token_usage = response.llm_output.get('token_usage', {})
|
|
|
|
if token_usage:
|
|
prompt_tokens = token_usage.get('prompt_tokens', 0)
|
|
completion_tokens = token_usage.get('completion_tokens', 0)
|
|
total_tokens = token_usage.get('total_tokens', 0)
|
|
|
|
print("--- [Token Usage] ---")
|
|
print(f" Prompt Tokens: {prompt_tokens}")
|
|
print(f" Completion Tokens: {completion_tokens}")
|
|
print(f" Total Tokens: {total_tokens}")
|
|
print("---------------------")
|