Flask创建流式返回的mock脚本

好的！文西2024-10-15 22:07

需求：对接口api/mocktest 接口进行流式返回的mock，其中要求接口请求方式为：api/mocktest/{first_time}/{total_time}，first_time为首token耗时、total_time为流式返回总耗时

python 复制代码

import sys
from flask import Flask, Response, request
import json
import time

app = Flask(__name__)

@app.route("/api/mocktest/<int:first_time>/<int:total_time>", methods=["POST"])
def process(first_time, total_time):
    if first_time >= total_time:
        return Response(json.dumps({"error": "First time is greater than total time"}), status=400)

    """
    /api接口定义，动态获取first token time和complete response time
    """
    stream = request.json.get("stream")
    stream_num = 10
    # 根据请求参数计算平均时间
    avg_time = (total_time - first_time) / (stream_num - 1)
    # 初始化流式数据
    mock_stream_data = [[json.dumps({"data": "this is a stream mock1", "status": 0}), first_time]]
    for i in range(2, stream_num + 1):
        res = [json.dumps({"data": "this is a stream mock" + str(i), "status": 0}), avg_time]
        mock_stream_data.append(res)

    # 初始化非流式数据
    no_stream_time = first_time
    mock_nostream_data = [
        [json.dumps({"data": "this is a nostream mock", "status": 0}), no_stream_time]
    ]

    def generate_stream():
        """
        流式返回数据
        """
        for each in mock_stream_data:
            v = each[0]
            t = each[1]
            if t > 0:
                time.sleep(t)  # 等待指定的时间后再发送下一条数据
            yield v + "\r\n"
            sys.stdout.flush()

    def generate_nostream():
        """
        非流式返回数据
        """
        v = mock_nostream_data[0][0]
        t = mock_nostream_data[0][1]
        if t > 0:
            time.sleep(t)
        yield v + "\r\n"

    if stream:
        return Response(generate_stream(), mimetype="text/plain")
    else:
        return Response(generate_nostream(), mimetype="text/plain")

app.run()

python运行后通过curl命令访问地址：

流式请求:

bash 复制代码

curl --location '127.0.0.1:5000/api/mocktest/3/10' \
--header 'Content-Type: application/json' \
--data '{
    "messages": [{"role": "user", "content": "发起请求"}],
    "stream": true
}'

结果：

非流式请求：

bash 复制代码

curl --location '127.0.0.1:5000/api/mocktest/3/10' \
--header 'Content-Type: application/json' \
--data '{
    "messages": [{"role": "user", "content": "发起请求"}],
    "stream": false
}'

结果：