481 lines
19 KiB
Python
481 lines
19 KiB
Python
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()
|