daren_project/test_stream_api.py
2025-04-29 10:22:57 +08:00

106 lines
3.6 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

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 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} 个数据包")