python 搭建sse服务

服务端

from fastapi import FastAPI
from fastapi.responses import StreamingResponse
import asyncio
import time
from pydantic import BaseModel

app = FastAPI()


# 定义一个Pydantic模型来描述JSON数据的结构
class Item(BaseModel):
    msg: str


@app.post("/stream")
async def stream(item: Item):

    async def event_stream():
        while True:
            # 模拟一个事件
            yield '{} data: {}\n\n'.format(item.msg, time.ctime())
            # 立即发送事件，而不是等待缓冲
            await asyncio.sleep(2)  # 每2秒发送一次事件

    return StreamingResponse(event_stream(), media_type="text/event-stream")


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

代码客户端

import requests
import json

# SSE URL
url = 'http://localhost:8000/stream'

# JSON payload
payload = {'msg': '你好'}

# Headers for SSE request
headers = {'Content-Type': 'application/json', 'Accept': 'text/event-stream'}

# Send SSE request
response = requests.post(url,
                         data=json.dumps(payload),
                         headers=headers,
                         stream=True)

# Check if the request was successful
if response.status_code == 200:
    # Process the SSE stream
    for line in response.iter_lines():
        if line:
            decoded_line = line.decode('utf-8')
            print(decoded_line)
else:
    print(f"Request failed with status code: {response.status_code}")

# Close the connection
response.close()

webui客户端

import random
import gradio as gr
import requests
import json
"""
用户输入后的回调函数 random_response
参数：
message: 用户此次输入的消息
history: 历史聊天记录，比如 [["use input 1", "assistant output 1"],["user input 2", "assistant output 2"]]
​
返回值：输出的内容
"""


def random_response(message, history):
    return random.choice(["Yes", "No"])


def chat(message, history):
    # SSE URL
    url = 'http://localhost:8000/stream'

    # JSON payload
    payload = {'msg': '你好'}

    # Headers for SSE request
    headers = {
        'Content-Type': 'application/json',
        'Accept': 'text/event-stream'
    }

    # Send SSE request
    response = requests.post(url,
                             data=json.dumps(payload),
                             headers=headers,
                             stream=True)

    # Check if the request was successful
    if response.status_code == 200:
        # Process the SSE stream
        for line in response.iter_lines():
            if line:
                decoded_line = line.decode('utf-8')
                yield decoded_line
    else:
        print(f"Request failed with status code: {response.status_code}")

    # Close the connection
    response.close()


gr.ChatInterface(chat).launch(server_name='0.0.0.0',
                                         server_port=8001)

调用openai服务示例

import requests
import json
def sse_send(
    url,
    model,
    api_key,
    prompt="You are a helpful assistant.",
    msg="",
    top_p=1.0,
    temperature=0.7,
):

    # JSON payload
    data = {
        "model": model,
        "messages": [
            {"role": "system", "content": prompt},
            {"role": "user", "content": msg},
        ],
        "stream": True,
        "stop": ["<|EOT|>", "<|im_end|>"],
        "top_p": top_p,
        "temperature": temperature,
        "repetition_penalty": 1.05,
    }

    # Headers for SSE request
    headers = {
        "Content-Type": "application/json",
        "Accept": "text/event-stream",
        "Authorization": f"Bearer {api_key}",
    }
    print("请求\n", data)
    data = json.dumps(data)
    # Send SSE request
    response = requests.post(url, data=data, headers=headers, stream=True, timeout=300)
    #print("响应\n", response.text)
    res = ""
    # Check if the request was successful
    if response.status_code == 200:
        # Process the SSE stream
        for line in response.iter_lines():
            if line:
                decoded_line = line.decode("utf-8")
                print(decoded_line)
                if "data: [DONE]" == decoded_line:
                    continue

                decoded_line = str(decoded_line).replace("data: ", "")
                decoded_line = json.loads(decoded_line)
                if len(decoded_line["choices"])==0:
                    continue
                if decoded_line["choices"][0]["delta"]["content"]==None:
                    continue
                decoded_line = decoded_line["choices"][0]["delta"]["content"]
                res += decoded_line
                print("新结果")
                print(res)
    else:
        print(f"Request failed with status code: {response.status_code}")

    # Close the connection
    response.close()
    return res