Function stream_websocket_responses

examples/server/server.py:15911–15939

(
        websocket: WebSocket,
        formatter: OpenAIFormatter,
        stream: CompletionStream,
        _: Callable[[], None],
        *,
        initial_payloads: Iterable[BaseModel | Dict[str, Any]],
        chunk_payloads: Callable[[CompletionChunk], Iterable[Any]],
        done_payloads: Callable[
            [Optional[OpenAICompletion]], Iterable[BaseModel | Dict[str, Any]]
        ],
    )

            disconnect_task.cancel()

    async def stream_websocket_responses(
        websocket: WebSocket,
        formatter: OpenAIFormatter,
        stream: CompletionStream,
        _: Callable[[], None],
        *,
        initial_payloads: Iterable[BaseModel | Dict[str, Any]],
        chunk_payloads: Callable[[CompletionChunk], Iterable[Any]],
        done_payloads: Callable[
            [Optional[OpenAICompletion]], Iterable[BaseModel | Dict[str, Any]]
        ],
    ) -> None:
        try:
            for payload in initial_payloads:
                await websocket.send_json(payload)
            while True:
                done, chunk, result = await asyncio.to_thread(
                    formatter.next_stream_output,
                    stream,
                )
                if done:
                    for payload in done_payloads(result):
                        await websocket.send_json(payload)
                    break
                assert chunk is not None
                for payload in chunk_payloads(chunk):
                    await websocket.send_json(payload)
        finally:
            stream.close()

    @app.post("/v1/completions")
    async def create_completion(  # pyright: ignore[reportUnusedFunction]