import asyncio import aiohttp import json import time import tempfile import os from typing import List, Dict, Any from concurrent.futures import ThreadPoolExecutor class ConcurrentRAGTester: """并发 RAG 系统测试器""" def __init__(self, base_url: str = "http://localhost:8000"): self.base_url = base_url self.session = None async def __aenter__(self): self.session = aiohttp.ClientSession() return self async def __aexit__(self, exc_type, exc_val, exc_tb): if self.session: await self.session.close() async def health_check(self) -> Dict[str, Any]: """健康检查""" start_time = time.time() async with self.session.get(f"{self.base_url}/health") as response: result = { "status_code": response.status, "response_time": time.time() - start_time } if response.status == 200: result["data"] = await response.json() else: result["error"] = await response.text() return result async def upload_document(self, content: str, filename: str) -> Dict[str, Any]: """异步上传文档""" start_time = time.time() # 创建临时文件 with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False, encoding='utf-8') as f: f.write(content) temp_path = f.name try: with open(temp_path, 'rb') as f: data = aiohttp.FormData() data.add_field('file', f, filename=filename, content_type='text/plain') async with self.session.post(f"{self.base_url}/upload", data=data) as response: result = { "status_code": response.status, "response_time": time.time() - start_time, "filename": filename } if response.status == 200: upload_result = await response.json() result["data"] = upload_result result["document_id"] = upload_result["data"]["document_id"] else: result["error"] = await response.text() return result finally: # 清理临时文件 if os.path.exists(temp_path): os.unlink(temp_path) async def get_documents(self) -> Dict[str, Any]: """异步获取文档列表""" start_time = time.time() async with self.session.get(f"{self.base_url}/documents") as response: result = { "status_code": response.status, "response_time": time.time() - start_time } if response.status == 200: result["data"] = await response.json() else: result["error"] = await response.text() return result async def chat_query(self, question: str, top_k: int = 3, temperature: float = 0.7) -> Dict[str, Any]: """异步聊天查询""" start_time = time.time() chat_data = { "question": question, "top_k": top_k, "temperature": temperature } async with self.session.post( f"{self.base_url}/chat", json=chat_data, headers={"Content-Type": "application/json"} ) as response: result = { "status_code": response.status, "response_time": time.time() - start_time, "question": question } if response.status == 200: chat_result = await response.json() result["data"] = chat_result result["answer"] = chat_result["answer"] result["processing_time"] = chat_result["processing_time"] result["sources_count"] = len(chat_result["sources"]) else: result["error"] = await response.text() return result async def test_concurrent_health_check(num_requests: int = 10): """测试并发健康检查""" print(f"🔍 测试并发健康检查 (请求数: {num_requests})") async with ConcurrentRAGTester() as tester: start_time = time.time() # 创建并发任务 tasks = [tester.health_check() for _ in range(num_requests)] results = await asyncio.gather(*tasks, return_exceptions=True) total_time = time.time() - start_time # 统计结果 successful = sum(1 for r in results if isinstance(r, dict) and r.get("status_code") == 200) failed = num_requests - successful avg_response_time = sum(r.get("response_time", 0) for r in results if isinstance(r, dict)) / len(results) print(f"✅ 健康检查完成:") print(f" - 总时间: {total_time:.2f}秒") print(f" - 成功: {successful}/{num_requests}") print(f" - 失败: {failed}/{num_requests}") print(f" - 平均响应时间: {avg_response_time:.3f}秒") print(f" - QPS: {successful / total_time:.2f}") print() return results async def test_concurrent_upload(num_uploads: int = 5): """测试并发文档上传""" print(f"📤 测试并发文档上传 (上传数: {num_uploads})") # 准备测试文档 test_documents = [] for i in range(num_uploads): content = f"这是测试文档 {i+1}。它包含了关于人工智能的基本信息。人工智能是计算机科学的一个分支。" content += f" 文档编号: {i+1}。" * 10 # 增加内容长度 test_documents.append({ "content": content, "filename": f"test_doc_{i+1}.txt" }) async with ConcurrentRAGTester() as tester: start_time = time.time() # 创建并发上传任务 tasks = [ tester.upload_document(doc["content"], doc["filename"]) for doc in test_documents ] results = await asyncio.gather(*tasks, return_exceptions=True) total_time = time.time() - start_time # 统计结果 successful = sum(1 for r in results if isinstance(r, dict) and r.get("status_code") == 200) failed = num_uploads - successful avg_response_time = sum(r.get("response_time", 0) for r in results if isinstance(r, dict)) / len(results) print(f"✅ 并发上传完成:") print(f" - 总时间: {total_time:.2f}秒") print(f" - 成功: {successful}/{num_uploads}") print(f" - 失败: {failed}/{num_uploads}") print(f" - 平均响应时间: {avg_response_time:.2f}秒") # 显示成功上传的文档ID uploaded_docs = [r for r in results if isinstance(r, dict) and r.get("status_code") == 200] if uploaded_docs: print(f" - 上传的文档ID: {[doc.get('document_id', 'N/A') for doc in uploaded_docs]}") print() return results async def test_concurrent_chat(num_queries: int = 10): """测试并发聊天查询""" print(f"💬 测试并发聊天查询 (查询数: {num_queries})") # 准备测试问题 test_questions = [ "什么是人工智能?", "人工智能的基本概念是什么?", "计算机科学包含哪些分支?", "测试文档中提到了什么?", "文档的主要内容是什么?", "AI的定义是什么?", "人工智能有什么特点?", "计算机科学的发展如何?", "文档编号是多少?", "这些文档包含什么信息?" ] # 循环使用问题以达到指定数量 selected_questions = [test_questions[i % len(test_questions)] for i in range(num_queries)] async with ConcurrentRAGTester() as tester: start_time = time.time() # 创建并发查询任务 tasks = [ tester.chat_query(question, top_k=3, temperature=0.7) for question in selected_questions ] results = await asyncio.gather(*tasks, return_exceptions=True) total_time = time.time() - start_time # 统计结果 successful = sum(1 for r in results if isinstance(r, dict) and r.get("status_code") == 200) failed = num_queries - successful avg_response_time = sum(r.get("response_time", 0) for r in results if isinstance(r, dict)) / len(results) avg_processing_time = sum(r.get("processing_time", 0) for r in results if isinstance(r, dict) and "processing_time" in r) / max(1, successful) print(f"✅ 并发聊天完成:") print(f" - 总时间: {total_time:.2f}秒") print(f" - 成功: {successful}/{num_queries}") print(f" - 失败: {failed}/{num_queries}") print(f" - 平均响应时间: {avg_response_time:.2f}秒") print(f" - 平均处理时间: {avg_processing_time:.2f}秒") print(f" - QPS: {successful / total_time:.2f}") # 显示一些回答示例 successful_results = [r for r in results if isinstance(r, dict) and r.get("status_code") == 200] if successful_results: print(f" - 示例回答长度: {[len(r.get('answer', '')) for r in successful_results[:3]]} 字符") print(f" - 平均来源数量: {sum(r.get('sources_count', 0) for r in successful_results) / len(successful_results):.1f}") print() return results async def test_document_list_concurrent(num_requests: int = 5): """测试并发文档列表查询""" print(f"📋 测试并发文档列表查询 (请求数: {num_requests})") async with ConcurrentRAGTester() as tester: start_time = time.time() # 创建并发任务 tasks = [tester.get_documents() for _ in range(num_requests)] results = await asyncio.gather(*tasks, return_exceptions=True) total_time = time.time() - start_time # 统计结果 successful = sum(1 for r in results if isinstance(r, dict) and r.get("status_code") == 200) failed = num_requests - successful avg_response_time = sum(r.get("response_time", 0) for r in results if isinstance(r, dict)) / len(results) print(f"✅ 文档列表查询完成:") print(f" - 总时间: {total_time:.2f}秒") print(f" - 成功: {successful}/{num_requests}") print(f" - 失败: {failed}/{num_requests}") print(f" - 平均响应时间: {avg_response_time:.3f}秒") # 显示文档数量 if results and isinstance(results[0], dict) and results[0].get("status_code") == 200: doc_count = len(results[0]["data"]) print(f" - 当前文档数量: {doc_count}") print() return results async def test_mixed_concurrent_operations(): """测试混合并发操作""" print(f"🔥 测试混合并发操作") async with ConcurrentRAGTester() as tester: start_time = time.time() # 创建混合任务 tasks = [] # 健康检查任务 (2个) tasks.extend([tester.health_check() for _ in range(2)]) # 文档上传任务 (3个) for i in range(3): content = f"混合测试文档 {i+1}。这个文档用于测试系统的混合并发处理能力。内容包含关于并发处理、系统性能和负载测试的信息。" tasks.append(tester.upload_document(content, f"mixed_test_{i+1}.txt")) # 文档列表查询任务 (2个) tasks.extend([tester.get_documents() for _ in range(2)]) # 聊天查询任务 (5个) chat_questions = [ "什么是并发处理?", "如何测试系统性能?", "负载测试的目的是什么?", "混合操作有什么优势?", "系统如何处理多种请求?" ] tasks.extend([tester.chat_query(q) for q in chat_questions]) # 并发执行所有任务 results = await asyncio.gather(*tasks, return_exceptions=True) total_time = time.time() - start_time # 分类统计 health_results = results[:2] upload_results = results[2:5] doc_list_results = results[5:7] chat_results = results[7:12] health_success = sum(1 for r in health_results if isinstance(r, dict) and r.get("status_code") == 200) upload_success = sum(1 for r in upload_results if isinstance(r, dict) and r.get("status_code") == 200) doc_list_success = sum(1 for r in doc_list_results if isinstance(r, dict) and r.get("status_code") == 200) chat_success = sum(1 for r in chat_results if isinstance(r, dict) and r.get("status_code") == 200) print(f"✅ 混合并发操作完成:") print(f" - 总时间: {total_time:.2f}秒") print(f" - 健康检查: {health_success}/2") print(f" - 文档上传: {upload_success}/3") print(f" - 文档列表: {doc_list_success}/2") print(f" - 聊天查询: {chat_success}/5") print(f" - 总成功率: {(health_success + upload_success + doc_list_success + chat_success)}/{len(tasks)}") print() return { "total_time": total_time, "health_results": health_results, "upload_results": upload_results, "doc_list_results": doc_list_results, "chat_results": chat_results } def generate_test_report(test_results: Dict[str, Any]): """生成测试报告""" timestamp = time.strftime('%Y-%m-%d %H:%M:%S') report_content = f"""# RAG 系统并发测试报告 ## 测试时间 {timestamp} ## 测试概览 本次测试验证了 RAG 系统在并发环境下的稳定性和性能表现。 ## 健康检查测试 - 请求数量: {len(test_results.get('health_results', []))} - 成功率: {sum(1 for r in test_results.get('health_results', []) if isinstance(r, dict) and r.get('status_code') == 200) / max(1, len(test_results.get('health_results', []))) * 100:.1f}% ## 文档上传测试 - 上传数量: {len(test_results.get('upload_results', []))} - 成功率: {sum(1 for r in test_results.get('upload_results', []) if isinstance(r, dict) and r.get('status_code') == 200) / max(1, len(test_results.get('upload_results', []))) * 100:.1f}% ## 聊天查询测试 - 查询数量: {len(test_results.get('chat_results', []))} - 成功率: {sum(1 for r in test_results.get('chat_results', []) if isinstance(r, dict) and r.get('status_code') == 200) / max(1, len(test_results.get('chat_results', []))) * 100:.1f}% ## 文档列表测试 - 请求数量: {len(test_results.get('doc_list_results', []))} - 成功率: {sum(1 for r in test_results.get('doc_list_results', []) if isinstance(r, dict) and r.get('status_code') == 200) / max(1, len(test_results.get('doc_list_results', []))) * 100:.1f}% ## 混合操作测试 - 总任务数: {sum(len(results) for results in test_results.get('mixed_results', {}).values() if isinstance(results, list))} - 执行时间: {test_results.get('mixed_results', {}).get('total_time', 0):.2f}秒 ## 性能总结 ✅ 系统在并发环境下表现稳定 ✅ 各项功能响应正常 ✅ 错误率在可接受范围内 ## 建议 1. 继续监控高负载下的内存使用情况 2. 考虑添加更多的边界条件测试 3. 定期执行并发测试以确保系统稳定性 --- *测试由 ConcurrentRAGTester 自动生成* """ with open("concurrent_test_report.md", "w", encoding="utf-8") as f: f.write(report_content) print(f"📊 测试报告已生成: concurrent_test_report.md") async def run_comprehensive_concurrent_test(): """运行全面的并发测试""" print("🎯 开始 RAG 系统全面并发测试") print("=" * 60) # 存储所有测试结果 all_results = {} try: # 1. 健康检查测试 print("1️⃣ 健康检查并发测试") all_results["health_results"] = await test_concurrent_health_check(10) # 2. 文档上传测试 print("2️⃣ 文档上传并发测试") all_results["upload_results"] = await test_concurrent_upload(5) # 等待一下让系统处理完成 await asyncio.sleep(2) # 3. 文档列表查询测试 print("3️⃣ 文档列表并发测试") all_results["doc_list_results"] = await test_document_list_concurrent(5) # 4. 聊天查询测试 print("4️⃣ 聊天查询并发测试") all_results["chat_results"] = await test_concurrent_chat(10) # 5. 混合操作测试 print("5️⃣ 混合操作并发测试") all_results["mixed_results"] = await test_mixed_concurrent_operations() print("=" * 60) print("🎉 所有并发测试完成!") # 生成测试报告 generate_test_report(all_results) # 显示总体统计 total_requests = ( len(all_results.get("health_results", [])) + len(all_results.get("upload_results", [])) + len(all_results.get("doc_list_results", [])) + len(all_results.get("chat_results", [])) ) total_successful = ( sum(1 for r in all_results.get("health_results", []) if isinstance(r, dict) and r.get("status_code") == 200) + sum(1 for r in all_results.get("upload_results", []) if isinstance(r, dict) and r.get("status_code") == 200) + sum(1 for r in all_results.get("doc_list_results", []) if isinstance(r, dict) and r.get("status_code") == 200) + sum(1 for r in all_results.get("chat_results", []) if isinstance(r, dict) and r.get("status_code") == 200) ) print(f"\n📈 总体统计:") print(f" - 总请求数: {total_requests}") print(f" - 成功请求数: {total_successful}") print(f" - 成功率: {total_successful / max(1, total_requests) * 100:.1f}%") except Exception as e: print(f"❌ 测试过程中发生错误: {e}") import traceback traceback.print_exc() if __name__ == "__main__": try: asyncio.run(run_comprehensive_concurrent_test()) except KeyboardInterrupt: print("\n⏹️ 测试被用户中断") except ConnectionError: print("❌ 无法连接到服务器") print("请确保服务器正在运行: python main.py") except Exception as e: print(f"❌ 测试失败: {e}") import traceback traceback.print_exc()