File size: 838 Bytes
77320e4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
import json

from sse_starlette import ServerSentEvent

from infiagent.schemas import ResponseBaseData


IGNORE_PING_COMMENT = {"comment": "IGNORE PING"}
DONE = "[DONE]"


async def async_sse_response_format(response_data_gen):
    async for content in response_data_gen:
        if content == DONE:
            sse_event = ServerSentEvent(data=DONE)
        else:
            data_dict = {
                "response": content,
                "ResponseBase": ResponseBaseData().dict()
            }
            sse_event = ServerSentEvent(data=json.dumps(data_dict, ensure_ascii=False))
        yield sse_event


def json_response_format(content):
    return {
        "response": content,
        "ResponseBase": ResponseBaseData().dict()
    }


def get_ignore_ping_comment():
    return lambda: ServerSentEvent(**IGNORE_PING_COMMENT)