"""性能测试:并发会话、长文档解析、OCR 耗时、API 响应时间基线。 用法: python3 test_performance.py 输出: perf_report.json + 控制台统计 """ import json import time import os import sys import threading from concurrent.futures import ThreadPoolExecutor, as_completed sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from fastapi.testclient import TestClient import api as api_module from api import app client = TestClient(app) def perf_single_session_ask(): """单会话 ask 响应时间(不含 Kimi,只测分诊+RAG)。""" # 清理限流桶 api_module._rate_buckets.clear() sid = client.post('/sessions', json={'mode': 'normal'}).json()['session_id'] times = [] questions = [ '员工工伤后能否开除', '合同违约金怎么约定', '客户欠款三个月不还', '收到税务通知要求补缴', '最近压力大想聊聊', ] for q in questions: start = time.time() r = client.post(f'/sessions/{sid}/ask', json={'question': q, 'scene': 'legal'}) elapsed = time.time() - start times.append(elapsed) assert r.status_code == 200, f"ask 失败: {r.status_code}" return { "test": "single_session_ask", "iterations": len(times), "avg_ms": round(sum(times) / len(times) * 1000, 1), "min_ms": round(min(times) * 1000, 1), "max_ms": round(max(times) * 1000, 1), } def perf_concurrent_sessions(): """并发会话创建 + 提问(10 并发)。""" api_module._rate_buckets.clear() def worker(i): start = time.time() sid = client.post('/sessions', json={'mode': 'normal'}).json()['session_id'] r = client.post(f'/sessions/{sid}/ask', json={ 'question': f'测试问题 {i}', 'scene': 'legal'}) elapsed = time.time() - start return elapsed, r.status_code with ThreadPoolExecutor(max_workers=10) as executor: futures = [executor.submit(worker, i) for i in range(10)] results = [f.result() for f in as_completed(futures)] times = [r[0] for r in results] status_codes = [r[1] for r in results] return { "test": "concurrent_sessions_10", "concurrency": 10, "avg_ms": round(sum(times) / len(times) * 1000, 1), "min_ms": round(min(times) * 1000, 1), "max_ms": round(max(times) * 1000, 1), "all_200": all(c == 200 for c in status_codes), "status_codes": status_codes, } def perf_long_document_parse(): """长文档解析:100KB / 1MB TXT 提取时间。""" api_module._rate_buckets.clear() sid = client.post('/sessions', json={'mode': 'normal'}).json()['session_id'] results = [] for size_kb in [10, 100, 500]: content = '合同条款测试内容。' * (size_kb * 1024 // 10) data = content.encode('utf-8') start = time.time() r = client.post(f'/sessions/{sid}/file', files={'file': (f'test_{size_kb}kb.txt', data)}) elapsed = time.time() - start assert r.status_code == 200, f"文件上传失败: {r.status_code}" results.append({ "size_kb": size_kb, "parse_ms": round(elapsed * 1000, 1), "text_length": r.json().get('text_length', 0), }) return {"test": "long_document_parse", "results": results} def perf_ocr_image(): """OCR 图片识别耗时。""" api_module._rate_buckets.clear() try: from PIL import Image, ImageDraw, ImageFont import io import pytesseract except ImportError: return {"test": "ocr_image", "skipped": "PIL/pytesseract not installed"} sid = client.post('/sessions', json={'mode': 'normal'}).json()['session_id'] # 生成不同复杂度的图片 results = [] for num_lines in [1, 5, 20]: img = Image.new('RGB', (800, 60 * num_lines + 40), 'white') d = ImageDraw.Draw(img) try: font = ImageFont.truetype('/System/Library/Fonts/STHeiti Medium.ttc', 36) except Exception: font = ImageFont.load_default() for i in range(num_lines): d.text((20, 20 + i * 60), f'合同第{i+1}条违约金条款', fill='black', font=font) buf = io.BytesIO() img.save(buf, format='PNG') data = buf.getvalue() start = time.time() r = client.post(f'/sessions/{sid}/file', files={'file': (f'test_{num_lines}lines.png', data)}) elapsed = time.time() - start if r.status_code == 200: results.append({ "lines": num_lines, "ocr_ms": round(elapsed * 1000, 1), "text_length": r.json().get('text_length', 0), }) else: results.append({"lines": num_lines, "error": r.status_code}) return {"test": "ocr_image", "results": results} def perf_rag_search(): """RAG 检索耗时:关键词 + 向量。""" api_module._rate_buckets.clear() queries = ['工伤解除', '合同违约金', '税务稽查', '欠款催收', '劳动仲裁'] keyword_times = [] vector_times = [] for q in queries: start = time.time() r = client.get(f'/knowledge/search?q={q}') keyword_times.append(time.time() - start) assert r.status_code == 200 start = time.time() r = client.get(f'/knowledge/vector-search?q={q}') vector_times.append(time.time() - start) assert r.status_code == 200 return { "test": "rag_search", "keyword_avg_ms": round(sum(keyword_times) / len(keyword_times) * 1000, 1), "vector_avg_ms": round(sum(vector_times) / len(vector_times) * 1000, 1), "queries": len(queries), } def perf_session_destroy(): """会话销毁耗时。""" times = [] for _ in range(20): api_module._rate_buckets.clear() sid = client.post('/sessions', json={'mode': 'ephemeral'}).json()['session_id'] start = time.time() client.post(f'/sessions/{sid}/destroy') times.append(time.time() - start) return { "test": "session_destroy", "iterations": len(times), "avg_ms": round(sum(times) / len(times) * 1000, 1), "min_ms": round(min(times) * 1000, 1), "max_ms": round(max(times) * 1000, 1), } def run_all(): """运行全部性能测试。""" print("=" * 60) print("BossAi 性能测试") print("=" * 60) tests = [ ("单会话 ask 响应时间", perf_single_session_ask), ("并发 10 会话", perf_concurrent_sessions), ("长文档解析", perf_long_document_parse), ("OCR 图片识别", perf_ocr_image), ("RAG 检索", perf_rag_search), ("会话销毁", perf_session_destroy), ] all_results = [] for name, test_func in tests: print(f"\n--- {name} ---") try: result = test_func() all_results.append(result) # 简洁输出 for k, v in result.items(): if k != "test" and not isinstance(v, list): print(f" {k}: {v}") elif isinstance(v, list) and k == "results": for item in v: print(f" {item}") except Exception as e: print(f" ERROR: {e}") all_results.append({"test": name, "error": str(e)}) # 写入报告 with open("perf_report.json", "w", encoding="utf-8") as f: json.dump(all_results, f, ensure_ascii=False, indent=2) print(f"\n{'=' * 60}") print(f"性能报告已写入 perf_report.json") print(f"{'=' * 60}") # 性能基线总结 print("\n性能基线总结:") for r in all_results: if "error" in r: print(f" {r['test']}: ERROR") elif "avg_ms" in r: print(f" {r['test']}: avg {r['avg_ms']}ms") elif "keyword_avg_ms" in r: print(f" {r['test']}: keyword {r['keyword_avg_ms']}ms, vector {r['vector_avg_ms']}ms") if __name__ == "__main__": run_all()