TokenUsage 回调函数概述

TokenUsage 是 LangChain 提供的实用回调函数，专门用于追踪和记录语言模型调用过程中的 token 消耗情况。可以帮助开发者：

  1. 精确监控 API 调用成本
  2. 优化提示词设计
  3. 分析模型使用效率

核心功能

TokenUsage 回调主要记录三类 token 数据：

  • 输入 token (prompt_tokens): 计算发送给模型的提示词 token 数量
  • 输出 token (completion_tokens): 计算模型返回结果的 token 数量
  • 总 token (total_tokens): 前两者的总和

实现方法

基础使用示例

from langchain.callbacks import TokenUsage
from langchain.llms import OpenAI

# 初始化回调
token_usage = TokenUsage()

# 创建LLM实例时传入回调
llm = OpenAI(
    temperature=0,
    callbacks=[token_usage]
)

# 执行查询
response = llm("请介绍一下Python语言")

# 获取token使用情况
print(f"输入token: {token_usage.prompt_tokens}")
print(f"输出token: {token_usage.completion_tokens}")
print(f"总token: {token_usage.total_tokens}")

链式调用中的使用

from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate

# 创建提示模板
prompt = PromptTemplate(
    input_variables=["topic"],
    template="请用中文详细解释以下主题: {topic}"
)

# 创建链
chain = LLMChain(
    llm=llm,
    prompt=prompt,
    callbacks=[token_usage]
)

# 执行链
result = chain.run(topic="机器学习")

高级应用场景

批量处理监控

topics = ["深度学习", "神经网络", "自然语言处理"]
total_tokens = 0

for topic in topics:
    chain.run(topic=topic)
    total_tokens += token_usage.total_tokens
    token_usage.reset()  # 重置计数器

print(f"批量处理总token消耗: {total_tokens}")

成本估算工具

def estimate_cost(total_tokens):
    # 假设使用GPT-3.5模型，每1000个token收费$0.002
    return total_tokens / 1000 * 0.002

cost = estimate_cost(total_tokens)
print(f"预计API调用成本: ${cost:.4f}")

注意事项

  1. 多回调组合：可以与其他回调(如StdOutCallback)同时使用
  2. 异步环境：在异步调用中需要确保回调函数的线程安全
  3. 重置机制：长时间运行的应用程序应定期重置计数器
  4. 模型差异：不同模型的token计算方式可能略有不同

最佳实践建议

  1. 在开发环境始终启用TokenUsage回调
  2. 为关键业务操作设置token消耗警报
  3. 定期分析历史token使用数据
  4. 根据token消耗优化提示词结构
  5. 考虑将token使用情况记录到监控系统

安装依赖

pip install -qU langchain-core langchain-openai

get_openai_callback 使用

get_openai_callback 是 LangChain 框架内置封装的实用工具，可以提供：

  1. 总 token 消耗量（包括输入和输出）
  2. 总 API 调用次数
  3. 预估的费用消耗
  4. 每次调用的详细记录

使用示例

from langchain.callbacks import get_openai_callback
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(
    model="gpt-3.5-turbo",
)

with get_openai_callback() as cb:
    result = llm.invoke("Tell me a joke")
    print(cb)

运行结果

Tokens Used: 26
    Prompt Tokens: 12
    Completion Tokens: 14
Successful Requests: 1
Total Cost (USD): $4.6e-05