实现http流式输出的最小实践

engine端

import asyncio
class AsyncStream:
    """A stream of RequestOutputs for a request that can be
    iterated over asynchronously."""

    def __init__(self) -> None:
        self._queue = asyncio.Queue()

    def put(self, item: str) -> None:
        self._queue.put_nowait(item)

    def __aiter__(self):
        return self

    async def __anext__(self) -> str:
        result = await self._queue.get()
        return result

class Engine():
    
    def __init__(self):
        self.items = ['my', 'name', 'is', 'lewis']
        self.stream = AsyncStream()
    async def generate(self):
        async for request_output in self.stream:
                print(request_output)
                yield request_output
        
def get_engine():
    return Engine()

server端

import argparse
import json
from typing import AsyncGenerator

from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, Response, StreamingResponse
import uvicorn

from api_engine import get_engine

TIMEOUT_KEEP_ALIVE = 5  # seconds.
TIMEOUT_TO_PREVENT_DEADLOCK = 1  # seconds.
app = FastAPI()
engine = get_engine()


@app.post("/generate")
async def generate(request: Request) -> Response:

    for item in engine.items:
        print('push:', item)
        engine.stream.put(item)
    results_generator = engine.generate()

    # Streaming case
    async def stream_results() -> AsyncGenerator[bytes, None]:
        async for request_output in results_generator:
            ret = {"text": request_output}
            yield (json.dumps(ret) + "\0").encode("utf-8")

    return StreamingResponse(stream_results())


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--port", type=int, default=8000)
    args = parser.parse_args()
    uvicorn.run(app,
                host='localhost',
                port=args.port,
                log_level="debug",
                timeout_keep_alive=TIMEOUT_KEEP_ALIVE)

client端

import requests
import json
from typing import Iterable, List

def get_streaming_response(response: requests.Response) -> Iterable[List[str]]:
    for chunk in response.iter_lines(chunk_size=8192,
                                     decode_unicode=False,
                                     delimiter=b"\0"):
        if chunk:
            data = json.loads(chunk.decode("utf-8"))
            output = data["text"]
            yield output

def post_http_request(
                      api_url: str,
                     ) -> requests.Response:
    headers = {"User-Agent": "Test Client"}
    response = requests.post(api_url, headers=headers, stream=True)
    return response

if __name__ == "__main__":
    api_url = f"http://localhost:8000/generate"
    response = post_http_request(api_url)
    num_printed_lines = 0
    for h in get_streaming_response(response):
        print(h)