daren_project/test_stream_api.py

106 lines
3.6 KiB
Python
Raw Normal View History

2025-04-29 10:22:57 +08:00
import requests
import json
import sys
import time
from datetime import datetime
# 接口地址 - 根据curl命令更新
API_URL = "http://127.0.0.1:8000/api/chat-history/" # 注意没有create路径
# 测试数据
test_data = {
"question": "总结下",
"conversation_id": "10b34248-2625-434b-a493-6d43520c837a",
"dataset_id_list": ["8390ca43-6e63-4df9-b0b9-6cb20e1b38af"],
"stream": True
}
# 请求头 - 添加认证令牌和其他头信息
headers = {
"Content-Type": "application/json",
"Authorization": "Token 7831a86588bc08d025e4c9bd668de3b7940f7634",
"User-Agent": "Apifox/1.0.0 (https://apifox.com)",
"Accept": "*/*",
"Host": "127.0.0.1:8000",
"Connection": "keep-alive"
}
print(f"API URL: {API_URL}")
print(f"发送请求: {json.dumps(test_data, ensure_ascii=False)}")
print("等待响应...")
start_time = time.time()
# 发送请求并获取流式响应
response = requests.post(
url=API_URL,
json=test_data,
headers=headers,
stream=True # 启用流式传输
)
print(f"响应状态码: {response.status_code}")
print(f"收到初始响应时间: {datetime.now().strftime('%H:%M:%S.%f')[:-3]}")
if response.status_code != 200 and response.status_code != 201:
print(f"错误: {response.status_code}, {response.text}")
sys.exit(1)
# 处理流式响应
print("\n----- 开始接收流式响应 -----\n")
buffer = ""
last_time = start_time
response_count = 0
for chunk in response.iter_content(chunk_size=1024):
if chunk:
current_time = time.time()
time_diff = current_time - last_time
last_time = current_time
elapsed = current_time - start_time
# 解码字节为字符串
chunk_str = chunk.decode('utf-8')
buffer += chunk_str
# 检查是否有完整的数据行
if '\n\n' in buffer:
lines = buffer.split('\n\n')
# 除了最后一行,其他都是完整的
for line in lines[:-1]:
if line.strip():
response_count += 1
timestamp = datetime.now().strftime('%H:%M:%S.%f')[:-3]
print(f"\n[{timestamp}] 响应 #{response_count} (距上次: {time_diff:.3f}s, 总计: {elapsed:.3f}s)")
print(f"{line}")
# 如果想解析JSON内容可以取消下面的注释
if line.startswith('data: '):
try:
json_str = line[6:] # 去掉 "data: " 前缀
data = json.loads(json_str)
if data.get('content'):
content = data.get('content')
# 如果内容太长只显示前30个字符
if len(content) > 30:
content = content[:30] + "..."
print(f"内容片段: {content}")
except json.JSONDecodeError:
pass
# 保留最后一个可能不完整的行
buffer = lines[-1]
# 重置计时器
last_time = time.time()
# 处理可能的剩余数据
if buffer.strip():
timestamp = datetime.now().strftime('%H:%M:%S.%f')[:-3]
elapsed = time.time() - start_time
print(f"\n[{timestamp}] 最终响应 (总计: {elapsed:.3f}s)")
print(f"{buffer}")
total_time = time.time() - start_time
print(f"\n----- 响应结束 -----")
print(f"总响应时间: {total_time:.3f}秒, 共接收 {response_count} 个数据包")