easy-rag/tests/test_concurrent.py

481 lines
19 KiB
Python
Raw Blame History

This file contains invisible Unicode characters

This file contains invisible Unicode characters that are indistinguishable to humans but may be processed differently by a computer. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

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()