sij/khoj
1
0
Fork 0
mirror of https://github.com/khoj-ai/khoj.git synced 2024-12-18 02:27:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 23
khoj/tests/test_conversation_utils.py

119 lines
5.3 KiB
Python
Raw Normal View History

Move message truncation logic into a separate function. Add unit tests with factory boy.
2023-06-06 01:58:29 +00:00
import factory
import tiktoken
Add isort to the pre-commit configuration and apply it to the whole project (#595) * Apply isort to the entire repository * Fix missing import issues in text_to_entries * Fix imports in migration files
2023-12-28 12:34:02 +00:00
from langchain.schema import ChatMessage
from khoj.processor.conversation import utils
Move message truncation logic into a separate function. Add unit tests with factory boy.
2023-06-06 01:58:29 +00:00
Run pre-commit script
2023-06-06 02:29:23 +00:00
Move message truncation logic into a separate function. Add unit tests with factory boy.
2023-06-06 01:58:29 +00:00
class ChatMessageFactory(factory.Factory):
class Meta:
model = ChatMessage
Run pre-commit script
2023-06-06 02:29:23 +00:00
content = factory.Faker("paragraph")
role = factory.Faker("name")
Move message truncation logic into a separate function. Add unit tests with factory boy.
2023-06-06 01:58:29 +00:00
class TestTruncateMessage:
max_prompt_size = 4096
Run pre-commit script
2023-06-06 02:29:23 +00:00
model_name = "gpt-3.5-turbo"
Move message truncation logic into a separate function. Add unit tests with factory boy.
2023-06-06 01:58:29 +00:00
encoder = tiktoken.encoding_for_model(model_name)
def test_truncate_message_all_small(self):
Handle truncation when single long non-system chat message Previously was assuming the system prompt is being always passed as the first message. So expected there to be at least 2 messages in logs. This broke chat actors querying with single long non system message. A more robust way to extract system prompt is via the message role instead
2024-03-15 09:22:29 +00:00
# Arrange
chat_history = ChatMessageFactory.build_batch(500)
Move message truncation logic into a separate function. Add unit tests with factory boy.
2023-06-06 01:58:29 +00:00
Handle truncation when single long non-system chat message Previously was assuming the system prompt is being always passed as the first message. So expected there to be at least 2 messages in logs. This broke chat actors querying with single long non system message. A more robust way to extract system prompt is via the message role instead
2024-03-15 09:22:29 +00:00
# Act
truncated_chat_history = utils.truncate_messages(chat_history, self.max_prompt_size, self.model_name)
tokens = sum([len(self.encoder.encode(message.content)) for message in truncated_chat_history])
Move message truncation logic into a separate function. Add unit tests with factory boy.
2023-06-06 01:58:29 +00:00
Handle truncation when single long non-system chat message Previously was assuming the system prompt is being always passed as the first message. So expected there to be at least 2 messages in logs. This broke chat actors querying with single long non system message. A more robust way to extract system prompt is via the message role instead
2024-03-15 09:22:29 +00:00
# Assert
Move message truncation logic into a separate function. Add unit tests with factory boy.
2023-06-06 01:58:29 +00:00
# The original object has been modified. Verify certain properties
Handle truncation when single long non-system chat message Previously was assuming the system prompt is being always passed as the first message. So expected there to be at least 2 messages in logs. This broke chat actors querying with single long non system message. A more robust way to extract system prompt is via the message role instead
2024-03-15 09:22:29 +00:00
assert len(chat_history) < 500
assert len(chat_history) > 1
Move message truncation logic into a separate function. Add unit tests with factory boy.
2023-06-06 01:58:29 +00:00
assert tokens <= self.max_prompt_size
def test_truncate_message_first_large(self):
Handle truncation when single long non-system chat message Previously was assuming the system prompt is being always passed as the first message. So expected there to be at least 2 messages in logs. This broke chat actors querying with single long non system message. A more robust way to extract system prompt is via the message role instead
2024-03-15 09:22:29 +00:00
# Arrange
chat_history = ChatMessageFactory.build_batch(25)
Fix failing unit tests by hard-coding model presence of expected search types
2023-06-13 23:32:47 +00:00
big_chat_message = ChatMessageFactory.build(content=factory.Faker("paragraph", nb_sentences=2000))
Move message truncation logic into a separate function. Add unit tests with factory boy.
2023-06-06 01:58:29 +00:00
big_chat_message.content = big_chat_message.content + "\n" + "Question?"
copy_big_chat_message = big_chat_message.copy()
Handle truncation when single long non-system chat message Previously was assuming the system prompt is being always passed as the first message. So expected there to be at least 2 messages in logs. This broke chat actors querying with single long non system message. A more robust way to extract system prompt is via the message role instead
2024-03-15 09:22:29 +00:00
chat_history.insert(0, big_chat_message)
tokens = sum([len(self.encoder.encode(message.content)) for message in chat_history])
Move message truncation logic into a separate function. Add unit tests with factory boy.
2023-06-06 01:58:29 +00:00
Handle truncation when single long non-system chat message Previously was assuming the system prompt is being always passed as the first message. So expected there to be at least 2 messages in logs. This broke chat actors querying with single long non system message. A more robust way to extract system prompt is via the message role instead
2024-03-15 09:22:29 +00:00
# Act
truncated_chat_history = utils.truncate_messages(chat_history, self.max_prompt_size, self.model_name)
tokens = sum([len(self.encoder.encode(message.content)) for message in truncated_chat_history])
Move message truncation logic into a separate function. Add unit tests with factory boy.
2023-06-06 01:58:29 +00:00
Handle truncation when single long non-system chat message Previously was assuming the system prompt is being always passed as the first message. So expected there to be at least 2 messages in logs. This broke chat actors querying with single long non system message. A more robust way to extract system prompt is via the message role instead
2024-03-15 09:22:29 +00:00
# Assert
Move message truncation logic into a separate function. Add unit tests with factory boy.
2023-06-06 01:58:29 +00:00
# The original object has been modified. Verify certain properties
Handle truncation when single long non-system chat message Previously was assuming the system prompt is being always passed as the first message. So expected there to be at least 2 messages in logs. This broke chat actors querying with single long non system message. A more robust way to extract system prompt is via the message role instead
2024-03-15 09:22:29 +00:00
assert len(chat_history) == 1
assert truncated_chat_history[0] != copy_big_chat_message
Move message truncation logic into a separate function. Add unit tests with factory boy.
2023-06-06 01:58:29 +00:00
assert tokens <= self.max_prompt_size
def test_truncate_message_last_large(self):
Handle truncation when single long non-system chat message Previously was assuming the system prompt is being always passed as the first message. So expected there to be at least 2 messages in logs. This broke chat actors querying with single long non system message. A more robust way to extract system prompt is via the message role instead
2024-03-15 09:22:29 +00:00
# Arrange
chat_history = ChatMessageFactory.build_batch(25)
chat_history[0].role = "system" # Mark the first message as system message
Run pre-commit script
2023-06-06 02:29:23 +00:00
big_chat_message = ChatMessageFactory.build(content=factory.Faker("paragraph", nb_sentences=1000))
Move message truncation logic into a separate function. Add unit tests with factory boy.
2023-06-06 01:58:29 +00:00
big_chat_message.content = big_chat_message.content + "\n" + "Question?"
copy_big_chat_message = big_chat_message.copy()
Run pre-commit script
2023-06-06 02:29:23 +00:00
Handle truncation when single long non-system chat message Previously was assuming the system prompt is being always passed as the first message. So expected there to be at least 2 messages in logs. This broke chat actors querying with single long non system message. A more robust way to extract system prompt is via the message role instead
2024-03-15 09:22:29 +00:00
chat_history.insert(0, big_chat_message)
initial_tokens = sum([len(self.encoder.encode(message.content)) for message in chat_history])
Move message truncation logic into a separate function. Add unit tests with factory boy.
2023-06-06 01:58:29 +00:00
Handle truncation when single long non-system chat message Previously was assuming the system prompt is being always passed as the first message. So expected there to be at least 2 messages in logs. This broke chat actors querying with single long non system message. A more robust way to extract system prompt is via the message role instead
2024-03-15 09:22:29 +00:00
# Act
truncated_chat_history = utils.truncate_messages(chat_history, self.max_prompt_size, self.model_name)
final_tokens = sum([len(self.encoder.encode(message.content)) for message in truncated_chat_history])
Move message truncation logic into a separate function. Add unit tests with factory boy.
2023-06-06 01:58:29 +00:00
Handle truncation when single long non-system chat message Previously was assuming the system prompt is being always passed as the first message. So expected there to be at least 2 messages in logs. This broke chat actors querying with single long non system message. A more robust way to extract system prompt is via the message role instead
2024-03-15 09:22:29 +00:00
# Assert
Add additional check for chat_messages length in UT
2023-08-01 16:25:52 +00:00
# The original object has been modified. Verify certain properties.
Handle truncation when single long non-system chat message Previously was assuming the system prompt is being always passed as the first message. So expected there to be at least 2 messages in logs. This broke chat actors querying with single long non system message. A more robust way to extract system prompt is via the message role instead
2024-03-15 09:22:29 +00:00
assert len(truncated_chat_history) == (
len(chat_history) + 1
Add additional check for chat_messages length in UT
2023-08-01 16:25:52 +00:00
) # Because the system_prompt is popped off from the chat_messages lsit
Handle truncation when single long non-system chat message Previously was assuming the system prompt is being always passed as the first message. So expected there to be at least 2 messages in logs. This broke chat actors querying with single long non system message. A more robust way to extract system prompt is via the message role instead
2024-03-15 09:22:29 +00:00
assert len(truncated_chat_history) < 26
assert len(truncated_chat_history) > 1
assert truncated_chat_history[0] != copy_big_chat_message
assert initial_tokens > self.max_prompt_size
assert final_tokens <= self.max_prompt_size
def test_truncate_single_large_non_system_message(self):
# Arrange
big_chat_message = ChatMessageFactory.build(content=factory.Faker("paragraph", nb_sentences=2000))
big_chat_message.content = big_chat_message.content + "\n" + "Question?"
big_chat_message.role = "user"
copy_big_chat_message = big_chat_message.copy()
chat_messages = [big_chat_message]
initial_tokens = sum([len(self.encoder.encode(message.content)) for message in chat_messages])
# Act
truncated_chat_history = utils.truncate_messages(chat_messages, self.max_prompt_size, self.model_name)
final_tokens = sum([len(self.encoder.encode(message.content)) for message in truncated_chat_history])
# Assert
# The original object has been modified. Verify certain properties
assert initial_tokens > self.max_prompt_size
assert final_tokens <= self.max_prompt_size
assert len(chat_messages) == 1
assert truncated_chat_history[0] != copy_big_chat_message
Handle msg truncation when question is larger than max prompt size Notice and truncate the question it self at this point
2024-03-31 10:07:29 +00:00
def test_truncate_single_large_question(self):
# Arrange
big_chat_message_content = " ".join(["hi"] * (self.max_prompt_size + 1))
big_chat_message = ChatMessageFactory.build(content=big_chat_message_content)
big_chat_message.role = "user"
copy_big_chat_message = big_chat_message.copy()
chat_messages = [big_chat_message]
initial_tokens = sum([len(self.encoder.encode(message.content)) for message in chat_messages])
# Act
truncated_chat_history = utils.truncate_messages(chat_messages, self.max_prompt_size, self.model_name)
final_tokens = sum([len(self.encoder.encode(message.content)) for message in truncated_chat_history])
# Assert
# The original object has been modified. Verify certain properties
assert initial_tokens > self.max_prompt_size
assert final_tokens <= self.max_prompt_size
assert len(chat_messages) == 1
assert truncated_chat_history[0] != copy_big_chat_message
Reference in a new issue Copy permalink