Token 使用数量功能开发总结
一、需求概述
实现一个每日 token 使用量追踪功能:
- 从 0:00 到 24:00,统计每个模型的 token 使用量
- 在模型切换器中显示 token 使用量
- 当超过 200 万 token 时给出提示
二、技术方案演进
方案一:从 API 响应提取 Token Usage(失败)
思路:从模型 API 的响应中直接提取 token usage 信息
尝试的方法:
- 修改
PatchedChatDeepSeek类,在_generate_with_cache和_agenerate_with_cache方法中提取response_metadata和llm_output中的 usage 信息 - 在
_stream方法中追踪流式响应的 token 使用量
遇到的问题:
- DeepSeek API 在流式模式 下不返回 token usage(返回
None) - 强制使用非流式模式会导致响应速度变慢,用户体验下降
- 修改
_convert_chunk_to_generation_chunk方法导致流式显示异常(数据竖着显示)
方案二:后端中间件追踪 + 前端 LocalStorage 存储(部分成功)
思路:使用 LangGraph 中间件在模型调用前后追踪 token usage
实现的组件:
TokenUsageMiddleware:在after_model和after_agent中追踪 token usageTokenUsageState:在ThreadState中添加 token_usage 字段token_usage_store.py:全局 token usage 存储- API 端点
/api/model/token-usage:供前端查询
遇到的问题:
- 中间件无法直接获取模型名
- 全局存储增加了复杂度
- API 端点需要额外维护
方案三:字符计数估算 Token(最终方案)
思路:通过统计消息内容的字符数来估算 token 数量
核心公式:
token数量 ≈ 总字符数 × 0.6
原理:
- 中文字符:约 1 个 token / 字符
- 英文字符:约 1 个 token / 4 字符
- 混合文本:综合估算系数约为 0.6
实现的组件:
TokenUsageMiddleware:使用字符计数估算 token- 前端
useTokenUsage()hook:LocalStorage 持久化 + 每日重置 - 前端
model-switcher/index.html:显示 token 使用量
三、最终实现
3.1 后端实现
TokenUsageMiddleware
文件:backend/packages/harness/deerflow/agents/middlewares/token_usage_middleware.py
核心逻辑:
after_model:在模型调用后,统计输入消息和输出消息的字符数,估算 token 数量after_agent:将估算的 token usage 写入线程状态
关键代码:
python
def estimate_tokens_from_text(text: str) -> int:
"""Estimate token count from text using character-based heuristic."""
if not text:
return 0
return int(len(text) * 0.6)
ThreadState 修改
文件:backend/packages/harness/deerflow/agents/thread_state.py
修改内容 :添加 TokenUsageState 类型定义和 token_usage 字段
3.2 前端实现
useTokenUsage Hook
文件:frontend/src/core/tokens/hooks.ts
功能:
- LocalStorage 持久化每日 token 使用量
- 每日自动重置
- 多窗口同步
- 订阅更新机制
useThreadStream 修改
文件:frontend/src/core/threads/hooks.ts
修改内容 :从 state.values.token_usage 提取 token 使用量并更新前端状态
Model Switcher 修改
文件:frontend/public/model-switcher/index.html
修改内容:
- 在模型名称下方显示 token 使用量
- 添加颜色阈值(正常/警告/危险)
- 添加测试按钮和清除数据按钮
四、文件变更清单
新增文件
| 文件路径 | 说明 |
|---|---|
backend/packages/harness/deerflow/agents/middlewares/token_usage_middleware.py |
Token 使用量追踪中间件 |
frontend/src/core/tokens/hooks.ts |
Token 使用量状态管理 Hook |
修改文件
| 文件路径 | 修改内容 |
|---|---|
backend/packages/harness/deerflow/agents/thread_state.py |
添加 TokenUsageState 和 token_usage 字段 |
backend/packages/harness/deerflow/agents/lead_agent/agent.py |
添加 TokenUsageMiddleware 到中间件链 |
backend/packages/harness/deerflow/agents/memory/updater.py |
移除 token_usage_store 引用 |
frontend/src/core/threads/hooks.ts |
提取 token_usage 并更新状态 |
frontend/public/model-switcher/index.html |
显示 token 使用量 |
删除文件
| 文件路径 | 说明 |
|---|---|
backend/packages/harness/deerflow/agents/token_usage_store.py |
全局 token usage 存储(已移除) |
五、经验教训
5.1 不要过度工程化
问题:最初尝试从 API 响应中精确提取 token usage,导致:
- 修改了模型类的核心方法
- 破坏了流式响应的用户体验
- 增加了代码复杂度
教训:对于非关键性指标,使用估算方法可以大大简化实现,同时满足需求。
5.2 字符计数估算的准确性
优点:
- 实现简单,不需要依赖额外的 tokenizer 库
- 对中文和英文都有较好的适应性
- 不会影响响应速度
局限性:
- 估算值与实际值有一定偏差(约 ±20%)
- 对于特殊字符和格式文本估算不够准确
5.3 中间件模式的优势
优势:
- 解耦:token 追踪逻辑与业务逻辑分离
- 可复用:可以轻松扩展到其他模型
- 透明:不影响现有代码的执行流程
六、待改进项
- 引入专业 tokenizer :使用
tiktoken或模型专用 tokenizer 提高准确性 - 后端持久化:将 token usage 存储到数据库,支持历史查询
- 告警机制:当接近阈值时给出渐进式警告
- 多模型支持:当前只在 lead_agent 中启用,需要扩展到其他 agent