关键点解析
使用上下文管理器
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"])