# -*- 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)