506 lines
18 KiB
Python
506 lines
18 KiB
Python
import argparse # 解析命令行参数
|
||
import asyncio # 异步事件循环
|
||
import json # 处理 JSON 文本
|
||
import os # 读取环境变量
|
||
import sys # 系统参数
|
||
import time # 时间统计
|
||
import uuid # 生成随机 chat/session id
|
||
|
||
import aiohttp # 异步 HTTP 客户端
|
||
import requests # 同步 HTTP 客户端
|
||
|
||
# 下面的默认参数可以通过环境变量覆盖，避免把机密写死在代码里。
|
||
API_URL = os.environ.get(
|
||
"RAGFLOW_URL",
|
||
"http://10.0.0.202:9099/system/ragflow/converse_with_chat_assistant",
|
||
)
|
||
AUTH_TOKEN = os.environ.get(
|
||
"RAGFLOW_TOKEN",
|
||
"Bearer ",
|
||
)
|
||
DEFAULT_CHAT_ID = os.environ.get(
|
||
"RAGFLOW_CHAT_ID", "a455abcadbcb11f0884b0242ac130006"
|
||
)
|
||
DEFAULT_SESSION_ID = os.environ.get(
|
||
"RAGFLOW_SESSION_ID", "3d3ec1e2b9344591a1486fc95c1e01ed"
|
||
)
|
||
DEFAULT_QUESTION = os.environ.get(
|
||
"RAGFLOW_QUESTION", "你好，请用简洁的语言介绍你自己"
|
||
)
|
||
LOGIN_URL = os.environ.get(
|
||
"RAGFLOW_LOGIN_URL",
|
||
"http://10.0.0.202:9099/login",
|
||
)
|
||
|
||
# 登录凭据
|
||
LOGIN_USERNAME = os.environ.get("RAGFLOW_USERNAME", "admin")
|
||
LOGIN_PASSWORD = os.environ.get("RAGFLOW_PASSWORD", "admin123")
|
||
|
||
|
||
class PerformanceTimer:
|
||
"""性能计时器类"""
|
||
|
||
def __init__(self):
|
||
self.start_time = None
|
||
self.end_time = None
|
||
self.elapsed_ms = 0
|
||
|
||
def start(self):
|
||
self.start_time = time.perf_counter()
|
||
return self
|
||
|
||
def restart(self):
|
||
"""重新开始计时"""
|
||
self.start_time = time.perf_counter()
|
||
self.end_time = None
|
||
self.elapsed_ms = 0
|
||
return self
|
||
|
||
def stop(self):
|
||
if self.start_time is None:
|
||
raise ValueError("计时器未启动")
|
||
self.end_time = time.perf_counter()
|
||
self.elapsed_ms = (self.end_time - self.start_time) * 1000
|
||
return self
|
||
|
||
def get_elapsed_ms(self) -> float:
|
||
if self.start_time is not None and self.end_time is not None:
|
||
return self.elapsed_ms
|
||
elif self.start_time is not None:
|
||
# 如果没有stop，返回当前经过的时间
|
||
return (time.perf_counter() - self.start_time) * 1000
|
||
return 0
|
||
|
||
def get_elapsed_seconds(self) -> float:
|
||
return self.get_elapsed_ms() / 1000
|
||
|
||
def format(self) -> str:
|
||
"""格式化输出时间"""
|
||
elapsed = self.get_elapsed_ms()
|
||
if elapsed < 1000:
|
||
return f"{elapsed:.2f}ms"
|
||
else:
|
||
return f"{self.get_elapsed_seconds():.2f}s ({elapsed:.0f}ms)"
|
||
|
||
|
||
async def login_get_token() -> str:
|
||
"""
|
||
登录获取token
|
||
"""
|
||
print("🔐 正在登录获取token...")
|
||
|
||
payload = {
|
||
"username": LOGIN_USERNAME,
|
||
"password": LOGIN_PASSWORD
|
||
}
|
||
|
||
timer = PerformanceTimer().start()
|
||
|
||
async with aiohttp.ClientSession() as session:
|
||
async with session.post(
|
||
LOGIN_URL, data=payload, timeout=60
|
||
) as resp:
|
||
timer.stop()
|
||
|
||
if resp.status != 200:
|
||
print(f"❌ 登录失败，HTTP状态码: {resp.status}")
|
||
error_text = await resp.text()
|
||
print(f"错误响应: {error_text}")
|
||
raise Exception(f"登录失败: {resp.status}")
|
||
|
||
response_data = await resp.json()
|
||
|
||
# 检查响应格式
|
||
if response_data.get("code") != 200:
|
||
print(f"❌ 登录失败，响应码: {response_data.get('code')}")
|
||
print(f"响应信息: {response_data.get('msg', '未知错误')}")
|
||
raise Exception(f"登录失败: {response_data.get('msg')}")
|
||
|
||
token = response_data.get("token")
|
||
|
||
if not token:
|
||
print("❌ 登录成功但未获取到token")
|
||
print(f"完整响应: {response_data}")
|
||
raise Exception("登录成功但未获取到token")
|
||
|
||
print(f"✅ 登录成功，耗时: {timer.format()}")
|
||
print(f"📝 Token预览: {token[:50]}...")
|
||
|
||
return token
|
||
|
||
|
||
async def stream_chat_async(
|
||
question: str,
|
||
chat_id: str,
|
||
session_id: str,
|
||
request_number: int = 1
|
||
) -> dict:
|
||
"""
|
||
异步版（aiohttp），适合已有事件循环的应用，不阻塞主线程。
|
||
返回完整回答和相关时间统计
|
||
"""
|
||
print(f"\n{'='*60}")
|
||
print(f"📤 请求 #{request_number}")
|
||
print(f"❓ 问题: {question}")
|
||
print(f"🆔 Chat ID: {chat_id}")
|
||
print(f"🔗 Session ID: {session_id}")
|
||
print(f"{'='*60}")
|
||
|
||
# 创建计时器
|
||
timer = PerformanceTimer().start()
|
||
|
||
token = await login_get_token()
|
||
|
||
# 重启计时器，只计算流式请求的时间（不包含登录时间）
|
||
timer.restart()
|
||
|
||
payload = {
|
||
"chatId": chat_id,
|
||
"question": question,
|
||
"stream": True,
|
||
"sessionId": session_id,
|
||
}
|
||
headers = {
|
||
"Authorization": f"Bearer {token}",
|
||
"Accept": "text/event-stream",
|
||
}
|
||
|
||
answer_parts = []
|
||
event_type = "message"
|
||
|
||
# 性能计时 - timer 已在登录前创建并 restart
|
||
first_token_time = None
|
||
connection_time = None
|
||
|
||
async with aiohttp.ClientSession() as session:
|
||
async with session.post(
|
||
API_URL, json=payload, headers=headers, timeout=120
|
||
) as resp:
|
||
print(
|
||
f"# status={resp.status} content-type={resp.headers.get('content-type')}")
|
||
resp.raise_for_status()
|
||
|
||
# 记录连接建立时间
|
||
connection_time = timer.get_elapsed_ms()
|
||
print(f"[debug] 连接建立时间: {connection_time:.0f}ms")
|
||
|
||
event_count = 0
|
||
while True:
|
||
raw_line = await resp.content.readline()
|
||
if not raw_line:
|
||
print(f"[debug] 收到 {event_count} 个事件，流结束")
|
||
break
|
||
line = raw_line.decode(errors="ignore").strip() # 逐行读取 SSE 数据
|
||
if not line or line.startswith(":"):
|
||
continue
|
||
if line.startswith("event:"):
|
||
event_type = line.split(":", 1)[1].strip() or "message"
|
||
print(f"[debug] 事件类型: {event_type}")
|
||
continue
|
||
if not line.startswith("data:"):
|
||
print(f"[skip] {line}")
|
||
continue
|
||
data_str = line[len("data:"):].strip()
|
||
if not data_str:
|
||
continue
|
||
event_count += 1
|
||
try:
|
||
payload_obj = json.loads(data_str)
|
||
except json.JSONDecodeError:
|
||
print(f"[{event_type}] {data_str}")
|
||
event_type = "message"
|
||
continue
|
||
|
||
# 打印所有收到的响应，便于调试
|
||
print(f"[debug] 收到事件 #{event_count}: {list(payload_obj.keys())}")
|
||
|
||
# 记录首Token时间（遇到任何包含内容的字段时）
|
||
if first_token_time is None:
|
||
# 检测缓存命中情况
|
||
if "from_cache" in payload_obj and payload_obj.get("from_cache") is True:
|
||
# 缓存命中！首Token就是连接建立时间
|
||
first_token_time = connection_time
|
||
cache_total_time = payload_obj.get("total_time")
|
||
print(f"[debug] 🔥 缓存命中! 连接时间: {connection_time:.0f}ms, 缓存总耗时: {cache_total_time}ms")
|
||
elif "answer" in payload_obj:
|
||
first_token_time = timer.get_elapsed_ms()
|
||
print(f"[debug] 首Token时间(answer): {first_token_time:.0f}ms")
|
||
elif "data" in payload_obj and payload_obj.get("data") is not True:
|
||
# 缓存命中时可能直接返回完整答案
|
||
first_token_time = timer.get_elapsed_ms()
|
||
print(f"[debug] 首Token时间(data): {first_token_time:.0f}ms")
|
||
elif payload_obj.get("status") == "completed" or event_type == "end":
|
||
# 如果一开始就收到completed，可能是即时响应
|
||
first_token_time = connection_time # 使用连接时间作为近似
|
||
print(f"[debug] 首Token时间(即时响应): {first_token_time:.0f}ms")
|
||
|
||
if event_type == "end" or payload_obj.get("status") == "completed":
|
||
print("\n[stream] completed")
|
||
break
|
||
if payload_obj.get("data") is True:
|
||
event_type = "message"
|
||
continue
|
||
if "answer" in payload_obj:
|
||
piece = payload_obj.get("answer", "")
|
||
answer_parts.append(piece)
|
||
# 流式输出：每个片段都单独打印
|
||
print(piece, end="", flush=True)
|
||
await asyncio.sleep(0.05) # 使用异步sleep而不是time.sleep
|
||
else:
|
||
print(f"[{event_type}] {payload_obj}")
|
||
event_type = "message"
|
||
|
||
timer.stop()
|
||
full_answer = "".join(answer_parts)
|
||
|
||
print(f"\n\n{'='*60}")
|
||
print(f"✅ 请求 #{request_number} 完成")
|
||
print(f"⏱️ 总耗时: {timer.format()}")
|
||
if connection_time is not None:
|
||
print(f"🔗 连接建立: {connection_time:.0f}ms")
|
||
if first_token_time is not None:
|
||
print(f"🚀 首Token延迟: {first_token_time:.0f}ms")
|
||
if connection_time is not None:
|
||
ttft = first_token_time - connection_time # Time To First Token
|
||
print(f"⏱️ 首Token耗时(不含连接): {ttft:.0f}ms")
|
||
else:
|
||
print(f"🚀 首Token延迟: 未记录")
|
||
print(f"📝 回答长度: {len(full_answer)} 字符")
|
||
print(f"{'='*60}")
|
||
|
||
# 返回完整统计信息
|
||
return {
|
||
"answer": full_answer,
|
||
"timer": timer,
|
||
"first_token_ms": first_token_time,
|
||
"connection_ms": connection_time
|
||
}
|
||
|
||
|
||
async def run_performance_test(
|
||
questions: list[str],
|
||
chat_id: str,
|
||
session_id: str,
|
||
delay_between_requests: float = 2.0
|
||
):
|
||
"""
|
||
运行性能测试，多次请求以测试缓存效果
|
||
|
||
Args:
|
||
questions: 问题列表，每个请求使用对应的问题
|
||
chat_id: 聊天会话ID
|
||
session_id: 会话ID
|
||
delay_between_requests: 请求之间的间隔时间(秒)
|
||
"""
|
||
num_requests = len(questions)
|
||
|
||
print("\n" + "🚀"*30)
|
||
print("🎯 性能测试开始")
|
||
print(f"📊 测试配置:")
|
||
print(f" - 请求次数: {num_requests}")
|
||
print(f" - 请求间隔: {delay_between_requests}s")
|
||
print("📝 测试问题:")
|
||
for i, q in enumerate(questions):
|
||
print(f" {i+1}. {q}")
|
||
print("🚀"*30)
|
||
|
||
results = []
|
||
|
||
for i in range(1, num_requests + 1):
|
||
question = questions[i-1]
|
||
try:
|
||
result = await stream_chat_async(
|
||
question=question,
|
||
chat_id=chat_id,
|
||
session_id=session_id,
|
||
request_number=i
|
||
)
|
||
results.append({
|
||
"request_number": i,
|
||
"question": question,
|
||
"elapsed_ms": result["timer"].get_elapsed_ms(),
|
||
"elapsed_seconds": result["timer"].get_elapsed_seconds(),
|
||
"first_token_ms": result["first_token_ms"],
|
||
"connection_ms": result["connection_ms"],
|
||
"answer_length": len(result["answer"]),
|
||
"from_cache": False # 后续可以根据响应判断
|
||
})
|
||
|
||
# 如果不是最后一次请求，等待一段时间
|
||
if i < num_requests:
|
||
print(f"\n⏳ 等待 {delay_between_requests}s 后进行下一次请求...")
|
||
await asyncio.sleep(delay_between_requests)
|
||
|
||
except Exception as e:
|
||
print(f"\n❌ 请求 #{i} 失败: {str(e)}")
|
||
results.append({
|
||
"request_number": i,
|
||
"question": question,
|
||
"elapsed_ms": -1,
|
||
"elapsed_seconds": -1,
|
||
"first_token_ms": None,
|
||
"connection_ms": None,
|
||
"answer_length": 0,
|
||
"error": str(e)
|
||
})
|
||
|
||
# 打印性能测试汇总
|
||
print("\n" + "📊"*30)
|
||
print("📈 性能测试结果汇总")
|
||
print("📊"*30)
|
||
|
||
valid_results = [r for r in results if r["elapsed_ms"] > 0]
|
||
|
||
if not valid_results:
|
||
print("❌ 所有请求都失败了")
|
||
return
|
||
|
||
elapsed_times = [r["elapsed_ms"] for r in valid_results]
|
||
avg_time = sum(elapsed_times) / len(elapsed_times)
|
||
min_time = min(elapsed_times)
|
||
max_time = max(elapsed_times)
|
||
|
||
print(f"\n{'请求#':<8} {'问题':<30} {'总耗时':<12} {'首Token':<10} {'速度评估'}")
|
||
print("-" * 80)
|
||
|
||
for r in results:
|
||
if r["elapsed_ms"] < 0:
|
||
print(f"#{r['request_number']:<6} {r['question'][:28]:<30} {'失败':<12} {'-':<10} ❌")
|
||
else:
|
||
# 速度评估
|
||
elapsed = r["elapsed_ms"]
|
||
if elapsed < 500:
|
||
speed = "🚀 极快 (缓存?)"
|
||
elif elapsed < 1500:
|
||
speed = "⚡ 较快"
|
||
elif elapsed < 3000:
|
||
speed = "⏱️ 正常"
|
||
else:
|
||
speed = "🐢 较慢"
|
||
|
||
first_token = f"{r['first_token_ms']:.0f}ms" if r["first_token_ms"] else "-"
|
||
question_short = r['question'][:28] + ".." if len(r['question']) > 28 else r['question']
|
||
print(f"#{r['request_number']:<6} {question_short:<30} {r['elapsed_ms']:.0f}ms{'':<5} {first_token:<10} {speed}")
|
||
|
||
print("-" * 80)
|
||
print(f"\n📊 统计信息:")
|
||
print(f" 平均耗时: {avg_time:.0f}ms ({avg_time/1000:.2f}s)")
|
||
print(f" 最快耗时: {min_time:.0f}ms")
|
||
print(f" 最慢耗时: {max_time:.0f}ms")
|
||
|
||
# 缓存效果分析
|
||
if len(valid_results) >= 2:
|
||
first_time = valid_results[0]["elapsed_ms"]
|
||
last_time = valid_results[-1]["elapsed_ms"]
|
||
|
||
print(f"\n💡 缓存效果分析:")
|
||
if last_time < first_time * 0.7:
|
||
print(f" ✅ 缓存可能生效! 后端请求({first_time:.0f}ms)比首次请求快 {((first_time - last_time) / first_time * 100):.1f}%")
|
||
elif last_time > first_time * 1.3:
|
||
print(f" ⚠️ 缓存可能未生效，后续请求比首次请求慢")
|
||
else:
|
||
print(f" 🤔 响应时间相近，缓存效果不确定")
|
||
|
||
# 首Token时间分析
|
||
first_tokens = [r["first_token_ms"] for r in valid_results if r["first_token_ms"] is not None]
|
||
if len(first_tokens) >= 2:
|
||
first_ttft = first_tokens[0]
|
||
last_ttft = first_tokens[-1]
|
||
print(f"\n🚀 首Token时间分析:")
|
||
if last_ttft < first_ttft * 0.5:
|
||
print(f" ✅ 后续请求首Token更快，可能命中缓存")
|
||
elif last_ttft > first_ttft * 1.5:
|
||
print(f" ⚠️ 后续请求首Token更慢，可能是新问题")
|
||
else:
|
||
print(f" 🤔 首Token时间相近")
|
||
|
||
print("📊"*30 + "\n")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
parser = argparse.ArgumentParser(
|
||
description="Call ragflow chat endpoint with SSE streaming and performance testing."
|
||
)
|
||
parser.add_argument(
|
||
"question",
|
||
nargs="?",
|
||
help="要提的问题，未提供则用环境变量 RAGFLOW_QUESTION",
|
||
)
|
||
parser.add_argument(
|
||
"-q",
|
||
"--question",
|
||
dest="question_opt",
|
||
help="与位置参数等价，优先级更高。",
|
||
)
|
||
parser.add_argument(
|
||
"--chat-id",
|
||
default=DEFAULT_CHAT_ID,
|
||
help="chatId，默认取 RAGFLOW_CHAT_ID。",
|
||
)
|
||
parser.add_argument(
|
||
"--session-id",
|
||
default=DEFAULT_SESSION_ID or uuid.uuid4().hex,
|
||
help="sessionId，默认取 RAGFLOW_SESSION_ID，否则随机生成。",
|
||
)
|
||
parser.add_argument(
|
||
"--url",
|
||
default=API_URL,
|
||
help="接口地址，默认取 RAGFLOW_URL。",
|
||
)
|
||
parser.add_argument(
|
||
"-n",
|
||
"--num-requests",
|
||
type=int,
|
||
default=3,
|
||
help="性能测试的请求次数，默认3次。",
|
||
)
|
||
parser.add_argument(
|
||
"-d",
|
||
"--delay",
|
||
type=float,
|
||
default=2.0,
|
||
help="请求之间的间隔时间(秒)，默认2秒。",
|
||
)
|
||
parser.add_argument(
|
||
"--single",
|
||
action="store_true",
|
||
help="单次请求模式(不进行性能测试)。",
|
||
)
|
||
args = parser.parse_args()
|
||
|
||
# Allow either positional or -q/--question to override env default.
|
||
question = args.question_opt or args.question or DEFAULT_QUESTION
|
||
|
||
# Override globals if provided via CLI.
|
||
API_URL = args.url # type: ignore
|
||
|
||
if args.single:
|
||
# 单次请求模式
|
||
asyncio.run(
|
||
stream_chat_async(
|
||
question=question,
|
||
chat_id=args.chat_id,
|
||
session_id=args.session_id,
|
||
)
|
||
)
|
||
else:
|
||
# 性能测试模式
|
||
if question and question != DEFAULT_QUESTION:
|
||
questions = [question, question, question]
|
||
print(f"📝 使用命令行参数的问题进行性能测试: {question}")
|
||
else:
|
||
questions = [
|
||
"公司有什么产品？", # 问题1
|
||
"公司有什么产品？", # 问题2（完全相同，测试精确缓存）
|
||
"公司有哪些产品？", # 问题3（表述略不同，测试语义缓存）
|
||
]
|
||
|
||
asyncio.run(
|
||
run_performance_test(
|
||
questions=questions,
|
||
chat_id=args.chat_id,
|
||
session_id=args.session_id,
|
||
delay_between_requests=args.delay
|
||
)
|
||
)
|