FastAPI中用於音頻stream的Websockets橋接器

Question

客觀的

我的目標是使用音頻 stream。 從邏輯上講，這是我的目標：

1. 音頻 stream 來自 WebSocket A（ FastAPI端點）
2. 音頻 stream 橋接到不同的 WebSocket，B，這將返回 JSON（ Rev-ai的 WebSocket）
3. Json 結果通過 WebSocket A 實時發回。 因此，雖然音頻 stream 仍在進來。

可能的解決方案

為了解決這個問題，我有很多想法，但最終我一直在嘗試將WebSocket A橋接到WebSocket B 。 到目前為止，我的嘗試涉及一個ConnectionManager class，其中包含一個Queue.queue 。 音頻 stream 的塊被添加到這個隊列中，這樣我們就不會直接從WebSocket A消費。

ConnectionManager還包含一個生成器方法，用於從隊列中產生所有值。

我的 FastAPI 實現從websocket A中消耗，如下所示：

@app.websocket("/ws")
async def predict_feature(websocket: WebSocket):
    await manager.connect(websocket)
    try:
        while True:
            chunk = await websocket.receive_bytes()
            manager.add_to_buffer(chunk)
    except KeyboardInterrupt:
        manager.disconnect()

在此攝取的同時，我想要一項任務，將我們的音頻 ZF7B44CFFAFD5C52223D5498196C8A2E7BZ 連接到WebSocket B ，並將獲得的值發送到WebSocket A 。 音頻 stream 可以通過上述generator方法消耗。

由於 WebSocket B 如何使用消息，因此生成器方法是必要的，如 Rev-ai 的示例所示：

streamclient = RevAiStreamingClient(access_token, config)
response_generator = streamclient.start(MEDIA_GENERATOR)
for response in response_generator:
    # return through websocket A this value
    print(response)

這是最大的挑戰之一，因為我們需要將數據消耗到生成器中並實時獲取結果。

最新嘗試

我一直在用asyncio運氣； 據我了解，一種可能性是創建一個在后台運行的協程。 我在這方面沒有成功，但聽起來很有希望。

我曾考慮通過FastAPI啟動事件觸發此事件，但我無法實現並發。 我嘗試使用event_loops ，但它給了我一個與nested event loop相關的錯誤。

警告

FastAPI can be optional if your insight deems so, and in a way so is WebSocket A. At the end of the day, the ultimate objective is to receive an audio stream through our own API endpoint, run it through Rev.ai's WebSocket, do一些額外的處理，並將結果發回。

Answer 1

websocket <-> websocket 的橋接器

Below is a simple example of websocket proxy, where websocket A and websocket B are both endpoints in the FastAPI app, but websocket B can be located in something else, just change its address ws_b_uri . 對於 websocket 客戶端，使用websockets庫。

為了進行數據轉發， A端點的代碼啟動了兩個任務forward和reverse ，並通過asyncio.gather()等待它們完成。 兩個方向的數據傳輸以並行方式發生。

import asyncio

from fastapi import FastAPI
from fastapi import WebSocket
import websockets
app = FastAPI()

ws_b_uri = "ws://localhost:8001/ws_b"


async def forward(ws_a: WebSocket, ws_b: websockets.WebSocketClientProtocol):
    while True:
        data = await ws_a.receive_bytes()
        print("websocket A received:", data)
        await ws_b.send(data)


async def reverse(ws_a: WebSocket, ws_b: websockets.WebSocketClientProtocol):
    while True:
        data = await ws_b.recv()
        await ws_a.send_text(data)
        print("websocket A sent:", data)


@app.websocket("/ws_a")
async def websocket_a(ws_a: WebSocket):
    await ws_a.accept()
    async with websockets.connect(ws_b_uri) as ws_b_client:
        fwd_task = asyncio.create_task(forward(ws_a, ws_b_client))
        rev_task = asyncio.create_task(reverse(ws_a, ws_b_client))
        await asyncio.gather(fwd_task, rev_task)


@app.websocket("/ws_b")
async def websocket_b(ws_b_server: WebSocket):
    await ws_b_server.accept()
    while True:
        data = await ws_b_server.receive_bytes()
        print("websocket B server recieved: ", data)
        await ws_b_server.send_text('{"response": "value from B server"}')

更新（橋 websocket <-> 同步發生器）

考慮到問題的最后更新，問題是 WebSocket B隱藏在同步發生器后面（實際上有兩個，一個用於輸入，另一個用於輸出），實際上，任務變成了如何在 WebSocket 和同步發電機之間架起一座橋梁。 由於我從未使用過rev-ai庫，因此我為streamclient.start制作了一個存根 function stream_client_start ，它采用生成器（原始MEDIA_GENERATOR ）並返回一個生成器（原始response_generator生成器）。

在這種情況下，我通過run_in_executor在單獨的線程中開始處理生成器，並且為了不重新發明輪子，我使用janus庫中的隊列進行通信，它允許您通過隊列綁定同步和異步代碼. 因此，也有兩個隊列，一個用於A -> B ，另一個用於B -> A 。

import asyncio
import time
from typing import Generator
from fastapi import FastAPI
from fastapi import WebSocket
import janus
import queue

app = FastAPI()


# Stub generator function (using websocket B in internal)
def stream_client_start(input_gen: Generator) -> Generator:
    for chunk in input_gen:
        time.sleep(1)
        yield f"Get {chunk}"


# queue to generator auxiliary adapter
def queue_to_generator(sync_queue: queue.Queue) -> Generator:
    while True:
        yield sync_queue.get()


async def forward(ws_a: WebSocket, queue_b):
    while True:
        data = await ws_a.receive_bytes()
        print("websocket A received:", data)
        await queue_b.put(data)


async def reverse(ws_a: WebSocket, queue_b):
    while True:
        data = await queue_b.get()
        await ws_a.send_text(data)
        print("websocket A sent:", data)


def process_b_client(fwd_queue, rev_queue):
    response_generator = stream_client_start(queue_to_generator(fwd_queue))
    for r in response_generator:
        rev_queue.put(r)


@app.websocket("/ws_a")
async def websocket_a(ws_a: WebSocket):
    loop = asyncio.get_event_loop()
    fwd_queue = janus.Queue()
    rev_queue = janus.Queue()
    await ws_a.accept()

    process_client_task = loop.run_in_executor(None, process_b_client, fwd_queue.sync_q, rev_queue.sync_q)
    fwd_task = asyncio.create_task(forward(ws_a, fwd_queue.async_q))
    rev_task = asyncio.create_task(reverse(ws_a, rev_queue.async_q))
    await asyncio.gather(process_client_task, fwd_task, rev_task)