Мост веб-сокетов для аудиопотока в FastAPI

Мост для веб-сокета <-> веб-сокет

Ниже приведен простой пример прокси-сервера веб-сокета, где веб-сокет A и веб-сокет B являются конечными точками в приложении FastAPI, но веб-сокет B может находиться где-то еще, просто измените его адрес 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 B спрятан за синхронным генератором (на самом деле их два, один на вход, другой на выход) и по сути задача превращается в то, как сделать мост между WebSocket и синхронным генератором. А так как я никогда не работал с библиотекой rev-ai, я сделал функцию-заглушку stream_client_start для streamclient.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)