-
Notifications
You must be signed in to change notification settings - Fork 0
/
utils.py
29 lines (23 loc) · 1.23 KB
/
utils.py
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
from modelhelper import num_tokens_from_messages
class MessageBuilder:
"""
A class for building and managing messages in a chat conversation.
Attributes:
message (list): A list of dictionaries representing chat messages.
model (str): The name of the ChatGPT model.
token_count (int): The total number of tokens in the conversation.
Methods:
__init__(self, system_content: str, chatgpt_model: str): Initializes the MessageBuilder instance.
append_message(self, role: str, content: str, index: int = 1): Appends a new message to the conversation.
"""
def __init__(self, system_content: str, chatgpt_model: str):
self.messages = [{'role': 'system', 'content': system_content}]
self.model = chatgpt_model
self.token_length = num_tokens_from_messages(
self.messages[-1], self.model)
def append_message(self, role: str, content: str, index: int = 1):
self.messages.insert(index, {'role': role, 'content': content})
self.token_length += num_tokens_from_messages(
self.messages[index], self.model)
def nonewlines(s: str) -> str:
return s.replace('\n', ' ').replace('\r', ' ')