diff --git a/libs/partners/mongodb/tests/integration_tests/test_chat_message_histories.py b/libs/partners/mongodb/tests/integration_tests/test_chat_message_histories.py index 5545ab455a..376422afff 100644 --- a/libs/partners/mongodb/tests/integration_tests/test_chat_message_histories.py +++ b/libs/partners/mongodb/tests/integration_tests/test_chat_message_histories.py @@ -31,7 +31,7 @@ def test_memory_with_message_store() -> None: assert "This is me, the AI" in messages_json assert "This is me, the human" in messages_json - # remove the record from Azure Cosmos DB, so the next test run won't pick it up + # remove the record from MongoDB, so the next test run won't pick it up memory.chat_memory.clear() assert memory.chat_memory.messages == [] diff --git a/libs/partners/mongodb/tests/unit_tests/test_chat_message_histories.py b/libs/partners/mongodb/tests/unit_tests/test_chat_message_histories.py new file mode 100644 index 0000000000..2c1889a43a --- /dev/null +++ b/libs/partners/mongodb/tests/unit_tests/test_chat_message_histories.py @@ -0,0 +1,40 @@ +import json + +from langchain.memory import ConversationBufferMemory +from langchain_core.messages import message_to_dict + +from langchain_mongodb.chat_message_histories import MongoDBChatMessageHistory +from tests.utils import MockCollection + + +class PatchedMongoDBChatMessageHistory(MongoDBChatMessageHistory): + def __init__(self) -> None: + self.session_id = "test-session" + self.database_name = "test-database" + self.collection_name = "test-collection" + self.collection = MockCollection() + + +def test_memory_with_message_store() -> None: + """Test the memory with a message store.""" + # setup MongoDB as a message store + message_history = PatchedMongoDBChatMessageHistory() + memory = ConversationBufferMemory( + memory_key="baz", chat_memory=message_history, return_messages=True + ) + + # add some messages + memory.chat_memory.add_ai_message("This is me, the AI") + memory.chat_memory.add_user_message("This is me, the human") + + # get the message history from the memory store and turn it into a json + messages = memory.chat_memory.messages + messages_json = json.dumps([message_to_dict(msg) for msg in messages]) + + assert "This is me, the AI" in messages_json + assert "This is me, the human" in messages_json + + # remove the record from MongoDB, so the next test run won't pick it up + memory.chat_memory.clear() + + assert memory.chat_memory.messages == [] diff --git a/libs/partners/mongodb/tests/utils.py b/libs/partners/mongodb/tests/utils.py index 1858ea1210..3716ac6947 100644 --- a/libs/partners/mongodb/tests/utils.py +++ b/libs/partners/mongodb/tests/utils.py @@ -169,18 +169,21 @@ class MockCollection(Collection): [k["_id"] for k in mongodb_inserts], acknowledged=True ) + def insert_one(self, to_insert: Any, *args, **kwargs) -> Any: # type: ignore + return self.insert_many([to_insert]) + def find_one(self, find_query: Dict[str, Any]) -> Optional[Dict[str, Any]]: # type: ignore + find = self.find(find_query) or [None] # type: ignore + return find[0] + + def find(self, find_query: Dict[str, Any]) -> Optional[List[Dict[str, Any]]]: # type: ignore def _is_match(item: Dict[str, Any]) -> bool: for key, match_val in find_query.items(): if item.get(key) != match_val: return False return True - # Return the first element to match - for document in self._data: - if _is_match(document): - return document - return None + return [document for document in self._data if _is_match(document)] def update_one( # type: ignore self,