Spring AI项目中流式模式下获取Token数量的解决方案

背景介绍

在Spring AI项目中，开发者使用ChatClient与大型语言模型交互时，经常需要获取API调用的Token消耗情况。Token数量是评估API使用成本和优化提示词设计的重要指标。然而，许多开发者在流式传输(Stream)模式下遇到了无法获取Token数量的问题。

问题现象

当使用流式模式调用ChatClient时，返回的ChatResponse中的usage字段显示所有Token计数均为0：

"usage": {
  "promptTokens": 0,
  "completionTokens": 0,
  "totalTokens": 0
}

而在非流式模式下，相同的调用却能正确返回Token消耗数据：

"usage": {
  "promptTokens": 1649,
  "completionTokens": 261,
  "totalTokens": 1910
}

技术原理分析

流式传输与非流式传输在底层实现上有本质区别：

1. 非流式传输：API调用会等待完整响应返回后才处理结果，此时系统可以准确计算整个交互过程的Token消耗。

2. 流式传输：数据是分块实时返回的，系统在第一个数据块返回时尚未完成全部处理，因此无法立即提供完整的Token计数。

解决方案

针对这一问题，Spring AI提供了明确的配置选项。开发者需要在流式模式下显式启用Token计数功能：

// 在创建ChatClient时设置streamUsage为true
ChatClient.builder(chatModel)
    .defaultOptions(ChatOptions.builder()
        .withStreamUsage(true)  // 关键配置
        .build())
    // 其他配置...
    .build();

深入理解

1. 性能考量：流式模式下实时计算Token会增加额外开销，因此默认关闭此功能以优化性能。

2. 实现机制：启用streamUsage后，系统会：

   • 实时跟踪每个数据块的Token消耗
   • 在流结束时汇总完整计数
   • 通过响应元数据提供最终统计

3. 使用建议：

   • 仅在确实需要Token统计时启用此选项
   • 注意监控启用后的性能影响
   • 考虑在开发调试阶段启用，生产环境根据需求决定

最佳实践

对于需要同时兼顾流式体验和Token统计的场景，可以采用以下模式：

ChatResponse response = chatClient.prompt()
    // 各种配置
    .stream()
    .collect(Collectors.collectingAndThen(
        Collectors.toList(),
        chunks -> {
            // 处理完整响应
            return aggregateResponse(chunks);
        }
    ));

这种方式既保持了流式处理的优势，又能在处理完整响应时获取准确的Token统计。

总结

Spring AI项目为开发者提供了灵活的API调用方式选择。理解流式与非流式模式下的差异，合理配置相关选项，能够帮助开发者更好地控制和优化与大型语言模型的交互过程。特别是在需要精确统计资源消耗的场景下，正确设置streamUsage参数至关重要。