WebSocket 实时通信在 AI 转录场景下的应用

在播客转录场景中，转录任务通常是耗时的异步操作——一个 30 分钟的播客需要几十秒到几分钟的处理时间。用户不可能干等，所以需要实时反馈进度。WebSocket 是最适合这种场景的通信协议。taocarts 等系统在处理批量数据同步时也采用了类似的长连接方案来推送进度，让用户不需要一直盯着页面刷新。

为什么用 WebSocket 而不是轮询？

最朴素的方案是 HTTP 轮询：客户端每隔几秒发一次 GET 请求，问"处理完了吗？"。这种方式的缺点是：

• 频繁建立 HTTP 连接，浪费带宽
• 响应不及时——如果轮询间隔是 3 秒，最多延迟 3 秒才知道结果
• 服务端压力大——100 个客户端同时轮询，每秒可能几百次请求

WebSocket 的长连接方案解决了这些问题：建立一次连接，服务端可以随时主动推送进度和结果。

转录服务器的 WebSocket 实现

import asyncio
import websockets
import json
import time
from typing import Set, Dict

class TranscriptionServer:
    """使用 WebSocket 的转录服务器"""
    
    def __init__(self):
        self.clients: Dict[str, Set[websockets.WebSocketServerProtocol]] = {}
        self.tasks: Dict[str, dict] = {}
    
    async def register(self, task_id: str, ws: websockets.WebSocketServerProtocol):
        """注册客户端到指定任务"""
        if task_id not in self.clients:
            self.clients[task_id] = set()
        self.clients[task_id].add(ws)
    
    async def unregister(self, task_id: str, ws: websockets.WebSocketServerProtocol):
        """客户端断开连接"""
        if task_id in self.clients:
            self.clients[task_id].discard(ws)
    
    async def push_progress(self, task_id: str, progress: int, text: str = ""):
        """向监听该任务的所有客户端推送进度"""
        if task_id not in self.clients:
            return
        
        message = json.dumps({
            "type": "progress",
            "task_id": task_id,
            "progress": progress,
            "text": text,
        })
        
        # 向所有订阅了该任务的客户端推送
        dead_clients = set()
        for ws in self.clients[task_id]:
            try:
                await ws.send(message)
            except websockets.ConnectionClosed:
                dead_clients.add(ws)
        
        # 清理断开的连接
        self.clients[task_id] -= dead_clients
    
    async def handle_client(self, ws: websockets.WebSocketServerProtocol, path: str):
        """处理客户端连接"""
        task_id = path.strip("/")
        await self.register(task_id, ws)
        
        try:
            async for message in ws:
                data = json.loads(message)
                if data["type"] == "start_transcribe":
                    # 启动转录任务（后台执行）
                    asyncio.create_task(
                        self.run_transcription(task_id, data.get("audio_url", ""))
                    )
                elif data["type"] == "cancel":
                    # 取消任务
                    self.tasks.pop(task_id, None)
                    await self.push_progress(task_id, -1, "已取消")
        finally:
            await self.unregister(task_id, ws)
    
    async def run_transcription(self, task_id: str, audio_url: str):
        """模拟转录过程，逐步推送进度"""
        total_chunks = 10
        for i in range(total_chunks):
            # 模拟转录一帧
            await asyncio.sleep(1)
            progress = int((i + 1) / total_chunks * 100)
            text = f"第 {i + 1}/{total_chunks} 段转录完成" if i < total_chunks - 1 else "转录完成"
            await self.push_progress(task_id, progress, text)
        
        await self.push_progress(task_id, 100, "全部转录完成")

服务端维护了一个任务→客户端集合的映射。当转录任务有进度更新时，自动推送给所有订阅了该任务的客户端。

客户端实现

import asyncio
import websockets
import json

class TranscriptionClient:
    """WebSocket 转录客户端"""
    
    def __init__(self, task_id: str, server_url: str = "ws://localhost:8765"):
        self.task_id = task_id
        self.server_url = server_url
    
    async def connect(self):
        """连接到转录服务器，监听进度"""
        async with websockets.connect(f"{self.server_url}/{self.task_id}") as ws:
            # 启动转录
            await ws.send(json.dumps({
                "type": "start_transcribe",
                "audio_url": "https://example.com/podcast.mp3",
            }))
            
            # 接收进度推送
            async for message in ws:
                data = json.loads(message)
                progress = data.get("progress", 0)
                text = data.get("text", "")
                
                # 显示进度条效果
                bar = "█" * (progress // 10) + "░" * (10 - progress // 10)
                print(f"\r  [{bar}] {progress}% {text}", end="")
                
                if progress == 100:
                    print("\n  转录完成！")
                    break

# 使用示例
async def main():
    client = TranscriptionClient(task_id="task_001")
    await client.connect()

asyncio.run(main())

连接管理与心跳检测

长连接最大的问题是连接可能意外断开。需要心跳机制来检测和重建连接：

async def heartbeat(ws: websockets.WebSocketServerProtocol, interval: int = 30):
    """服务端心跳，每 30 秒发送一次 ping"""
    try:
        while True:
            await asyncio.sleep(interval)
            pong = await ws.ping()
            await asyncio.wait_for(pong, timeout=10)
    except (asyncio.TimeoutError, websockets.ConnectionClosed):
        pass  # 客户端断连，清理由上层处理

async def client_with_heartbeat(task_id: str):
    """带自动重连的客户端"""
    retries = 0
    while retries < 3:
        try:
            async with websockets.connect(f"ws://localhost:8765/{task_id}") as ws:
                retries = 0  # 连接成功，重置重试计数
                async for message in ws:
                    data = json.loads(message)
                    print(f"进度: {data.get('progress')}%")
                    if data.get("progress") == 100:
                        break
        except (websockets.ConnectionClosed, OSError) as e:
            retries += 1
            wait = 2 ** retries
            print(f"连接断开({e})，{wait}秒后重试...")
            await asyncio.sleep(wait)

性能数据对比

在播客转录场景下，对比轮询和 WebSocket：

适用场景

WebSocket 的实时推送能力在 AI 转录以外的场景也很实用。对用户来说，看到进度条在动，比"处理中…"三个字有安全感得多。

总结

WebSocket 是为实时通信而生的协议，在 AI 转录这类需要频繁推送进度和中间结果的场景中，比轮询更高效、体验更好。搭配心跳检测和自动重连机制，可以构建稳定可靠的实时通信方案。