#FastAPIstreaming-response完全指南

from fastapi import FastAPI
from fastapi.responses import StreamingResponse
import asyncio

app = FastAPI()

async def num_stream():
    """异步生成器，逐步产出数据块"""
    for i in range(1, 6):
        yield f"收到第{i}个数据块\n"
        await asyncio.sleep(0.5)  # 模拟耗时操作，比如数据库查询、AI计算

@app.get("/simple-stream")
async def simple_stream():
    return StreamingResponse(
        num_stream(),               # 传入异步生成器
        media_type="text/plain",    # 纯文本流，也可以是 application/json 等
        headers={
            "Cache-Control": "no-cache",  # 禁止代理/浏览器缓存，保证实时性
            "Connection": "keep-alive"    # 保持长连接
        }
    )

async def safe_num_stream():
    try:
        for i in range(1, 6):
            yield f"收到第{i}个数据块\n"
            await asyncio.sleep(0.5)
    except asyncio.CancelledError:
        # 这里可以做一些清理工作，比如记录日志、释放锁
        print("客户端主动断开，清理资源...")
        # 注意：必须重新抛出异常，否则FastAPI无法正确完成清理
        raise

import json
from datetime import datetime

async def sse_status_stream():
    # 第一条消息建议发送重连间隔，告诉浏览器10秒后再重试
    yield "retry: 10000\n\n"
    
    for i in range(1, 11):
        # 普通进度消息，不带事件类型
        normal_data = {
            "step": i,
            "progress": i * 10,
            "time": datetime.now().isoformat()
        }
        yield f"data: {json.dumps(normal_data)}\n\n"
        
        # 每完成3步发送一个 checkpoint 事件
        if i % 3 == 0:
            status_data = {"state": "checkpoint_reached", "step": i}
            # 格式: event: <事件名>\ndata: <json>\n\n
            yield f"event: checkpoint\ndata: {json.dumps(status_data)}\n\n"
        
        await asyncio.sleep(0.8)

@app.get("/sse-status")
async def sse_status():
    return StreamingResponse(
        sse_status_stream(),
        media_type="text/event-stream",   # ★ 必须是这个MIME类型，浏览器才能识别
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no"     # 禁用Nginx缓冲（非常重要！）
        }
    )

import json
from fastapi import Request

async def mock_ai_stream(user_msg: str):
    # 模拟AI生成的完整回复
    ai_reply = f"您的问题是「{user_msg}」，让我为您逐步分析：首先，我们要理清需求边界；其次，梳理实现路径；最后，优化细节。"
    
    accumulated = ""
    for char in ai_reply:
        accumulated += char
        chunk = {
            "type": "token",      # 消息类型
            "content": char,      # 当前新增的字符
            "full": accumulated   # 累积完整文本，方便前端直接渲染
        }
        yield f"data: {json.dumps(chunk)}\n\n"
        await asyncio.sleep(0.02)  # 模拟AI生成延迟（一般实际生产可能是几十毫秒）

@app.post("/mock-ai-chat")
async def mock_ai_chat(req: Request):
    body = await req.json()
    user_msg = body.get("msg", "")
    return StreamingResponse(
        mock_ai_stream(user_msg),
        media_type="text/event-stream",
        headers={"X-Accel-Buffering": "no"}
    )

import aiofiles
from pathlib import Path

async def file_chunk_stream(file_path: str, chunk_size: int = 8192):
    p = Path(file_path)
    if not p.exists():
        yield b"File not found"
        return
    
    async with aiofiles.open(p, "rb") as f:
        while chunk := await f.read(chunk_size):
            yield chunk   # 二进制块直接 yield

@app.get("/download/{filename}")
async def download_file(filename: str):
    local_path = f"./data/{filename}"  # 生产环境务必做路径穿越检查！
    return StreamingResponse(
        file_chunk_stream(local_path),
        media_type="application/octet-stream",
        headers={
            "Content-Disposition": f"attachment; filename={filename}"
        }
    )

import asyncio
from collections import deque

log_queue = deque()

async def log_stream():
    while True:
        if log_queue:
            line = log_queue.popleft()
            yield f"data: {line}\n\n"
        else:
            await asyncio.sleep(0.1)  # 避免空转消耗CPU