| 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) |