model-evaluation / conversation.py
AlekseyKorshuk's picture
updates
cb80e0b
class Conversation:
bot_id: str
memory: str
prompt: str
bot_label: str
user_label: str
messages: list
def __init__(self, bot_config):
self.bot_id = bot_config.get("bot_id")
self.memory = bot_config.get("memory", "")
self.prompt = bot_config.get("prompt", "")
self.bot_label = bot_config.get("botLabel", "Character")
self.user_label = bot_config.get("userLabel", "User")
self.first_message = bot_config.get("firstMessage", f"Hi, my name is {self.bot_label}!")
self.reset_conversation()
def reset_conversation(self):
self.messages = [
{
"from": self.bot_label,
"value": self.first_message
}
]
def set_chat_history(self, chat_history):
messages = []
for conversation_pair in chat_history:
for item in conversation_pair:
if item:
messages.append(item)
self.messages = []
for i, message in enumerate(messages):
label = self.bot_label if i % 2 == 0 else self.user_label
self.messages.append(
{
"from": label,
"value": message.strip()
}
)
def add_user_message(self, message):
self.messages.append(
{
"from": self.user_label,
"value": message.strip()
}
)
def reset_last_message(self, message):
self.messages[-1]["value"] = message.strip()