大模型流式输出简谈(websocket版)

百年੭ ᐕ)੭*⁾⁾2026-06-20 8:39

大模型流式输出是绕不开的一环,本文我将简单写一个示例,带你了解并简单上手

python代码准备

  • 本次需要用到的包

要用到StreamingResponse来处理流式输出

python 复制代码 
from fastapi import FastAPI,WebSocket
from pydantic import BaseModel
from fastapi.middleware.cors import CORSMiddleware
from langchain.chat_models import init_chat_model
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from dotenv import load_dotenv
import uvicorn
import os

from starlette.websockets import WebSocketDisconnect

  • 流式接口编写

    • 构建自定义请求体ChatRequest,接收用户输入的字符串

    • 建立websocket握手

    • 接收前端传入的文本信息

    • 大模型对传入的文本信息进行流式响应并发送给客户端

    • 设置结束标记

python 复制代码 
class ChatRequest(BaseModel):
    user_input: str
@app.websocket('/chat/ws')
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept() #important:建立握手请求
    try:
        while True:
            #接受客户端消息
            data = await websocket.receive_text() #tips:异步等待客户端发送文本消息

            #解析json
            try:
                request=ChatRequest.model_validate_json(data) #tips:验证传入的data是否符合定义的Pydantic请求模型ChatRequest
            except Exception as e:
                await websocket.send_text(f'[ERROR] Invalid JSON {str(e)}')
                continue

            async for chunk in chain.astream({
                'user_input': request.user_input
            }):
                if chunk:
                    await websocket.send_text(chunk)

            #发送结束标记
            await websocket.send_text('[DONE]')

    except WebSocketDisconnect:
        print('Client disconnected')

    except Exception as e:
        print(f"Error: {e}")
        await websocket.send_text(f"[ERROR] {str(e)}")

前端代码准备

  • 主要讲解一下javascript中流式输出方法的实现
    • 构建connect函数并新建WebSocket对象ws与fastapi后端接口对接，管理websocket的开启，接收信息，关闭以及异常处理
    • 设置事件监听对象.onmessage判断是否接收到终止符，接收到就给反馈显示回答结束，否则就在输出栏添加event.data
    • 构建sendMessage函数实现用户消息输入与响应
html 复制代码 
<script>
let ws;

function connect() {
    ws = new WebSocket("ws://127.0.0.1:8000/chat/ws");

    ws.onopen = () => {
        console.log("WebSocket connected");
    };

    ws.onmessage = (event) => {
        const output = document.getElementById('output');
        const text = event.data;
        if (text === "[DONE]") {
            output.innerHTML += "<br><b>--- 回答结束 ---</b><br>";
        } else {
            output.textContent += text; // 或用 innerHTML += text.replace(/\n/g, '<br>')
        }
        output.scrollTop = output.scrollHeight;
    };

    ws.onclose = () => {
        console.log("WebSocket disconnected");
    };

    ws.onerror = (err) => {
        console.error("WebSocket error:", err);
    };
}

function sendMessage() {
    const input = document.getElementById('user_input');
    const msg = input.value.trim();
    if (!msg) return;

    if (ws?.readyState === WebSocket.OPEN) {
        ws.send(JSON.stringify({ user_input: msg }));
        document.getElementById('output').textContent += `\n\n你: ${msg}\nAI: `;
        input.value = '';
    } else {
        alert("尚未连接到服务器！");
    }
}

// 页面加载时自动连接
window.onload = connect;
</script>

整体代码

python

python 复制代码 
from fastapi import FastAPI,WebSocket
from pydantic import BaseModel
from fastapi.middleware.cors import CORSMiddleware
from langchain.chat_models import init_chat_model
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from dotenv import load_dotenv
import uvicorn
import os

from starlette.websockets import WebSocketDisconnect

load_dotenv()

llm=init_chat_model(
    model='glm-4.7',
    model_provider='openai',
    api_key=os.getenv('zhipu_key'),
    base_url=os.getenv('zhipu_base_url')
)

app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=['*'],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"]
)

chat_prompt_template = ChatPromptTemplate(
    messages=[
        ('system','你现在是一位小说家'),
        ('human','{user_input}')
    ]
)

chain = chat_prompt_template|llm|StrOutputParser()

# 定义请求体模型（用于验证）
class ChatRequest(BaseModel):
    user_input: str


@app.websocket('/chat/ws')
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept() #important:建立握手请求
    try:
        while True:
            #接受客户端消息
            data = await websocket.receive_text() #tips:异步等待客户端发送文本消息

            #解析json
            try:
                request=ChatRequest.model_validate_json(data) #tips:验证传入的data是否符合定义的Pydantic请求模型ChatRequest
            except Exception as e:
                await websocket.send_text(f'[ERROR] Invalid JSON {str(e)}')
                continue

            async for chunk in chain.astream({
                'user_input': request.user_input
            }):
                if chunk:
                    await websocket.send_text(chunk)

            #发送结束标记
            await websocket.send_text('[DONE]')

    except WebSocketDisconnect:
        print('Client disconnected')

    except Exception as e:
        print(f"Error: {e}")
        await websocket.send_text(f"[ERROR] {str(e)}")

if __name__ == '__main__':
    uvicorn.run(app, host='127.0.0.1', port=8000)

前端

html 复制代码 
<!DOCTYPE html>
<html>
<head>
    <meta charset="UTF-8">
    <title>WebSocket Chat</title>
</head>
<body>
<div id="output" style="width:500px;height:300px;background:#f0f0f0;overflow:auto;padding:10px;"></div>
<input id="user_input" type="text" placeholder="输入你的问题..." style="width:400px;">
<button onclick="sendMessage()">发送</button>

<script>
let ws;

function connect() {
    ws = new WebSocket("ws://127.0.0.1:8000/chat/ws");

    ws.onopen = () => {
        console.log("WebSocket connected");
    };

    ws.onmessage = (event) => {
        const output = document.getElementById('output');
        const text = event.data;
        if (text === "[DONE]") {
            output.innerHTML += "<br><b>--- 回答结束 ---</b><br>";
        } else {
            output.textContent += text; // 或用 innerHTML += text.replace(/\n/g, '<br>')
        }
        output.scrollTop = output.scrollHeight;
    };

    ws.onclose = () => {
        console.log("WebSocket disconnected");
    };

    ws.onerror = (err) => {
        console.error("WebSocket error:", err);
    };
}

function sendMessage() {
    const input = document.getElementById('user_input');
    const msg = input.value.trim();
    if (!msg) return;

    if (ws?.readyState === WebSocket.OPEN) {
        ws.send(JSON.stringify({ user_input: msg }));
        document.getElementById('output').textContent += `\n\n你: ${msg}\nAI: `;
        input.value = '';
    } else {
        alert("尚未连接到服务器！");
    }
}

// 页面加载时自动连接
window.onload = connect;
</script>
</body>
</html>

websocket实现流式输出相对于http/https实现的方式还是不太一样的,希望本篇能够对你有所启发。

热门推荐
012026年7月AI圈大地震：GPT-5.6被政府限制、Claude入驻Slack、Anthropic自研芯片02GitHub 镜像站点03如何新建文件夹？ 电脑新建文件夹的4种方法04幻兽帕鲁 - 服务器管理员权限与 GM 命令完全指南05AI科技热点日报 | 2026年07月01日06国内可直接用、免费额度/永久免费的大模型API清单（含 SiliconFlow、火山、阿里、智谱、百度、Kimi、DeepSeek、DMXAPI 等）07【AI前线观察】2026年国产开源大模型全面横评：从 DeepSeek V4 到 Kimi K3，谁才是开发者的最优选择？082026年AI技术突破与产业落地全景：从GPT-5到多模态智能体的新纪元092026 国产 AI 大模型横评：DeepSeek、通义千问、Kimi、文心一言、星火、豆包谁更能打？10全面体验 Grok API 中转站（2025 · Grok 4 系列最新版）