File size: 2,157 Bytes
da27d7f
e994f7a
5484efc
43c8935
da27d7f
43c8935
 
 
 
 
da27d7f
e994f7a
da27d7f
9b03fdf
 
e994f7a
9b03fdf
e994f7a
da27d7f
e994f7a
 
 
 
 
da27d7f
 
921b216
43c8935
 
da27d7f
7f5564f
da27d7f
6f83d4a
 
 
 
7f5564f
 
6f83d4a
 
 
 
 
 
 
 
 
 
 
 
 
7f5564f
6f83d4a
 
 
 
 
 
 
7f5564f
6f83d4a
 
 
 
 
 
 
7f5564f
6f83d4a
 
 
7f5564f
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
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
import json
import time
import tiktoken

class OpenaiStreamOutputer:
    """
    Create chat completion - OpenAI API Documentation
    * https://platform.openai.com/docs/api-reference/chat/create
    """

    def __init__(self):
        current_time = int(time.time())
        self.default_data = {
            "id": "chatcmpl-hugginface",
            "object": "chat.completion.chunk",
            "created": current_time,
            "model": "hugginface",
            "system_fingerprint": "fp_44709d6fcb",
            "choices": [],
            "usage": {
                "prompt_tokens": 0,
                "completion_tokens": 0,
                "total_tokens": 0
            }
        }

    def data_to_string(self, data={}, content_type=""):
        data_str = f"{json.dumps(data)}"
        return data_str

    def output(self, content=None, content_type="Completions") -> str:
        data = self.default_data.copy()
        if content_type == "Role":
            data["choices"] = [
                {
                    "index": 0,
                    "delta": {"role": "assistant"},
                    "finish_reason": None,
                }
            ]
        elif content_type in [
            "Completions",
            "InternalSearchQuery",
            "InternalSearchResult",
            "SuggestedResponses",
        ]:
            if content_type in ["InternalSearchQuery", "InternalSearchResult"]:
                content += "\n"
            data["choices"] = [
                {
                    "index": 0,
                    "delta": {"content": content},
                    "finish_reason": None,
                }
            ]
        elif content_type == "Finished":
            data["choices"] = [
                {
                    "index": 0,
                    "delta": {},
                    "finish_reason": "stop",
                }
            ]
        else:
            data["choices"] = [
                {
                    "index": 0,
                    "delta": {},
                    "finish_reason": None,
                }
            ]
        return self.data_to_string(data, content_type)