klbr/khoj
1
0
Fork 0
mirror of https://github.com/khoaliber/khoj.git synced 2026-03-02 13:18:18 +00:00
Code Issues Packages Projects Releases Wiki Activity

Rename truncate messages method and update unit tests to simplify assertion logic

Browse Source
This commit is contained in:
Saba
2023-06-06 23:25:43 -07:00
parent 7119ed0849
commit 5d5ebcbf7c
2 changed files with 9 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
tests/test_conversation_utils.py
View File

@@ -19,18 +19,15 @@ class TestTruncateMessage:
def test_truncate_message_all_small(self):
chat_messages = ChatMessageFactory.build_batch(500)
assert len(chat_messages) == 500
tokens = sum([len(self.encoder.encode(message.content)) for message in chat_messages])
assert tokens > self.max_prompt_size
prompt = utils.truncate_message(chat_messages, self.max_prompt_size, self.model_name)
prompt = utils.truncate_messages(chat_messages, self.max_prompt_size, self.model_name)
tokens = sum([len(self.encoder.encode(message.content)) for message in prompt])
# The original object has been modified. Verify certain properties
assert len(chat_messages) < 500
assert len(chat_messages) > 1
assert prompt == chat_messages
tokens = sum([len(self.encoder.encode(message.content)) for message in prompt])
assert tokens <= self.max_prompt_size
def test_truncate_message_first_large(self):
@@ -39,18 +36,14 @@ class TestTruncateMessage:
big_chat_message.content = big_chat_message.content + "\n" + "Question?"
copy_big_chat_message = big_chat_message.copy()
chat_messages.insert(0, big_chat_message)
assert len(chat_messages) == 26
tokens = sum([len(self.encoder.encode(message.content)) for message in chat_messages])
assert tokens > self.max_prompt_size
prompt = utils.truncate_message(chat_messages, self.max_prompt_size, self.model_name)
prompt = utils.truncate_messages(chat_messages, self.max_prompt_size, self.model_name)
tokens = sum([len(self.encoder.encode(message.content)) for message in prompt])
# The original object has been modified. Verify certain properties
assert len(chat_messages) < 26
assert len(chat_messages) == 1
assert prompt[0] != copy_big_chat_message
tokens = sum([len(self.encoder.encode(message.content)) for message in prompt])
assert tokens <= self.max_prompt_size
def test_truncate_message_last_large(self):
@@ -60,16 +53,13 @@ class TestTruncateMessage:
copy_big_chat_message = big_chat_message.copy()
chat_messages.append(big_chat_message)
assert len(chat_messages) == 26
tokens = sum([len(self.encoder.encode(message.content)) for message in chat_messages])
assert tokens > self.max_prompt_size
prompt = utils.truncate_message(chat_messages, self.max_prompt_size, self.model_name)
prompt = utils.truncate_messages(chat_messages, self.max_prompt_size, self.model_name)
tokens = sum([len(self.encoder.encode(message.content)) for message in prompt])
# The original object has been modified. Verify certain properties
assert len(chat_messages) < 26
assert len(chat_messages) > 1
assert prompt[0] != copy_big_chat_message
tokens = sum([len(self.encoder.encode(message.content)) for message in prompt])
assert tokens < self.max_prompt_size
assert tokens <= self.max_prompt_size