119 lines
3.7 KiB
Python
119 lines
3.7 KiB
Python
"""
|
||
Token 用量查询 API（v0.3 新增）
|
||
|
||
提供 LLM token 消耗的实时统计与成本估算。
|
||
内部用途，销售给客户的版本可移除此 blueprint 注册。
|
||
"""
|
||
|
||
from flask import jsonify, request
|
||
from . import usage_bp
|
||
from ..utils import token_tracker
|
||
from ..utils.logger import get_logger
|
||
|
||
logger = get_logger('foresight.api.usage')
|
||
|
||
|
||
@usage_bp.route('/summary', methods=['GET'])
|
||
def get_usage_summary():
|
||
"""
|
||
获取当前累计 token 用量摘要
|
||
|
||
返回：
|
||
{
|
||
"success": true,
|
||
"data": {
|
||
"stages": {
|
||
"<stage_name>": {
|
||
"by_model": {...},
|
||
"prompt_tokens": int,
|
||
"completion_tokens": int,
|
||
"calls": int,
|
||
"estimated_cost_cny": float
|
||
}
|
||
},
|
||
"total": {...},
|
||
"current_stage": "..."
|
||
}
|
||
}
|
||
"""
|
||
try:
|
||
return jsonify({
|
||
"success": True,
|
||
"data": token_tracker.get_summary(),
|
||
})
|
||
except Exception as e:
|
||
logger.error(f"获取 token 用量失败: {e}")
|
||
return jsonify({"success": False, "error": str(e)}), 500
|
||
|
||
|
||
@usage_bp.route('/reset', methods=['POST'])
|
||
def reset_usage():
|
||
"""
|
||
清空 token 用量统计
|
||
|
||
Body (可选):
|
||
{"stage": "<stage_name>"} // 不填则全部清空
|
||
"""
|
||
try:
|
||
data = request.get_json(silent=True) or {}
|
||
stage = data.get("stage")
|
||
token_tracker.reset(stage)
|
||
return jsonify({
|
||
"success": True,
|
||
"data": {"reset_stage": stage or "all"},
|
||
})
|
||
except Exception as e:
|
||
logger.error(f"重置 token 用量失败: {e}")
|
||
return jsonify({"success": False, "error": str(e)}), 500
|
||
|
||
|
||
@usage_bp.route('/estimate-simulation', methods=['GET'])
|
||
def estimate_simulation():
|
||
"""
|
||
估算 OASIS 模拟子进程的 token 消耗（无法精确追踪，只能估算）
|
||
|
||
Query params:
|
||
rounds: 模拟轮数 (默认 15)
|
||
active_agents_per_round: 每轮平均激活 agent 数 (默认 10)
|
||
llm_calls_per_action: 每动作 LLM 调用次数 (默认 4)
|
||
prompt_tokens: 平均 prompt token (默认 800)
|
||
completion_tokens: 平均 completion token (默认 200)
|
||
model: 模型名 (默认 glm-4-flash)
|
||
"""
|
||
try:
|
||
rounds = request.args.get('rounds', 15, type=int)
|
||
agents = request.args.get('active_agents_per_round', 10, type=int)
|
||
calls = request.args.get('llm_calls_per_action', 4, type=int)
|
||
prompt = request.args.get('prompt_tokens', 800, type=int)
|
||
completion = request.args.get('completion_tokens', 200, type=int)
|
||
model = request.args.get('model', 'glm-4-flash')
|
||
|
||
result = token_tracker.estimate_simulation_cost(
|
||
rounds=rounds,
|
||
active_agents_per_round_avg=agents,
|
||
llm_calls_per_action=calls,
|
||
avg_prompt_tokens_per_call=prompt,
|
||
avg_completion_tokens_per_call=completion,
|
||
model=model,
|
||
)
|
||
return jsonify({"success": True, "data": result})
|
||
except Exception as e:
|
||
logger.error(f"估算模拟 token 失败: {e}")
|
||
return jsonify({"success": False, "error": str(e)}), 500
|
||
|
||
|
||
@usage_bp.route('/set-stage', methods=['POST'])
|
||
def set_stage():
|
||
"""
|
||
手动设置当前 stage（用于测试或非自动埋点的场景）
|
||
|
||
Body: {"stage": "<stage_name>"}
|
||
"""
|
||
try:
|
||
data = request.get_json(silent=True) or {}
|
||
stage = data.get("stage")
|
||
token_tracker.set_stage(stage)
|
||
return jsonify({"success": True, "data": {"current_stage": stage}})
|
||
except Exception as e:
|
||
return jsonify({"success": False, "error": str(e)}), 500
|