LLM 流式输出实现：SSE 与 WebSocket

from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
import httpx, json, asyncio

app = FastAPI()
LLM_API_URL = "https://api.openai.com/v1/chat/completions"

@app.post("/api/chat/stream")
async def chat_stream(request: Request):
    body = await request.json()
    
    async def event_generator():
        headers = {
            "Authorization": f"Bearer {body['api_key']}",
            "Content-Type": "application/json"
        }
        payload = {
            "model": body["model"],
            "messages": body["messages"],
            "stream": True,
            "max_tokens": 4096
        }
        
        try:
            async with httpx.AsyncClient(timeout=120) as client:
                async with client.stream("POST", LLM_API_URL,
                                          headers=headers, json=payload) as resp:
                    async for line in resp.aiter_lines():
                        if not line.startswith("data: "):
                            continue
                        data = line[6:]
                        if data.strip() == "[DONE]":
                            yield "data: " + json.dumps({"type": "done"}) + "\n\n"
                            break
                        chunk = json.loads(data)
                        delta = chunk["choices"][0]["delta"]
                        if "content" in delta:
                            token = delta["content"]
                            yield "data: " + json.dumps({"type": "token", "content": token}) + "\n\n"
        except httpx.ReadTimeout:
            yield "data: " + json.dumps({"type": "error", "msg": "LLM timeout"}) + "\n\n"
        except Exception as e:
            yield "data: " + json.dumps({"type": "error", "msg": str(e)}) + "\n\n"
    
    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no", "Connection": "keep-alive"}
    )

class SSEClient {
  constructor(url) { this.url = url; this.controller = null; }
  
  async stream(messages, onToken, onDone, onError) {
    this.controller = new AbortController();
    try {
      const resp = await fetch(this.url, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({ messages }),
        signal: this.controller.signal
      });
      const reader = resp.body.getReader();
      const decoder = new TextDecoder();
      let buffer = '';
      while (true) {
        const { done, value } = await reader.read();
        if (done) break;
        buffer += decoder.decode(value, { stream: true });
        const lines = buffer.split('\n');
        buffer = lines.pop();
        for (const line of lines) {
          if (!line.startsWith('data: ')) continue;
          const data = JSON.parse(line.slice(6));
          if (data.type === 'token') onToken(data.content);
          else if (data.type === 'done') onDone();
          else if (data.type === 'error') onError(data.msg);
        }
      }
    } catch (e) {
      if (e.name !== 'AbortError') onError(e.message);
    }
  }
  abort() { this.controller?.abort(); }
}

// 使用
const client = new SSEClient('/api/chat/stream');
let fullText = '';
await client.stream(
  [{ role: 'user', content: '解释 RAG 架构' }],
  (token) => { fullText += token; renderMarkdown(fullText); },
  () => console.log('完成'),
  (err) => console.error('错误:', err)
);

from fastapi import WebSocket, WebSocketDisconnect

@app.websocket("/ws/chat")
async def ws_chat(ws: WebSocket):
    await ws.accept()
    try:
        while True:
            msg = await ws.receive_json()
            if msg.get("action") == "abort":
                await ws.send_json({"type": "aborted"})
                continue
            async for token in call_llm_stream(msg["messages"]):
                await ws.send_json({"type": "token", "content": token})
            await ws.send_json({"type": "done"})
    except WebSocketDisconnect:
        print("Client disconnected")

class ReconnectingSSE {
  constructor(url, opts = {}) {
    this.url = url;
    this.maxRetries = opts.maxRetries || 3;
    this.retryDelay = opts.retryDelay || 1000;
    this.retries = 0;
  }
  async connect(messages) {
    while (this.retries < this.maxRetries) {
      try {
        await this.sseStream(messages);
        this.retries = 0;
        return;
      } catch (e) {
        this.retries++;
        const delay = this.retryDelay * Math.pow(2, this.retries);
        await new Promise(r => setTimeout(r, delay));
      }
    }
    throw new Error('Max retries exceeded');
  }
}