關鍵點解析
使用上下文管理器
with get_openai_callback() as cb:
這一行是核心。cb
會自動收集本次調用的 prompt tokens、completion tokens 以及 total tokens。自動統計
在上下文退出時,cb
中已經包含了這次調用的消耗情況,無需額外手動計算。累加到全局狀態
通過:state["token_count"].prompt_tokens += cb.prompt_tokens
就能把每次調用的消耗累加起來,便于后續統一展示。
附加信息管理
示例中我還加了metadata
,比如生成的名字和時間戳,這樣就能把統計和業務邏輯自然結合起來。
from langchain_openai import ChatOpenAI
from langchain.callbacks import get_openai_callback# ========== 初始化模型 ==========
llm = ChatOpenAI(model="gpt-3.5-turbo", # 你也可以換成 gpt-4o-mini 或其他模型temperature=0
)# 模擬全局狀態,記錄 Token 累計情況
state = {"token_count": {"prompt_tokens": 0,"completion_tokens": 0,"total_tokens": 0,}
}def call_llm(prompt: str):"""調用 LLM 并統計 token 消耗"""with get_openai_callback() as cb:response = llm.invoke(prompt)# 打印本次調用消耗print(f"本次調用消耗: prompt={cb.prompt_tokens}, "f"completion={cb.completion_tokens}, total={cb.total_tokens}")# 累加到全局 statestate["token_count"]["prompt_tokens"] += cb.prompt_tokensstate["token_count"]["completion_tokens"] += cb.completion_tokensstate["token_count"]["total_tokens"] += cb.total_tokensreturn response.content# ========== 示例運行 ==========
if __name__ == "__main__":result1 = call_llm("請寫一首五言絕句,主題是春天。")print("模型回復:", result1, "\n")result2 = call_llm("請用一句話總結一下這首詩。")print("模型回復:", result2, "\n")print("累計 Token 使用情況:", state["token_count"])