96 lines
3.1 KiB
Python
96 lines
3.1 KiB
Python
"""
|
|
测试 Token 使用统计的索引功能
|
|
"""
|
|
import asyncio
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
# 添加 backend 到路径
|
|
sys.path.insert(0, str(Path(__file__).parent / "backend"))
|
|
|
|
from backend.services.token_usage_service import token_usage_service
|
|
from backend.models.internal import TokenUsageStatus
|
|
|
|
|
|
async def test_token_usage_indexes():
|
|
"""测试 Token 使用统计的索引功能"""
|
|
|
|
print("=" * 80)
|
|
print("Token 使用统计 - 索引功能测试")
|
|
print("=" * 80)
|
|
|
|
# 1. 记录一些测试数据
|
|
print("\n📝 记录测试数据...")
|
|
|
|
test_records = [
|
|
{
|
|
"chat_id": "测试角色/测试聊天1",
|
|
"role_name": "测试角色",
|
|
"chat_name": "测试聊天1",
|
|
"prompt_tokens": 100,
|
|
"completion_tokens": 200,
|
|
"total_tokens": 300,
|
|
"api_url": "https://api.openai.com/v1"
|
|
},
|
|
{
|
|
"chat_id": "测试角色/测试聊天2",
|
|
"role_name": "测试角色",
|
|
"chat_name": "测试聊天2",
|
|
"prompt_tokens": 150,
|
|
"completion_tokens": 250,
|
|
"total_tokens": 400,
|
|
"api_url": "https://api.deepseek.com/v1"
|
|
},
|
|
{
|
|
"chat_id": "测试角色/测试聊天1",
|
|
"role_name": "测试角色",
|
|
"chat_name": "测试聊天1",
|
|
"prompt_tokens": 120,
|
|
"completion_tokens": 180,
|
|
"total_tokens": 300,
|
|
"api_url": "https://api.openai.com/v1"
|
|
}
|
|
]
|
|
|
|
for record_data in test_records:
|
|
await token_usage_service.record_usage(**record_data)
|
|
print(f" ✅ 记录: {record_data['api_url']} - {record_data['total_tokens']} tokens")
|
|
|
|
# 2. 测试 API URL 索引查询
|
|
print("\n🔍 测试 API URL 索引查询...")
|
|
api_url_stats = await token_usage_service.get_api_url_stats()
|
|
|
|
if api_url_stats:
|
|
print(f"\n 找到 {len(api_url_stats)} 个 API URL:")
|
|
for url, stats in api_url_stats.items():
|
|
print(f"\n 📊 {url}:")
|
|
print(f" - 总调用次数: {stats['count']}")
|
|
print(f" - 总 Token 数: {stats['totalTokens']}")
|
|
print(f" - Prompt Tokens: {stats['totalPromptTokens']}")
|
|
print(f" - Completion Tokens: {stats['totalCompletionTokens']}")
|
|
else:
|
|
print(" ⚠️ 没有找到 API URL 统计数据")
|
|
|
|
# 3. 测试每日索引查询
|
|
print("\n🔍 测试每日索引查询...")
|
|
from datetime import datetime
|
|
now = datetime.now()
|
|
daily_stats = await token_usage_service.get_daily_stats(now.year, now.month)
|
|
|
|
if daily_stats:
|
|
print(f"\n 找到 {len(daily_stats)} 天的数据:")
|
|
for day, stats in list(daily_stats.items())[:5]: # 只显示前5天
|
|
print(f" 📅 {day}:")
|
|
print(f" - 调用次数: {stats['count']}")
|
|
print(f" - 总 Token 数: {stats['totalTokens']}")
|
|
else:
|
|
print(" ⚠️ 没有找到每日统计数据")
|
|
|
|
print("\n" + "=" * 80)
|
|
print("✅ 测试完成!")
|
|
print("=" * 80)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(test_token_usage_indexes())
|