File size: 1,310 Bytes
bf41ce7 1fac9b0 bf41ce7 | 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 | """chat response model"""
from typing import Any
from Brain.src.model.message_model import MessageModel
class ChatResponseModel:
def __init__(self, data: Any):
self.id = data["id"]
self.object = data["object"]
self.model = data["model"]
self.usage = data["usage"]
"""'choices': [
{
'message': {
'role': 'assistant',
'content': 'The 2020 World Series was played in Arlington, Texas at the Globe Life Field, which was the new home stadium for the Texas Rangers.'},
'finish_reason': 'stop',
'index': 0
}
]"""
self.choices = []
for choice in data["choices"]:
self.choices.append(ChoiceModel(choice))
def get_one_message_item(self) -> MessageModel:
if len(self.choices) > 0:
return self.choices.__getitem__(0).get_message()
else:
return MessageModel("", "")
class ChoiceModel:
def __init__(self, data: Any):
tmp_message = data["message"]
self.message = MessageModel(tmp_message["role"], tmp_message["content"])
self.finish_reason = data["finish_reason"]
self.index = data["index"]
def get_message(self):
return self.message
|