170 lines
5.4 KiB
Python
170 lines
5.4 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
SSE流式响应缓冲测试脚本
|
||
用于验证FastAPI StreamingResponse的缓冲机制，不修改现有代码
|
||
"""
|
||
|
||
import asyncio
|
||
import time
|
||
import json
|
||
from typing import AsyncGenerator
|
||
from fastapi import FastAPI, Response
|
||
from fastapi.responses import StreamingResponse
|
||
import uvicorn
|
||
|
||
|
||
app = FastAPI()
|
||
|
||
|
||
async def slow_data_generator() -> AsyncGenerator[dict, None]:
|
||
"""模拟慢速数据生成，类似于RAGFlow的流式响应"""
|
||
print("Generator started at:", time.time())
|
||
|
||
for i in range(5):
|
||
print(f"Generating chunk {i} at {time.time()}")
|
||
data = {
|
||
"data": {
|
||
"answer": f"这是第{i+1}个数据块，生成时间: {time.time()}",
|
||
"timestamp": time.time()
|
||
}
|
||
}
|
||
yield data
|
||
await asyncio.sleep(1) # 模拟每1秒产生一个数据块
|
||
print(f"Yielded chunk {i} at {time.time()}")
|
||
|
||
|
||
@app.get("/test/streaming")
|
||
async def test_streaming():
|
||
"""测试1: FastAPI StreamingResponse的默认行为"""
|
||
start_time = time.time()
|
||
|
||
async def stream_response():
|
||
print(f"Stream response started at {start_time}")
|
||
|
||
async for chunk in slow_data_generator():
|
||
print(f"About to yield: {chunk}")
|
||
payload = json.dumps(chunk, ensure_ascii=False)
|
||
sse_data = f"data: {payload}\n\n"
|
||
print(f"Yielding SSE data at {time.time()}")
|
||
yield sse_data
|
||
|
||
print(f"Stream response ended at {time.time()}")
|
||
|
||
return StreamingResponse(
|
||
stream_response(),
|
||
media_type='text/event-stream',
|
||
headers={
|
||
'Cache-Control': 'no-cache',
|
||
'Connection': 'keep-alive',
|
||
'X-Accel-Buffering': 'no'
|
||
}
|
||
)
|
||
|
||
|
||
@app.get("/test/plain")
|
||
async def test_plain():
|
||
"""测试2: 对比使用PlainTextResponse的行为"""
|
||
async def generate_text():
|
||
async for chunk in slow_data_generator():
|
||
payload = json.dumps(chunk, ensure_ascii=False)
|
||
sse_data = f"data: {payload}\n\n"
|
||
yield sse_data
|
||
await asyncio.sleep(0.1) # 短暂延迟
|
||
|
||
return Response(generate_text(), media_type='text/plain')
|
||
|
||
|
||
@app.get("/test/flush")
|
||
async def test_flush():
|
||
"""测试3: 尝试强制刷新缓冲区"""
|
||
|
||
async def stream_with_flush():
|
||
async for chunk in slow_data_generator():
|
||
payload = json.dumps(chunk, ensure_ascii=False)
|
||
sse_data = f"data: {payload}\n\n"
|
||
|
||
# 尝试多种方法强制刷新
|
||
print(f"Yielding with flush attempt: {chunk}")
|
||
yield sse_data
|
||
|
||
# 让出控制权
|
||
await asyncio.sleep(0)
|
||
|
||
return StreamingResponse(
|
||
stream_with_flush(),
|
||
media_type='text/event-stream'
|
||
)
|
||
|
||
|
||
# 客户端测试脚本
|
||
async def test_client():
|
||
"""客户端测试函数，用于验证服务端响应行为"""
|
||
|
||
print("=" * 50)
|
||
print("SSE流式响应缓冲测试")
|
||
print("=" * 50)
|
||
|
||
import aiohttp
|
||
|
||
async with aiohttp.ClientSession() as session:
|
||
# 测试1: StreamingResponse
|
||
print("\n📡 测试1: FastAPI StreamingResponse")
|
||
print("请求URL: http://localhost:8000/test/streaming")
|
||
|
||
try:
|
||
async with session.get('http://localhost:8000/test/streaming') as response:
|
||
print(f"响应状态: {response.status}")
|
||
print(f"响应头: {dict(response.headers)}")
|
||
|
||
chunk_count = 0
|
||
async for line in response.content:
|
||
chunk_count += 1
|
||
print(f"客户端收到数据 {chunk_count}: {line.decode().strip()} | 时间: {time.time()}")
|
||
|
||
except Exception as e:
|
||
print(f"测试1失败: {e}")
|
||
|
||
await asyncio.sleep(2)
|
||
|
||
# 测试2: PlainTextResponse
|
||
print("\n📡 测试2: PlainTextResponse")
|
||
print("请求URL: http://localhost:8000/test/plain")
|
||
|
||
try:
|
||
async with session.get('http://localhost:8000/test/plain') as response:
|
||
print(f"响应状态: {response.status}")
|
||
|
||
chunk_count = 0
|
||
async for line in response.content:
|
||
chunk_count += 1
|
||
print(f"客户端收到数据 {chunk_count}: {line.decode().strip()} | 时间: {time.time()}")
|
||
|
||
except Exception as e:
|
||
print(f"测试2失败: {e}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
print("""
|
||
SSE流式响应缓冲测试脚本
|
||
|
||
这个脚本会启动一个测试服务器，模拟RAGFlow的流式响应行为，
|
||
用于验证FastAPI StreamingResponse是否存在缓冲问题。
|
||
|
||
使用方法:
|
||
1. 启动测试服务器: python test_sse_buffer_test.py server
|
||
2. 在另一个终端运行客户端测试: python test_sse_buffer_test.py client
|
||
|
||
或者直接运行主程序进行集成测试
|
||
""")
|
||
|
||
import sys
|
||
if len(sys.argv) > 1:
|
||
if sys.argv[1] == "server":
|
||
print("启动测试服务器...")
|
||
uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info")
|
||
elif sys.argv[1] == "client":
|
||
print("运行客户端测试...")
|
||
asyncio.run(test_client())
|
||
else:
|
||
print("运行集成测试...")
|
||
print("请在另一个终端运行 'python test_sse_buffer_test.py client' 来测试") |