274 lines
10 KiB
Python
274 lines
10 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
RAG 性能分析脚本
|
||
|
||
分析 RAG 流程各阶段耗时，帮助定位性能瓶颈
|
||
"""
|
||
|
||
import sys
|
||
import os
|
||
import json
|
||
import time
|
||
import argparse
|
||
from datetime import datetime
|
||
from typing import Dict
|
||
|
||
# 添加项目根目录到路径
|
||
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||
sys.path.insert(0, PROJECT_ROOT)
|
||
|
||
|
||
def call_rag_stream_api(
|
||
question: str,
|
||
kb_name: str = "public_kb",
|
||
base_url: str = "http://localhost:5001",
|
||
) -> Dict:
|
||
"""
|
||
调用 RAG 流式 API 并收集各阶段事件
|
||
|
||
Returns:
|
||
包含各阶段耗时信息的字典
|
||
"""
|
||
import requests
|
||
|
||
url = f"{base_url.rstrip('/')}/rag"
|
||
headers = {
|
||
"Content-Type": "application/json",
|
||
"Authorization": "Bearer mock-token-admin"
|
||
}
|
||
data = {
|
||
"message": question,
|
||
"collections": [kb_name] if kb_name else []
|
||
}
|
||
|
||
events = []
|
||
timing = {
|
||
'total_duration_ms': 0,
|
||
'search_time_ms': 0,
|
||
'rerank_time_ms': 0,
|
||
'llm_time_ms': 0,
|
||
'stages': []
|
||
}
|
||
|
||
try:
|
||
start_time = time.time()
|
||
response = requests.post(url, json=data, headers=headers, stream=True, timeout=300)
|
||
response.raise_for_status()
|
||
|
||
for line in response.iter_lines():
|
||
if line:
|
||
line_str = line.decode('utf-8')
|
||
if line_str.startswith('data: '):
|
||
try:
|
||
event = json.loads(line_str[6:])
|
||
event['_received_at'] = time.time()
|
||
events.append(event)
|
||
|
||
# 分析 finish 事件中的耗时信息
|
||
if event.get('type') == 'finish':
|
||
timing['total_duration_ms'] = event.get('duration_ms', 0)
|
||
if 'timing' in event:
|
||
timing['search_time_ms'] = event['timing'].get('total_search_ms', 0)
|
||
timing['rerank_time_ms'] = event['timing'].get('rerank_ms', 0)
|
||
timing['rerank_cached'] = event['timing'].get('rerank_cached', False)
|
||
except json.JSONDecodeError:
|
||
continue
|
||
|
||
end_time = time.time()
|
||
timing['actual_elapsed_ms'] = int((end_time - start_time) * 1000)
|
||
|
||
# 计算 LLM 生成时间
|
||
if timing['total_duration_ms'] > 0 and timing['search_time_ms'] > 0:
|
||
timing['llm_time_ms'] = timing['total_duration_ms'] - timing['search_time_ms']
|
||
|
||
finish_event = next((event for event in events if event.get('type') == 'finish'), None)
|
||
error_event = next((event for event in events if event.get('type') == 'error'), None)
|
||
return {
|
||
'success': finish_event is not None and error_event is None,
|
||
'error': error_event.get('message', '') if error_event else (
|
||
'' if finish_event else 'SSE 流结束但未收到 finish 事件'
|
||
),
|
||
'events': events,
|
||
'timing': timing,
|
||
'answer': finish_event.get('answer', '') if finish_event else '',
|
||
}
|
||
|
||
except requests.exceptions.RequestException as e:
|
||
return {
|
||
'success': False,
|
||
'error': str(e),
|
||
'events': events,
|
||
'timing': timing
|
||
}
|
||
|
||
|
||
def analyze_performance(result: Dict, question: str) -> None:
|
||
"""分析并打印性能数据"""
|
||
print("\n" + "=" * 80)
|
||
print(f"问题: {question}")
|
||
print("=" * 80)
|
||
|
||
if not result['success']:
|
||
print(f"[X] 请求失败: {result.get('error', '未知错误')}")
|
||
return
|
||
|
||
timing = result['timing']
|
||
events = result['events']
|
||
|
||
# 打印各阶段事件时间线
|
||
print("\n[事件时间线]")
|
||
print("-" * 80)
|
||
first_event_time = None
|
||
for event in events:
|
||
event_type = event.get('type', 'unknown')
|
||
received_at = event.get('_received_at', 0)
|
||
|
||
if first_event_time is None:
|
||
first_event_time = received_at
|
||
relative_time = 0
|
||
else:
|
||
relative_time = (received_at - first_event_time) * 1000
|
||
|
||
if event_type == 'finish':
|
||
print(f" {relative_time:>8.0f}ms | {event_type:20s} | 总耗时: {event.get('duration_ms', 0)}ms")
|
||
elif event_type == 'sources':
|
||
sources = event.get('sources', [])
|
||
print(f" {relative_time:>8.0f}ms | {event_type:20s} | 找到 {len(sources)} 个来源")
|
||
elif event_type == 'chunks_retrieved':
|
||
chunks = event.get('data', {}).get('chunks', [])
|
||
print(f" {relative_time:>8.0f}ms | {event_type:20s} | 召回 {len(chunks)} 个切片")
|
||
elif event_type == 'chunk':
|
||
# 流式输出，只显示第一个
|
||
if not hasattr(analyze_performance, '_chunk_printed'):
|
||
print(f" {relative_time:>8.0f}ms | {event_type:20s} | 开始流式输出...")
|
||
analyze_performance._chunk_printed = True
|
||
else:
|
||
print(f" {relative_time:>8.0f}ms | {event_type:20s}")
|
||
|
||
if hasattr(analyze_performance, '_chunk_printed'):
|
||
delattr(analyze_performance, '_chunk_printed')
|
||
|
||
# 打印耗时统计
|
||
print("\n[耗时统计]")
|
||
print("-" * 80)
|
||
total = timing['total_duration_ms']
|
||
search = timing['search_time_ms']
|
||
rerank = timing['rerank_time_ms']
|
||
llm = timing['llm_time_ms']
|
||
actual = timing.get('actual_elapsed_ms', total)
|
||
|
||
if total > 0:
|
||
print(f" 总耗时 (API报告): {total:>8}ms ({total/1000:.2f}s)")
|
||
print(f" 实际耗时 (本地测量): {actual:>8}ms ({actual/1000:.2f}s)")
|
||
print()
|
||
print(f" 检索阶段: {search:>8}ms ({search/total*100:.1f}%)")
|
||
if rerank > 0:
|
||
cached_flag = " [缓存]" if timing.get('rerank_cached') else ""
|
||
print(f" 重排序阶段: {rerank:>8}ms ({rerank/total*100:.1f}%){cached_flag}")
|
||
print(f" LLM生成阶段: {llm:>8}ms ({llm/total*100:.1f}%)")
|
||
|
||
# 性能诊断
|
||
print("\n[性能诊断]")
|
||
print("-" * 80)
|
||
if total > 10000:
|
||
print(" [!] 总耗时超过 10 秒，需要优化")
|
||
|
||
if search > 3000:
|
||
print(f" [!] 检索耗时过长 ({search}ms)，可能原因:")
|
||
print(" - 向量库数据量过大")
|
||
print(" - 未命中查询缓存")
|
||
print(" - BM25 索引加载慢")
|
||
|
||
if rerank > 2000 and not timing.get('rerank_cached'):
|
||
print(f" [!] 重排序耗时过长 ({rerank}ms)，可能原因:")
|
||
print(" - Rerank 模型计算量大")
|
||
print(" - 候选切片数量过多")
|
||
|
||
if llm > 5000:
|
||
print(f" [!] LLM生成耗时过长 ({llm}ms)，可能原因:")
|
||
print(" - 模型生成速度慢")
|
||
print(" - 输出 token 数量多")
|
||
print(" - 网络延迟")
|
||
print(" 建议: 检查 MiMo thinking 是否关闭、上下文长度及 max_tokens")
|
||
|
||
if total < 5000:
|
||
print(" [OK] 性能良好")
|
||
else:
|
||
print(" [!] 未能获取耗时数据")
|
||
|
||
|
||
def run_performance_tests(
|
||
base_url: str = "http://localhost:5001",
|
||
kb_name: str = "public_kb",
|
||
question: str = None,
|
||
):
|
||
"""执行性能测试"""
|
||
# 测试问题（不同类型）
|
||
test_questions = [
|
||
{"question": "智启科技成立于哪一年？", "type": "精确匹配"},
|
||
{"question": "公司的愿景是什么？", "type": "语义理解"},
|
||
{"question": "公司有哪些分公司，分别在哪些城市？", "type": "跨文档关联"},
|
||
{"question": "一个入职3年的员工，累计病假1个月，能拿到多少病假工资？", "type": "复杂推理"},
|
||
{"question": "P4级工程师的年薪范围是多少？", "type": "表格数据"},
|
||
]
|
||
if question:
|
||
test_questions = [{"question": question, "type": "自定义问题"}]
|
||
|
||
print("=" * 80)
|
||
print("RAG Performance Analysis")
|
||
print(f"Time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
|
||
print("=" * 80)
|
||
|
||
results = []
|
||
for i, test in enumerate(test_questions):
|
||
print(f"\n\n[{i+1}/{len(test_questions)}] 类型: {test['type']}")
|
||
result = call_rag_stream_api(test['question'], kb_name=kb_name, base_url=base_url)
|
||
analyze_performance(result, test['question'])
|
||
results.append({
|
||
'question': test['question'],
|
||
'type': test['type'],
|
||
'success': result['success'],
|
||
'timing': result['timing']
|
||
})
|
||
|
||
# 汇总统计
|
||
print("\n\n" + "=" * 80)
|
||
print("[Performance Summary]")
|
||
print("=" * 80)
|
||
|
||
successful = [r for r in results if r['success']]
|
||
if successful:
|
||
total_times = [r['timing']['total_duration_ms'] for r in successful]
|
||
search_times = [r['timing']['search_time_ms'] for r in successful]
|
||
llm_times = [r['timing']['llm_time_ms'] for r in successful]
|
||
|
||
print(f"\n成功请求数: {len(successful)}/{len(results)}")
|
||
print(f"\n平均总耗时: {sum(total_times)/len(total_times):.0f}ms")
|
||
print(f"平均检索耗时: {sum(search_times)/len(search_times):.0f}ms")
|
||
print(f"平均LLM耗时: {sum(llm_times)/len(llm_times):.0f}ms")
|
||
print(f"\n最快总耗时: {min(total_times)}ms")
|
||
print(f"最慢总耗时: {max(total_times)}ms")
|
||
|
||
# 保存详细报告
|
||
report_dir = os.path.join(PROJECT_ROOT, ".data", "performance")
|
||
os.makedirs(report_dir, exist_ok=True)
|
||
report_file = os.path.join(
|
||
report_dir,
|
||
f"performance_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
|
||
)
|
||
with open(report_file, 'w', encoding='utf-8') as f:
|
||
json.dump({
|
||
'test_time': datetime.now().isoformat(),
|
||
'results': results
|
||
}, f, ensure_ascii=False, indent=2)
|
||
print(f"\n详细报告已保存: {report_file}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
parser = argparse.ArgumentParser(description="RAG SSE 链路耗时分析")
|
||
parser.add_argument("--base-url", default="http://localhost:5001")
|
||
parser.add_argument("--kb", default="public_kb", help="目标知识库名称")
|
||
parser.add_argument("--question", help="只测试一个自定义问题")
|
||
args = parser.parse_args()
|
||
run_performance_tests(args.base_url, args.kb, args.question)
|