流式响应实现详解

from fastapi import FastAPI
from fastapi.responses import StreamingResponse
import json

app = FastAPI()

@app.post("/chat/stream")
async def chat_stream(request: dict):
    async def event_stream():
        # 调用LLM的流式接口
        async for chunk in llm.astream(
            messages=[{"role": "user", "content": request["message"]}]
        ):
            data = json.dumps({
                "content": chunk.content,
                "role": "assistant"
            }, ensure_ascii=False)
            yield f"data: {data}\n\n"
        
        # 发送结束标记
        yield f"data: [DONE]\n\n"
    
    return StreamingResponse(
        event_stream(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",  # Nginx不缓冲
        }
    )

// 浏览器端SSE
const eventSource = new EventSource('/chat/stream');

eventSource.onmessage = (event) => {
    if (event.data === '[DONE]') {
        eventSource.close();
        return;
    }
    const data = JSON.parse(event.data);
    document.getElementById('output').innerHTML += data.content;
};

eventSource.onerror = (error) => {
    console.error('SSE error:', error);
    eventSource.close();
};

import httpx
import json

async def stream_chat(url, message):
    async with httpx.AsyncClient() as client:
        async with client.stream(
            "POST",
            f"{url}/chat/stream",
            json={"message": message},
            timeout=60.0
        ) as response:
            async for line in response.aiter_lines():
                if line.startswith("data: "):
                    data = line[6:]
                    if data == "[DONE]":
                        break
                    chunk = json.loads(data)
                    print(chunk["content"], end="", flush=True)

from fastapi import FastAPI, WebSocket

app = FastAPI()

@app.websocket("/ws/chat")
async def websocket_chat(ws: WebSocket):
    await ws.accept()
    
    try:
        while True:
            message = await ws.receive_text()
            
            # 流式生成
            async for chunk in llm.astream(
                messages=[{"role": "user", "content": message}]
            ):
                await ws.send_json({
                    "type": "token",
                    "content": chunk.content
                })
            
            await ws.send_json({"type": "done"})
    
    except WebSocketDisconnect:
        print("Client disconnected")

const ws = new WebSocket('ws://localhost:8000/ws/chat');

ws.onopen = () => {
    ws.send(JSON.stringify({message: '你好'}));
};

ws.onmessage = (event) => {
    const data = JSON.parse(event.data);
    if (data.type === 'token') {
        output.innerHTML += data.content;
    } else if (data.type === 'done') {
        console.log('Generation complete');
    }
};

async def typewriter_stream(text, delay=0.03):
    """模拟打字机效果的流式输出"""
    for char in text:
        yield char
        await asyncio.sleep(delay)

class CancellableStream:
    def __init__(self):
        self.cancelled = False
    
    def cancel(self):
        self.cancelled = True
    
    async def stream(self, llm, messages):
        async for chunk in llm.astream(messages):
            if self.cancelled:
                break
            yield chunk

async def parallel_stream(prompts, llm):
    """同时流式生成多个响应"""
    async def stream_one(prompt, index):
        result = []
        async for chunk in llm.astream([{"role": "user", "content": prompt}]):
            result.append({"index": index, "content": chunk.content})
        return result
    
    tasks = [stream_one(p, i) for i, p in enumerate(prompts)]
    results = await asyncio.gather(*tasks)
    
    # 交错输出
    for i in range(max(len(r) for r in results)):
        for result in results:
            if i < len(result):
                yield result[i]

# Nginx配置：禁用缓冲
location /chat/stream {
    proxy_pass http://backend;
    proxy_buffering off;        # 关闭代理缓冲
    proxy_cache off;            # 关闭缓存
    chunked_transfer_encoding on;
    proxy_read_timeout 300s;
}

# Token级流式（推荐）：每个token一个chunk
async for token in llm.astream(messages):
    yield token

# 字符级流式：将token拆分为字符
async for token in llm.astream(messages):
    for char in token.content:
        yield char
        await asyncio.sleep(0.01)  # 添加微小延迟