Final touchups

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

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.

test/chat_message_stores/in_memory/test_chat_message_store.py

@@ -6,20 +6,29 @@
 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() == []
         assert store.write_messages([]) == 0
         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": {},
             "type": "haystack_experimental.chat_message_stores.in_memory.chat_message_store.InMemoryChatMessageStore"
         }
 
     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?")])

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?")])
 

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()