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

116 lines
5 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 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 TestTruncateMessage:
Update truncation test to reduce flakyness in cloud tests Removed dependency on faker, factory for the truncation tests as that seems to be the point of flakiness
2024-06-07 14:10:53 +00:00
max_prompt_size = 10
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
Update truncation test to reduce flakyness in cloud tests Removed dependency on faker, factory for the truncation tests as that seems to be the point of flakiness
2024-06-07 14:10:53 +00:00
chat_history = generate_chat_history(50)
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
Update truncation test to reduce flakyness in cloud tests Removed dependency on faker, factory for the truncation tests as that seems to be the point of flakiness
2024-06-07 14:10:53 +00:00
assert len(chat_history) < 50
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
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
Update truncation test to reduce flakyness in cloud tests Removed dependency on faker, factory for the truncation tests as that seems to be the point of flakiness
2024-06-07 14:10:53 +00:00
chat_history = generate_chat_history(5)
big_chat_message = ChatMessage(role="user", content=f"{generate_content(6)}\nQuestion?")
Move message truncation logic into a separate function. Add unit tests with factory boy.
2023-06-06 01:58:29 +00:00
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
Update truncation test to reduce flakyness in cloud tests Removed dependency on faker, factory for the truncation tests as that seems to be the point of flakiness
2024-06-07 14:10:53 +00:00
chat_history = generate_chat_history(5)
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[0].role = "system" # Mark the first message as system message
Update truncation test to reduce flakyness in cloud tests Removed dependency on faker, factory for the truncation tests as that seems to be the point of flakiness
2024-06-07 14:10:53 +00:00
big_chat_message = ChatMessage(role="user", content=f"{generate_content(11)}\nQuestion?")
Move message truncation logic into a separate function. Add unit tests with factory boy.
2023-06-06 01:58:29 +00:00
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.
Update truncation test to reduce flakyness in cloud tests Removed dependency on faker, factory for the truncation tests as that seems to be the point of flakiness
2024-06-07 14:10:53 +00:00
assert (
len(truncated_chat_history) == len(chat_history) + 1
) # Because the system_prompt is popped off from the chat_messages list
assert len(truncated_chat_history) < 10
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) > 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
Update truncation test to reduce flakyness in cloud tests Removed dependency on faker, factory for the truncation tests as that seems to be the point of flakiness
2024-06-07 14:10:53 +00:00
big_chat_message = ChatMessage(role="user", content=f"{generate_content(11)}\nQuestion?")
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
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))
Update truncation test to reduce flakyness in cloud tests Removed dependency on faker, factory for the truncation tests as that seems to be the point of flakiness
2024-06-07 14:10:53 +00:00
big_chat_message = ChatMessage(role="user", content=big_chat_message_content)
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
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
Update truncation test to reduce flakyness in cloud tests Removed dependency on faker, factory for the truncation tests as that seems to be the point of flakiness
2024-06-07 14:10:53 +00:00
def generate_content(count):
return " ".join([f"{index}" for index, _ in enumerate(range(count))])
def generate_chat_history(count):
return [
ChatMessage(role="user" if index % 2 == 0 else "assistant", content=f"{index}")
for index, _ in enumerate(range(count))
]
Reference in a new issue Copy permalink