|
|
import datetime
|
|
|
|
|
|
from config import API_COST_LIMITS, MONOGODB_DB_NAME
|
|
|
|
|
|
from .chat_messages import ChatMessages
|
|
|
from .mongodb_init import mongo_client as client
|
|
|
|
|
|
def record_api_cost(messages: ChatMessages):
|
|
|
"""记录API调用费用"""
|
|
|
|
|
|
db = client[MONOGODB_DB_NAME]
|
|
|
collection = db['api_cost']
|
|
|
|
|
|
cost_data = {
|
|
|
'created_at': datetime.datetime.now(),
|
|
|
'model': messages.model,
|
|
|
'cost': messages.cost,
|
|
|
'currency_symbol': messages.currency_symbol,
|
|
|
'input_tokens': messages[:-1].count_message_tokens(),
|
|
|
'output_tokens': messages[-1:].count_message_tokens(),
|
|
|
'total_tokens': messages.count_message_tokens()
|
|
|
}
|
|
|
collection.insert_one(cost_data)
|
|
|
|
|
|
def get_model_cost_stats(start_date: datetime.datetime, end_date: datetime.datetime) -> list:
|
|
|
"""获取指定时间段内的模型调用费用统计"""
|
|
|
pipeline = [
|
|
|
{
|
|
|
'$match': {
|
|
|
'created_at': {
|
|
|
'$gte': start_date,
|
|
|
'$lte': end_date
|
|
|
}
|
|
|
}
|
|
|
},
|
|
|
{
|
|
|
'$group': {
|
|
|
'_id': '$model',
|
|
|
'total_cost': { '$sum': '$cost' },
|
|
|
'total_calls': { '$sum': 1 },
|
|
|
'total_input_tokens': { '$sum': '$input_tokens' },
|
|
|
'total_output_tokens': { '$sum': '$output_tokens' },
|
|
|
'total_tokens': { '$sum': '$total_tokens' },
|
|
|
'avg_cost_per_call': { '$avg': '$cost' },
|
|
|
'currency_symbol': { '$first': '$currency_symbol' }
|
|
|
}
|
|
|
},
|
|
|
{
|
|
|
'$project': {
|
|
|
'model': '$_id',
|
|
|
'total_cost': { '$round': ['$total_cost', 4] },
|
|
|
'total_calls': 1,
|
|
|
'total_input_tokens': 1,
|
|
|
'total_output_tokens': 1,
|
|
|
'total_tokens': 1,
|
|
|
'avg_cost_per_call': { '$round': ['$avg_cost_per_call', 4] },
|
|
|
'currency_symbol': 1,
|
|
|
'_id': 0
|
|
|
}
|
|
|
},
|
|
|
{
|
|
|
'$sort': { 'total_cost': -1 }
|
|
|
}
|
|
|
]
|
|
|
|
|
|
|
|
|
db = client[MONOGODB_DB_NAME]
|
|
|
collection = db['api_cost']
|
|
|
|
|
|
stats = list(collection.aggregate(pipeline))
|
|
|
return stats
|
|
|
|
|
|
|
|
|
def print_cost_report(days: int = 30, hours: int = 0):
|
|
|
"""打印最近N天的费用报告"""
|
|
|
end_date = datetime.datetime.now()
|
|
|
start_date = end_date - datetime.timedelta(days=days, hours=hours)
|
|
|
|
|
|
stats = get_model_cost_stats(start_date, end_date)
|
|
|
|
|
|
print(f"\n=== API Cost Report ({start_date.date()} to {end_date.date()}) ===")
|
|
|
for model_stat in stats:
|
|
|
print(f"\nModel: {model_stat['model']}")
|
|
|
print(f"Total Cost: {model_stat['currency_symbol']}{model_stat['total_cost']:.4f}")
|
|
|
print(f"Total Calls: {model_stat['total_calls']}")
|
|
|
print(f"Total Tokens: {model_stat['total_tokens']:,}")
|
|
|
print(f"Avg Cost/Call: {model_stat['currency_symbol']}{model_stat['avg_cost_per_call']:.4f}")
|
|
|
|
|
|
def check_cost_limits() -> bool:
|
|
|
"""
|
|
|
检查API调用费用是否超过限制
|
|
|
返回: 如果未超过限制返回True，否则返回False
|
|
|
"""
|
|
|
now = datetime.datetime.now()
|
|
|
hour_ago = now - datetime.timedelta(hours=1)
|
|
|
day_ago = now - datetime.timedelta(days=1)
|
|
|
|
|
|
|
|
|
hour_stats = get_model_cost_stats(hour_ago, now)
|
|
|
day_stats = get_model_cost_stats(day_ago, now)
|
|
|
|
|
|
|
|
|
hour_total_rmb = sum(
|
|
|
stat['total_cost'] * (API_COST_LIMITS['USD_TO_RMB_RATE'] if stat['currency_symbol'] == '$' else 1)
|
|
|
for stat in hour_stats
|
|
|
)
|
|
|
day_total_rmb = sum(
|
|
|
stat['total_cost'] * (API_COST_LIMITS['USD_TO_RMB_RATE'] if stat['currency_symbol'] == '$' else 1)
|
|
|
for stat in day_stats
|
|
|
)
|
|
|
|
|
|
|
|
|
if hour_total_rmb >= API_COST_LIMITS['HOURLY_LIMIT_RMB']:
|
|
|
print(f"警告：最近1小时API费用（￥{hour_total_rmb:.2f}）超过限制（￥{API_COST_LIMITS['HOURLY_LIMIT_RMB']}）")
|
|
|
raise Exception("最近1小时内API调用费用超过设定上限！")
|
|
|
|
|
|
if day_total_rmb >= API_COST_LIMITS['DAILY_LIMIT_RMB']:
|
|
|
print(f"警告：最近24小时API费用（￥{day_total_rmb:.2f}）超过限制（￥{API_COST_LIMITS['DAILY_LIMIT_RMB']}）")
|
|
|
raise Exception("最近1天内API调用费用超过设定上限！")
|
|
|
|
|
|
return True |
