File size: 1,327 Bytes
da27d7f ed570c3 da27d7f ed570c3 da27d7f ed570c3 da27d7f ed570c3 da27d7f ed570c3 da27d7f ed570c3 |
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 |
import json
import time
class OpenaiStreamOutputer:
def __init__(self):
self.default_data = {
"id": "chatcmpl-123",
"object": "chat.completion",
"created": None, # Will be set to current time in output
"model": "gpt-3.5-turbo-0613",
"system_fingerprint": "fp_44709d6fcb",
"choices": [],
"usage": {
"prompt_tokens": 0,
"completion_tokens": 0,
"total_tokens": 0
}
}
def data_to_string(self, data):
return json.dumps(data, indent=2)
def output(self, content=None, role="assistant", prompt_tokens=0, completion_tokens=0):
data = self.default_data.copy()
# Set the current Unix timestamp
data["created"] = int(time.time())
data["choices"] = [{
"index": 0,
"message": {
"role": role,
"content": content,
},
"logprobs": None,
"finish_reason": "stop"
}]
# Update token usage
data["usage"]["prompt_tokens"] = prompt_tokens
data["usage"]["completion_tokens"] = completion_tokens
data["usage"]["total_tokens"] = prompt_tokens + completion_tokens
return self.data_to_string(data) |