messagers/message_outputer.py

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",
            "model": "hugginface",
            "system_fingerprint": "fp_44709d6fcb",
            "choices": [],
            "usage": {
                "prompt_tokens": 0,
                "completion_tokens": 0,
                "total_tokens": 0
            }
        }

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

    def output(self, content=None, role="assistant", content_type="Completions", tokens_info=None) -> str:
        data = self.default_data.copy()
        message = {"role": role, "content": content}

        if tokens_info is None:
            tokens_info = {"prompt_tokens": 9, "completion_tokens": 12, "total_tokens": 21}

        data["choices"] = [
            {
                "index": 0,
                "message": message,
                "logprobs": None,
                "finish_reason": "stop" if content_type == "Finished" else None
            }
        ]

        data["usage"] = tokens_info
        return self.data_to_string(data)