154 lines
4.6 KiB
Python
154 lines
4.6 KiB
Python
#!/usr/bin/env python
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
WebSocket流式输出测试脚本
|
||
"""
|
||
import websocket
|
||
import json
|
||
import sys
|
||
import time
|
||
import uuid
|
||
from datetime import datetime
|
||
import threading
|
||
import ssl
|
||
import argparse
|
||
|
||
# 测试配置
|
||
WS_URL = "ws://127.0.0.1:8000/ws/chat/stream/" # WebSocket URL
|
||
TOKEN = "7831a86588bc08d025e4c9bd668de3b7940f7634" # 替换为你的实际认证令牌
|
||
|
||
# 测试数据
|
||
test_data = {
|
||
"question": "什么是流式输出?",
|
||
"conversation_id": str(uuid.uuid4()), # 随机生成一个会话ID
|
||
"dataset_id_list": ["8390ca43-6e63-4df9-b0b9-6cb20e1b38af"] # 替换为实际的知识库ID
|
||
}
|
||
|
||
# 全局变量
|
||
response_count = 0
|
||
start_time = None
|
||
full_content = ""
|
||
is_connected = False
|
||
|
||
|
||
def on_message(ws, message):
|
||
"""
|
||
处理接收到的WebSocket消息
|
||
"""
|
||
global response_count, full_content
|
||
|
||
try:
|
||
# 解析JSON响应
|
||
response_count += 1
|
||
data = json.loads(message)
|
||
|
||
# 获取当前时间和距开始的时间
|
||
current_time = time.time()
|
||
elapsed = current_time - start_time
|
||
timestamp = datetime.now().strftime('%H:%M:%S.%f')[:-3]
|
||
|
||
# 打印基本信息
|
||
print(f"\n[{timestamp}] 响应 #{response_count} (总计: {elapsed:.3f}s)")
|
||
|
||
# 检查消息类型
|
||
msg_type = data.get('message', '')
|
||
if msg_type == '开始流式传输':
|
||
print("=== 开始接收流式内容 ===")
|
||
elif msg_type == 'partial':
|
||
# 显示部分内容
|
||
if 'data' in data and 'content' in data['data']:
|
||
content = data['data']['content']
|
||
full_content += content
|
||
# 如果内容太长，只显示前50个字符
|
||
display_content = content[:50] + "..." if len(content) > 50 else content
|
||
print(f"部分内容: {display_content}")
|
||
elif msg_type == '完成':
|
||
# 显示完整信息
|
||
if 'data' in data:
|
||
if 'title' in data['data']:
|
||
print(f"标题: {data['data']['title']}")
|
||
if 'content' in data['data']:
|
||
print(f"完整内容长度: {len(data['data']['content'])} 字符")
|
||
print("=== 流式传输完成 ===")
|
||
|
||
# 如果是错误消息
|
||
if data.get('code') == 500:
|
||
print(f"错误: {data.get('message')}")
|
||
ws.close()
|
||
|
||
except json.JSONDecodeError as e:
|
||
print(f"JSON解析错误: {str(e)}")
|
||
except Exception as e:
|
||
print(f"处理消息时出错: {str(e)}")
|
||
|
||
|
||
def on_error(ws, error):
|
||
"""处理WebSocket错误"""
|
||
print(f"发生错误: {str(error)}")
|
||
|
||
|
||
def on_close(ws, close_status_code, close_msg):
|
||
"""处理WebSocket连接关闭"""
|
||
global is_connected
|
||
is_connected = False
|
||
total_time = time.time() - start_time
|
||
print(f"\n===== 连接关闭 =====")
|
||
print(f"状态码: {close_status_code}, 消息: {close_msg}")
|
||
print(f"总响应时间: {total_time:.3f}秒, 共接收 {response_count} 个数据包")
|
||
print(f"接收到的完整内容长度: {len(full_content)} 字符")
|
||
|
||
|
||
def on_open(ws):
|
||
"""处理WebSocket连接成功"""
|
||
global start_time, is_connected
|
||
is_connected = True
|
||
print("WebSocket连接已建立")
|
||
print(f"发送测试数据: {json.dumps(test_data, ensure_ascii=False)}")
|
||
|
||
# 记录开始时间
|
||
start_time = time.time()
|
||
|
||
# 发送测试数据
|
||
ws.send(json.dumps(test_data))
|
||
print("数据已发送，等待响应...")
|
||
|
||
|
||
def main():
|
||
"""主函数"""
|
||
parser = argparse.ArgumentParser(description='WebSocket流式输出测试工具')
|
||
parser.add_argument('--url', type=str, default=WS_URL, help='WebSocket URL')
|
||
parser.add_argument('--token', type=str, default=TOKEN, help='认证令牌')
|
||
parser.add_argument('--question', type=str, default=test_data['question'], help='要发送的问题')
|
||
args = parser.parse_args()
|
||
|
||
# 更新测试数据
|
||
url = f"{args.url}?token={args.token}"
|
||
test_data['question'] = args.question
|
||
|
||
print(f"连接到: {url}")
|
||
|
||
# 设置更详细的日志级别（可选）
|
||
# websocket.enableTrace(True)
|
||
|
||
# 创建WebSocket连接
|
||
ws = websocket.WebSocketApp(
|
||
url,
|
||
on_open=on_open,
|
||
on_message=on_message,
|
||
on_error=on_error,
|
||
on_close=on_close
|
||
)
|
||
|
||
# 设置运行超时（可选）
|
||
# 如果需要SSL连接
|
||
# ws.run_forever(sslopt={"cert_reqs": ssl.CERT_NONE})
|
||
|
||
# 启动WebSocket连接
|
||
ws.run_forever()
|
||
|
||
# 等待一小段时间以确保所有消息都被处理
|
||
time.sleep(1)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |