File size: 1,936 Bytes
e2a9bea |
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 |
import json
class OpenaiStreamOutputer:
"""
Create chat completion - OpenAI API Documentation
* https://platform.openai.com/docs/api-reference/chat/create
"""
def __init__(self):
self.default_data = {
"created": 1700000000,
"id": "chatcmpl-hugginface",
"object": "chat.completion.chunk",
# "content_type": "Completions",
"model": "hugginface",
"choices": [],
}
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)
|