From e21212394e0a9e7dbdcb94902dbfef6593a51c50 Mon Sep 17 00:00:00 2001 From: Vladimir Blagojevic Date: Mon, 29 Jul 2024 16:13:58 +0100 Subject: [PATCH] Final touchups --- .../retrievers/chat_message_retriever.py | 2 +- .../components/writers/chat_message_writer.py | 4 ++-- .../in_memory/test_chat_message_store.py | 18 ++++++++++++++++++ .../retrievers/test_chat_message_retriever.py | 13 ++++++++++++- .../writers/test_chat_message_writer.py | 12 ++++++++++++ 5 files changed, 45 insertions(+), 4 deletions(-) diff --git a/haystack_experimental/components/retrievers/chat_message_retriever.py b/haystack_experimental/components/retrievers/chat_message_retriever.py index b2a88949..32acb804 100644 --- a/haystack_experimental/components/retrievers/chat_message_retriever.py +++ b/haystack_experimental/components/retrievers/chat_message_retriever.py @@ -16,7 +16,7 @@ @component class ChatMessageRetriever: """ - Retrieves chat messages. + Retrieves chat messages from the underlying ChatMessageStore. Usage example: ```python diff --git a/haystack_experimental/components/writers/chat_message_writer.py b/haystack_experimental/components/writers/chat_message_writer.py index deaa67ee..97e82f57 100644 --- a/haystack_experimental/components/writers/chat_message_writer.py +++ b/haystack_experimental/components/writers/chat_message_writer.py @@ -16,7 +16,7 @@ @component class ChatMessageWriter: """ - Writes chat messages to a ChatMessageStore. + Writes chat messages to an underlying ChatMessageStore. Usage example: ```python @@ -99,7 +99,7 @@ def run(self, messages: List[ChatMessage]): :param messages: A list of chat messages to write to the store. :returns: - - `messages_written`: Number of messages written + - `messages_written`: Number of messages written to the ChatMessageStore. :raises ValueError: If the specified message store is not found. diff --git a/test/chat_message_stores/in_memory/test_chat_message_store.py b/test/chat_message_stores/in_memory/test_chat_message_store.py index e4975881..8084d736 100644 --- a/test/chat_message_stores/in_memory/test_chat_message_store.py +++ b/test/chat_message_stores/in_memory/test_chat_message_store.py @@ -6,6 +6,9 @@ class TestChatMessageStore: def test_init(self): + """ + Test that the InMemoryChatMessageStore can be initialized and that it works as expected. + """ store = InMemoryChatMessageStore() assert store.count_messages() == 0 assert store.retrieve() == [] @@ -13,6 +16,9 @@ def test_init(self): assert not store.delete_messages() def test_to_dict(self): + """ + Test that the InMemoryChatMessageStore can be serialized to a dictionary. + """ store = InMemoryChatMessageStore() assert store.to_dict() == { "init_parameters": {}, @@ -20,6 +26,9 @@ def test_to_dict(self): } def test_from_dict(self): + """ + Test that the InMemoryChatMessageStore can be deserialized from a dictionary. + """ data = { "init_parameters": {}, "type": "haystack_experimental.chat_message_stores.in_memory.chat_message_store.InMemoryChatMessageStore" @@ -28,6 +37,9 @@ def test_from_dict(self): assert store.to_dict() == data def test_count_messages(self): + """ + Test that the InMemoryChatMessageStore can count the number of messages in the store correctly. + """ store = InMemoryChatMessageStore() assert store.count_messages() == 0 store.write_messages(messages=[ChatMessage.from_user(content="Hello, how can I help you?")]) @@ -38,6 +50,9 @@ def test_count_messages(self): assert store.count_messages() == 3 def test_retrieve(self): + """ + Test that the InMemoryChatMessageStore can retrieve all messages from the store correctly. + """ store = InMemoryChatMessageStore() assert store.retrieve() == [] store.write_messages(messages=[ChatMessage.from_user(content="Hello, how can I help you?")]) @@ -55,6 +70,9 @@ def test_retrieve(self): ] def test_delete_messages(self): + """ + Test that the InMemoryChatMessageStore can delete all messages from the store correctly. + """ store = InMemoryChatMessageStore() assert store.count_messages() == 0 store.write_messages(messages=[ChatMessage.from_user(content="Hello, how can I help you?")]) diff --git a/test/components/retrievers/test_chat_message_retriever.py b/test/components/retrievers/test_chat_message_retriever.py index 82ea6271..eb9b313b 100644 --- a/test/components/retrievers/test_chat_message_retriever.py +++ b/test/components/retrievers/test_chat_message_retriever.py @@ -8,7 +8,9 @@ class TestChatMessageRetriever: def test_init(self): - + """ + Test that the ChatMessageRetriever component can be initialized with a valid message store. + """ messages = [ ChatMessage.from_user(content="Hello, how can I help you?"), ChatMessage.from_user(content="Hallo, wie kann ich Ihnen helfen?") @@ -22,6 +24,9 @@ def test_init(self): assert retriever.run() == {"messages": messages} def test_to_dict(self): + """ + Test that the ChatMessageRetriever component can be serialized to a dictionary. + """ message_store = InMemoryChatMessageStore() retriever = ChatMessageRetriever(message_store) @@ -37,6 +42,9 @@ def test_to_dict(self): } def test_from_dict(self): + """ + Test that the ChatMessageRetriever component can be deserialized from a dictionary. + """ data = { "type": "haystack_experimental.components.retrievers.chat_message_retriever.ChatMessageRetriever", "init_parameters": { @@ -53,6 +61,9 @@ def test_from_dict(self): } def test_chat_message_retriever_pipeline(self): + """ + Test that the ChatMessageRetriever can be used in a pipeline and that it works as expected. + """ store = InMemoryChatMessageStore() store.write_messages([ChatMessage.from_assistant("Hello, how can I help you?")]) diff --git a/test/components/writers/test_chat_message_writer.py b/test/components/writers/test_chat_message_writer.py index 0fe5e87d..4fed9ab2 100644 --- a/test/components/writers/test_chat_message_writer.py +++ b/test/components/writers/test_chat_message_writer.py @@ -8,6 +8,9 @@ class TestChatMessageRetriever: def test_init(self): + """ + Test that the ChatMessageWriter component can be initialized with a valid message store. + """ messages = [ ChatMessage.from_user(content="Hello, how can I help you?"), ChatMessage.from_user(content="Hallo, wie kann ich Ihnen helfen?") @@ -21,6 +24,9 @@ def test_init(self): assert retriever.run(messages=messages) == {"messages_written": 2} def test_to_dict(self): + """ + Test that the ChatMessageWriter component can be serialized to a dictionary. + """ message_store = InMemoryChatMessageStore() retriever = ChatMessageWriter(message_store) @@ -36,6 +42,9 @@ def test_to_dict(self): } def test_from_dict(self): + """ + Test that the ChatMessageWriter component can be deserialized from a dictionary. + """ data = { "type": "haystack_experimental.components.writers.chat_message_writer.ChatMessageWriter", "init_parameters": { @@ -52,6 +61,9 @@ def test_from_dict(self): } def test_chat_message_writer_pipeline(self): + """ + Test that the ChatMessageWriter can be used in a pipeline and that it works as expected. + """ store = InMemoryChatMessageStore() pipe = Pipeline()