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

engine端

import asyncio
class AsyncStream:
    #实现了__aiter__和__anext__表示这是一个异步队列，而不是同步队列，所以不能用for _ in 访问
    #内部元素，只能用async for _ in 访问
 
    def __init__(self, request_id: str) -> None:
        self.request_id = request_id
        self._queue = asyncio.Queue()
        self._finished = False
 
    def put(self, item) -> None:
        if self._finished:
            return
        self._queue.put_nowait(item)
 
    def finish(self) -> None:
        self._queue.put_nowait(StopIteration)
        self._finished = True
 
    @property
    def finished(self) -> bool:
        return self._finished
 
    def __aiter__(self):
        return self
 
    async def __anext__(self):
        result = await self._queue.get()
        if result is StopIteration:
            raise StopAsyncIteration
        elif isinstance(result, Exception):
            raise result
        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)