MetaGPT

1
# -*- coding: utf-8 -*-
2
# @Time    : 2024/3/27 10:00
3
# @Author  : leiwu30
4
# @File    : stream_pipe.py
5
# @Version : None
6
# @Description : None
7

8
import json
9
import time
10
from multiprocessing import Pipe
11

12

13
class StreamPipe:
14
    parent_conn, child_conn = Pipe()
15
    finish: bool = False
16

17
    format_data = {
18
        "id": "chatcmpl-96bVnBOOyPFZZxEoTIGbdpFcVEnur",
19
        "object": "chat.completion.chunk",
20
        "created": 1711361191,
21
        "model": "gpt-3.5-turbo-0125",
22
        "system_fingerprint": "fp_3bc1b5746c",
23
        "choices": [
24
            {"index": 0, "delta": {"role": "assistant", "content": "content"}, "logprobs": None, "finish_reason": None}
25
        ],
26
    }
27

28
    def set_message(self, msg):
29
        self.parent_conn.send(msg)
30

31
    def get_message(self, timeout: int = 3):
32
        if self.child_conn.poll(timeout):
33
            return self.child_conn.recv()
34
        else:
35
            return None
36

37
    def msg2stream(self, msg):
38
        self.format_data["created"] = int(time.time())
39
        self.format_data["choices"][0]["delta"]["content"] = msg
40
        return f"data: {json.dumps(self.format_data, ensure_ascii=False)}\n".encode("utf-8")
41