400 lines
15 KiB
Python
400 lines
15 KiB
Python
"""
|
||
MaxKB RAG 性能测试 - 测试首个token接收时间
|
||
"""
|
||
import asyncio
|
||
import aiohttp
|
||
import time
|
||
import sys
|
||
import os
|
||
import json
|
||
|
||
# 添加项目路径
|
||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
|
||
# MaxKB 配置
|
||
MAXKB_API_URL = "http://10.0.0.202:8080"
|
||
API_KEY = "application-b758e71ecb8d4237f91795365b433c35"
|
||
CSRF_TOKEN = "bDzz45UzuDGB33YCykegwEG9QMBuLba3cisIDsL5Z4tsXTrntxZq8nhVJf7qT5Pq"
|
||
|
||
# DeepSeek 配置
|
||
DEEPSEEK_API_BASE = "https://api.deepseek.com"
|
||
DEEPSEEK_API_KEY = "sk-56b608b26a6949e4b09b5bf5f11c8f5b"
|
||
DEEPSEEK_MODEL = "deepseek-chat"
|
||
|
||
# 测试问题
|
||
TEST_QUESTION = "你好，请介绍一下你自己"
|
||
|
||
|
||
async def test_deepseek_first_token(question: str = TEST_QUESTION):
|
||
"""测试 DeepSeek 直接调用的首个token接收时间"""
|
||
|
||
print("=" * 70)
|
||
print("DeepSeek API 性能测试 - 首个Token接收时间")
|
||
print("=" * 70)
|
||
|
||
url = f"{DEEPSEEK_API_BASE}/chat/completions"
|
||
headers = {
|
||
"Content-Type": "application/json",
|
||
"Authorization": f"Bearer {DEEPSEEK_API_KEY}"
|
||
}
|
||
data = {
|
||
"model": DEEPSEEK_MODEL,
|
||
"messages": [
|
||
{"role": "user", "content": question}
|
||
],
|
||
"stream": True
|
||
}
|
||
|
||
print(f"\n[1/2] 发送测试问题: \"{question}\"")
|
||
print(" 等待响应...")
|
||
|
||
total_start_time = time.time()
|
||
first_token_time = None
|
||
token_count = 0
|
||
first_token_received = False
|
||
|
||
async with aiohttp.ClientSession() as session:
|
||
async with session.post(url, json=data, headers=headers) as response:
|
||
if response.status != 200:
|
||
print(f" ✗ 请求失败: {response.status}")
|
||
error_text = await response.text()
|
||
print(f" 错误信息: {error_text}")
|
||
return
|
||
|
||
print(f" ✓ 连接建立，开始接收数据...")
|
||
print(f"\n[2/2] 接收流式数据...")
|
||
print("-" * 70)
|
||
|
||
async for chunk in response.content:
|
||
current_time = time.time()
|
||
elapsed = current_time - total_start_time
|
||
|
||
# 解析 chunk 数据
|
||
chunk_str = chunk.decode('utf-8')
|
||
|
||
if not chunk_str.strip():
|
||
continue
|
||
|
||
# 记录首个token时间
|
||
if not first_token_received and 'data:' in chunk_str:
|
||
first_token_time = elapsed
|
||
first_token_received = True
|
||
print(f"\n🎯 首个Token接收时间: {first_token_time:.4f} 秒")
|
||
print("-" * 70)
|
||
|
||
# 解析并显示内容
|
||
for line in chunk_str.split('\n'):
|
||
if line.startswith('data:'):
|
||
data_str = line[5:].strip()
|
||
if data_str == '[DONE]':
|
||
continue
|
||
|
||
try:
|
||
data_json = json.loads(data_str)
|
||
# 从 choices[0].delta.content 获取实际内容（DeepSeek响应格式）
|
||
choices = data_json.get('choices', [])
|
||
if choices:
|
||
delta = choices[0].get('delta', {})
|
||
content = delta.get('content', '')
|
||
|
||
if content:
|
||
token_count += 1
|
||
# 显示内容（前40个字符）
|
||
display_content = content[:40] + '...' if len(content) > 40 else content
|
||
print(f"[{elapsed:.3f}s] Token #{token_count}: {display_content}")
|
||
|
||
except json.JSONDecodeError:
|
||
pass
|
||
|
||
total_time = time.time() - total_start_time
|
||
|
||
# 输出测试结果
|
||
print("\n" + "=" * 70)
|
||
print("📊 DeepSeek API 性能测试结果")
|
||
print("=" * 70)
|
||
print(f" 测试问题: {question}")
|
||
print(f" 首个Token时间: {first_token_time:.4f} 秒" if first_token_time else " 首个Token时间: N/A")
|
||
print(f" 总响应时间: {total_time:.4f} 秒")
|
||
print(f" 接收Token数: {token_count}")
|
||
if first_token_time and total_time > 0:
|
||
print(f" 平均Token速率: {token_count / (total_time - first_token_time):.2f} tokens/秒")
|
||
print("=" * 70)
|
||
|
||
return first_token_time, total_time
|
||
|
||
|
||
async def get_chat_id():
|
||
"""获取一个新的 chat_id"""
|
||
url = f"{MAXKB_API_URL}/chat/api/open"
|
||
headers = {
|
||
"accept": "*/*",
|
||
"Authorization": f"Bearer {API_KEY}"
|
||
}
|
||
|
||
async with aiohttp.ClientSession() as session:
|
||
async with session.get(url, headers=headers) as response:
|
||
if response.status == 200:
|
||
data = await response.json()
|
||
# data 直接就是 chat_id 字符串
|
||
return data.get("data")
|
||
else:
|
||
raise Exception(f"获取 chat_id 失败: {response.status}")
|
||
|
||
|
||
async def test_first_token_time(question: str = TEST_QUESTION):
|
||
"""测试首个token的接收时间"""
|
||
|
||
print("=" * 70)
|
||
print("MaxKB RAG 性能测试 - 首个Token接收时间")
|
||
print("=" * 70)
|
||
|
||
# 1. 获取 chat_id
|
||
print("\n[1/3] 获取 chat_id...")
|
||
chat_id_start_time = time.time()
|
||
try:
|
||
chat_id = await get_chat_id()
|
||
chat_id_time = time.time() - chat_id_start_time
|
||
print(f" ✓ chat_id: {chat_id} ({chat_id_time:.4f}s)")
|
||
except Exception as e:
|
||
print(f" ✗ 获取 chat_id 失败: {e}")
|
||
return
|
||
|
||
# 2. 发起流式对话请求
|
||
print(f"\n[2/3] 发送测试问题: \"{question}\"")
|
||
print(" 等待响应...")
|
||
|
||
url = f"{MAXKB_API_URL}/chat/api/chat_message/{chat_id}"
|
||
headers = {
|
||
"accept": "*/*",
|
||
"Authorization": f"Bearer {API_KEY}",
|
||
"Content-Type": "application/json",
|
||
"X-CSRFTOKEN": CSRF_TOKEN
|
||
}
|
||
data = {
|
||
"message": question,
|
||
"stream": True,
|
||
"re_chat": False
|
||
}
|
||
|
||
total_start_time = time.time()
|
||
first_token_time = None
|
||
token_count = 0
|
||
first_token_received = False
|
||
|
||
async with aiohttp.ClientSession() as session:
|
||
async with session.post(url, json=data, headers=headers) as response:
|
||
if response.status != 200:
|
||
print(f" ✗ 请求失败: {response.status}")
|
||
error_text = await response.text()
|
||
print(f" 错误信息: {error_text}")
|
||
return
|
||
|
||
print(f" ✓ 连接建立，开始接收数据...")
|
||
print(f"\n[3/3] 接收流式数据...")
|
||
print("-" * 70)
|
||
|
||
async for chunk in response.content:
|
||
current_time = time.time()
|
||
elapsed = current_time - total_start_time
|
||
|
||
# 解析 chunk 数据
|
||
chunk_str = chunk.decode('utf-8')
|
||
|
||
if not chunk_str.strip():
|
||
continue
|
||
|
||
# 记录首个token时间
|
||
if not first_token_received and chunk_str.startswith('data:'):
|
||
first_token_time = elapsed
|
||
first_token_received = True
|
||
print(f"\n🎯 首个Token接收时间: {first_token_time:.4f} 秒")
|
||
print("-" * 70)
|
||
|
||
# 解析并显示内容
|
||
for line in chunk_str.split('\n'):
|
||
if line.startswith('data:'):
|
||
data_str = line[5:].strip()
|
||
if data_str == '[DONE]':
|
||
continue
|
||
|
||
try:
|
||
import json
|
||
data_json = json.loads(data_str)
|
||
# 从 content 字段获取实际内容（MaxKB响应格式）
|
||
content = data_json.get('content', '')
|
||
|
||
if content:
|
||
if not first_token_received:
|
||
first_token_time = elapsed
|
||
first_token_received = True
|
||
print(f"\n🎯 首个Token接收时间: {first_token_time:.4f} 秒")
|
||
print("-" * 70)
|
||
|
||
token_count += 1
|
||
# 显示内容（前40个字符）
|
||
display_content = content[:40] + '...' if len(content) > 40 else content
|
||
print(f"[{elapsed:.3f}s] Token #{token_count}: {display_content}")
|
||
|
||
except json.JSONDecodeError:
|
||
pass
|
||
|
||
total_time = time.time() - total_start_time
|
||
|
||
# 3. 输出测试结果
|
||
print("\n" + "=" * 70)
|
||
print("📊 性能测试结果")
|
||
print("=" * 70)
|
||
print(f" 测试问题: {question}")
|
||
print(f" 首个Token时间: {first_token_time:.4f} 秒" if first_token_time else " 首个Token时间: N/A")
|
||
print(f" 总响应时间: {total_time:.4f} 秒")
|
||
print(f" 接收Token数: {token_count}")
|
||
if first_token_time and total_time > 0:
|
||
print(f" 平均Token速率: {token_count / (total_time - first_token_time):.2f} tokens/秒")
|
||
print("=" * 70)
|
||
|
||
return first_token_time, total_time
|
||
|
||
|
||
async def run_multiple_tests(n: int = 3, question: str = None):
|
||
"""运行多次测试取平均值"""
|
||
if question is None:
|
||
question = TEST_QUESTION
|
||
|
||
print(f"\n🔄 准备运行 {n} 次性能测试...")
|
||
print(f" 测试问题: {question}")
|
||
print()
|
||
|
||
first_token_times = []
|
||
total_times = []
|
||
|
||
for i in range(n):
|
||
print(f"\n{'='*70}")
|
||
print(f"测试 #{i + 1}/{n}")
|
||
print('='*70)
|
||
|
||
# 重新获取 chat_id（每次测试使用新的会话）
|
||
chat_id_start_time = time.time()
|
||
chat_id = await get_chat_id()
|
||
chat_id_time = time.time() - chat_id_start_time
|
||
print(f" 获取chat_id: {chat_id_time:.4f}s")
|
||
|
||
url = f"{MAXKB_API_URL}/chat/api/chat_message/{chat_id}"
|
||
headers = {
|
||
"accept": "*/*",
|
||
"Authorization": f"Bearer {API_KEY}",
|
||
"Content-Type": "application/json",
|
||
"X-CSRFTOKEN": CSRF_TOKEN
|
||
}
|
||
data = {
|
||
"message": question,
|
||
"stream": True,
|
||
"re_chat": False
|
||
}
|
||
|
||
total_start_time = time.time()
|
||
first_token_time = None
|
||
token_count = 0
|
||
first_token_received = False
|
||
|
||
async with aiohttp.ClientSession() as session:
|
||
async with session.post(url, json=data, headers=headers) as response:
|
||
if response.status != 200:
|
||
print(f"请求失败: {response.status}")
|
||
continue
|
||
|
||
async for chunk in response.content:
|
||
current_time = time.time()
|
||
elapsed = current_time - total_start_time
|
||
|
||
chunk_str = chunk.decode('utf-8')
|
||
if not chunk_str.strip():
|
||
continue
|
||
|
||
if not first_token_received and chunk_str.startswith('data:'):
|
||
first_token_time = elapsed
|
||
first_token_received = True
|
||
|
||
if chunk_str.startswith('data:'):
|
||
data_str = chunk_str[5:].strip()
|
||
if data_str == '[DONE]':
|
||
continue
|
||
try:
|
||
import json
|
||
data_json = json.loads(data_str)
|
||
# 从 content 字段获取实际内容
|
||
if data_json.get('content'):
|
||
if not first_token_received:
|
||
first_token_time = elapsed
|
||
first_token_received = True
|
||
token_count += 1
|
||
except:
|
||
pass
|
||
|
||
total_time = time.time() - total_start_time
|
||
|
||
if first_token_time:
|
||
first_token_times.append(first_token_time)
|
||
total_times.append(total_time)
|
||
|
||
print(f"\n 结果 #{i + 1}: 首Token={first_token_time:.4f}s, 总时间={total_time:.4f}s, Tokens={token_count}")
|
||
await asyncio.sleep(1) # 等待1秒再进行下次测试
|
||
|
||
# 计算平均值
|
||
print(f"\n{'='*70}")
|
||
print("📈 多次测试统计结果")
|
||
print('='*70)
|
||
if first_token_times:
|
||
avg_first_token = sum(first_token_times) / len(first_token_times)
|
||
min_first_token = min(first_token_times)
|
||
max_first_token = max(first_token_times)
|
||
print(f" 首个Token时间:")
|
||
print(f" - 平均值: {avg_first_token:.4f} 秒")
|
||
print(f" - 最小值: {min_first_token:.4f} 秒")
|
||
print(f" - 最大值: {max_first_token:.4f} 秒")
|
||
|
||
if total_times:
|
||
avg_total = sum(total_times) / len(total_times)
|
||
print(f" 总响应时间:")
|
||
print(f" - 平均值: {avg_total:.4f} 秒")
|
||
print('='*70)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
import argparse
|
||
|
||
parser = argparse.ArgumentParser(description='MaxKB RAG 性能测试')
|
||
parser.add_argument('question', nargs='?', default=TEST_QUESTION, help='测试问题（直接写在命令后面）')
|
||
parser.add_argument('--times', type=int, default=1, help='测试次数 (默认1次)')
|
||
parser.add_argument('--deepseek', action='store_true', help='同时测试 DeepSeek API 性能')
|
||
|
||
args = parser.parse_args()
|
||
|
||
if args.deepseek:
|
||
# 先测试 DeepSeek
|
||
print("\n" + "=" * 70)
|
||
print("🆚 性能对比测试")
|
||
print("=" * 70)
|
||
print("\n>>> 开始测试 DeepSeek API...")
|
||
ds_first, ds_total = asyncio.run(test_deepseek_first_token(args.question))
|
||
|
||
print("\n\n>>> 开始测试 MaxKB...")
|
||
maxkb_first, maxkb_total = asyncio.run(test_first_token_time(args.question))
|
||
|
||
# 对比结果
|
||
print("\n" + "=" * 70)
|
||
print("📊 性能对比结果")
|
||
print("=" * 70)
|
||
print(f"\n {'服务':<15} {'首个Token':<15} {'总响应时间':<15}")
|
||
print(f" {'-'*45}")
|
||
print(f" {'DeepSeek':<15} {f'{ds_first:.4f}s' if ds_first else 'N/A':<15} {f'{ds_total:.4f}s' if ds_total else 'N/A':<15}")
|
||
print(f" {'MaxKB':<15} {f'{maxkb_first:.4f}s' if maxkb_first else 'N/A':<15} {f'{maxkb_total:.4f}s' if maxkb_total else 'N/A':<15}")
|
||
print(f" {'-'*45}")
|
||
if ds_first and maxkb_first:
|
||
diff = maxkb_first - ds_first
|
||
print(f"\n ⚡ MaxKB 比 DeepSeek 慢: {diff:.4f} 秒 ({diff/ds_first*100:.1f}%)")
|
||
print("=" * 70)
|
||
elif args.times == 1:
|
||
asyncio.run(test_first_token_time(args.question))
|
||
else:
|
||
asyncio.run(run_multiple_tests(args.times, args.question))
|