LLM项目中的Token使用统计抽象化设计

在LLM项目中，开发者Simon Willison针对不同AI模型API返回的Token使用统计信息进行了抽象化设计，旨在简化Token使用统计的实现方式，并为后续的成本计算功能奠定基础。

Token统计的多样性挑战

当前各大AI模型提供商在API响应中返回的Token使用统计信息格式各异：

1. OpenAI：提供最详细的统计信息，包括输入Token、输出Token、总Token数，以及细分的缓存Token、音频Token等分类
2. Gemini 1.5 Pro：使用promptTokenCount、candidatesTokenCount和totalTokenCount三个字段
3. Anthropic：采用最简单的input_tokens和output_tokens结构
4. xAI/grok-beta和lambdalabs/hermes3-405b：与OpenAI类似但细节较少

这种多样性给统一处理Token统计带来了挑战。

设计方案演进

开发者最初考虑使用数据类来封装Token统计信息：

@dataclass
class Usage:
    model_id: str
    input_tokens: int
    output_tokens: int
    details: Dict[str, int]

但最终采用了更直接的方案，在Response类中添加了set_usage()方法：

def set_usage(input: int, output: int, details: dict = None) -> None:

这种方法更简洁，且与现有代码结构更契合。

详细统计信息的处理策略

对于OpenAI等提供的详细统计信息，项目采用了智能简化策略：

1. 移除所有值为0的键
2. 移除所有内容为空的嵌套对象
3. 保留核心的input_tokens和output_tokens作为独立字段
4. 将简化后的详细信息存储在单独的JSON字段中

例如，OpenAI的详细响应：

{
  "completion_tokens": 421,
  "prompt_tokens": 30791,
  "total_tokens": 31212,
  "prompt_tokens_details": {
    "cached_tokens": 30592,
    "audio_tokens": 0
  },
  "completion_tokens_details": {
    "reasoning_tokens": 0,
    "audio_tokens": 0,
    "accepted_prediction_tokens": 0,
    "rejected_prediction_tokens": 0
  }
}

会被简化为：

{
  "prompt_tokens_details": {"cached_tokens": 30592}
}

这种处理方式既保留了有价值的信息，又避免了存储冗余数据。

命令行工具集成

为了提升用户体验，项目还增加了命令行选项来直接查看Token使用情况：

llm prompt -u "你的提示内容"

这会在执行提示后立即显示Token使用统计，方便开发者实时了解资源消耗。

未来扩展考虑

设计时还考虑了未来可能的需求：

1. 批量模式支持：为后续实现批量处理时的成本计算预留接口
2. 成本计算基础：统一的Token统计为按模型定价计算实际花费奠定了基础
3. 扩展性：设计足够灵活以容纳未来可能出现的新统计维度

技术实现细节

简化统计信息的核心算法如下：

def simplify_usage_dict(d):
    def remove_empty_and_zero(obj):
        if isinstance(obj, dict):
            cleaned = {
                k: remove_empty_and_zero(v)
                for k, v in obj.items()
                if v != 0 and v != {}
            }
            return {k: v for k, v in cleaned.items() if v is not None and v != {}}
        return obj
    return remove_empty_and_zero(d) or {}

这个递归函数能够深入嵌套结构，智能地移除所有空值和零值，同时保留有意义的数据。

总结

LLM项目通过抽象化Token使用统计信息，解决了不同AI模型API返回格式不一致的问题。这一设计既考虑了当前需求，又为未来功能扩展预留了空间，体现了良好的软件工程实践。统一的接口简化了上层应用的开发，使开发者能够更专注于业务逻辑而非数据格式处理。