Implemented appending arbitrary messages

langchain/memory/chat_message_histories/cassandra.py

@@ -3,10 +3,8 @@
 from typing import List
 
 from langchain.schema import (
-    AIMessage,
     BaseChatMessageHistory,
     BaseMessage,
-    HumanMessage,
     _message_to_dict,
     messages_from_dict,
 )
@@ -143,13 +141,7 @@ def messages(self) -> List[BaseMessage]:  # type: ignore
 
         return messages
 
-    def add_user_message(self, message: str) -> None:
-        self.append(HumanMessage(content=message))
-
-    def add_ai_message(self, message: str) -> None:
-        self.append(AIMessage(content=message))
-
-    def append(self, message: BaseMessage) -> None:
+    def add_message(self, message: BaseMessage) -> None:
         """Append the message to the record in Cassandra"""
 
         import uuid

langchain/memory/chat_message_histories/cosmos_db.py

@@ -6,10 +6,8 @@
 from typing import TYPE_CHECKING, Any, List, Optional, Type
 
 from langchain.schema import (
-    AIMessage,
     BaseChatMessageHistory,
     BaseMessage,
-    HumanMessage,
     messages_from_dict,
     messages_to_dict,
 )
@@ -140,18 +138,13 @@ def load_messages(self) -> None:
         if "messages" in item and len(item["messages"]) > 0:
             self.messages = messages_from_dict(item["messages"])
 
-    def add_user_message(self, message: str) -> None:
-        """Add a user message to the memory."""
-        self.upsert_messages(HumanMessage(content=message))
-
-    def add_ai_message(self, message: str) -> None:
-        """Add a AI message to the memory."""
-        self.upsert_messages(AIMessage(content=message))
+    def add_message(self, message: BaseMessage) -> None:
+        """Add a self-created message to the store"""
+        self.messages.append(message)
+        self.upsert_messages()
 
-    def upsert_messages(self, new_message: Optional[BaseMessage] = None) -> None:
+    def upsert_messages(self) -> None:
         """Update the cosmosdb item."""
-        if new_message:
-            self.messages.append(new_message)
         if not self._container:
             raise ValueError("Container not initialized")
         self._container.upsert_item(

langchain/memory/chat_message_histories/dynamodb.py

@@ -2,10 +2,8 @@
 from typing import List
 
 from langchain.schema import (
-    AIMessage,
     BaseChatMessageHistory,
     BaseMessage,
-    HumanMessage,
     _message_to_dict,
     messages_from_dict,
     messages_to_dict,
@@ -53,13 +51,7 @@ def messages(self) -> List[BaseMessage]:  # type: ignore
         messages = messages_from_dict(items)
         return messages
 
-    def add_user_message(self, message: str) -> None:
-        self.append(HumanMessage(content=message))
-
-    def add_ai_message(self, message: str) -> None:
-        self.append(AIMessage(content=message))
-
-    def append(self, message: BaseMessage) -> None:
+    def add_message(self, message: BaseMessage) -> None:
         """Append the message to the record in DynamoDB"""
         from botocore.exceptions import ClientError
 

langchain/memory/chat_message_histories/file.py

@@ -4,10 +4,8 @@
 from typing import List
 
 from langchain.schema import (
-    AIMessage,
     BaseChatMessageHistory,
     BaseMessage,
-    HumanMessage,
     messages_from_dict,
     messages_to_dict,
 )
@@ -36,13 +34,7 @@ def messages(self) -> List[BaseMessage]:  # type: ignore
         messages = messages_from_dict(items)
         return messages
 
-    def add_user_message(self, message: str) -> None:
-        self.append(HumanMessage(content=message))
-
-    def add_ai_message(self, message: str) -> None:
-        self.append(AIMessage(content=message))
-
-    def append(self, message: BaseMessage) -> None:
+    def add_message(self, message: BaseMessage) -> None:
         """Append the message to the record in the local file"""
         messages = messages_to_dict(self.messages)
         messages.append(messages_to_dict([message])[0])

langchain/memory/chat_message_histories/firestore.py

@@ -5,10 +5,8 @@
 from typing import TYPE_CHECKING, List, Optional
 
 from langchain.schema import (
-    AIMessage,
     BaseChatMessageHistory,
     BaseMessage,
-    HumanMessage,
     messages_from_dict,
     messages_to_dict,
 )
@@ -81,18 +79,12 @@ def load_messages(self) -> None:
             if "messages" in data and len(data["messages"]) > 0:
                 self.messages = messages_from_dict(data["messages"])
 
-    def add_user_message(self, message: str) -> None:
-        """Add a user message to the memory."""
-        self.upsert_messages(HumanMessage(content=message))
-
-    def add_ai_message(self, message: str) -> None:
-        """Add a AI message to the memory."""
-        self.upsert_messages(AIMessage(content=message))
+    def add_message(self, message: BaseMessage) -> None:
+        self.messages.append(message)
+        self.upsert_messages()
 
     def upsert_messages(self, new_message: Optional[BaseMessage] = None) -> None:
         """Update the Firestore document."""
-        if new_message:
-            self.messages.append(new_message)
         if not self._document:
             raise ValueError("Document not initialized")
         self._document.set(

langchain/memory/chat_message_histories/in_memory.py

@@ -3,21 +3,17 @@
 from pydantic import BaseModel
 
 from langchain.schema import (
-    AIMessage,
     BaseChatMessageHistory,
     BaseMessage,
-    HumanMessage,
 )
 
 
 class ChatMessageHistory(BaseChatMessageHistory, BaseModel):
     messages: List[BaseMessage] = []
 
-    def add_user_message(self, message: str) -> None:
-        self.messages.append(HumanMessage(content=message))
-
-    def add_ai_message(self, message: str) -> None:
-        self.messages.append(AIMessage(content=message))
+    def add_message(self, message: BaseMessage) -> None:
+        """Add a self-created message to the store"""
+        self.messages.append(message)
 
     def clear(self) -> None:
         self.messages = []

langchain/memory/chat_message_histories/momento.py

@@ -5,10 +5,8 @@
 from typing import TYPE_CHECKING, Any, Optional
 
 from langchain.schema import (
-    AIMessage,
     BaseChatMessageHistory,
     BaseMessage,
-    HumanMessage,
     _message_to_dict,
     messages_from_dict,
 )
@@ -143,23 +141,7 @@ def messages(self) -> list[BaseMessage]:  # type: ignore[override]
         else:
             raise Exception(f"Unexpected response: {fetch_response}")
 
-    def add_user_message(self, message: str) -> None:
-        """Store a user message in the cache.
-
-        Args:
-            message (str): The message to store.
-        """
-        self.__add_message(HumanMessage(content=message))
-
-    def add_ai_message(self, message: str) -> None:
-        """Store an AI message in the cache.
-
-        Args:
-            message (str): The message to store.
-        """
-        self.__add_message(AIMessage(content=message))
-
-    def __add_message(self, message: BaseMessage) -> None:
+    def add_message(self, message: BaseMessage) -> None:
         """Store a message in the cache.
 
         Args:

langchain/memory/chat_message_histories/mongodb.py

@@ -3,10 +3,8 @@
 from typing import List
 
 from langchain.schema import (
-    AIMessage,
     BaseChatMessageHistory,
     BaseMessage,
-    HumanMessage,
     _message_to_dict,
     messages_from_dict,
 )
@@ -68,13 +66,7 @@ def messages(self) -> List[BaseMessage]:  # type: ignore
         messages = messages_from_dict(items)
         return messages
 
-    def add_user_message(self, message: str) -> None:
-        self.append(HumanMessage(content=message))
-
-    def add_ai_message(self, message: str) -> None:
-        self.append(AIMessage(content=message))
-
-    def append(self, message: BaseMessage) -> None:
+    def add_message(self, message: BaseMessage) -> None:
         """Append the message to the record in MongoDB"""
         from pymongo import errors
 

langchain/memory/chat_message_histories/postgres.py

@@ -3,10 +3,8 @@
 from typing import List
 
 from langchain.schema import (
-    AIMessage,
     BaseChatMessageHistory,
     BaseMessage,
-    HumanMessage,
     _message_to_dict,
     messages_from_dict,
 )
@@ -55,13 +53,7 @@ def messages(self) -> List[BaseMessage]:  # type: ignore
         messages = messages_from_dict(items)
         return messages
 
-    def add_user_message(self, message: str) -> None:
-        self.append(HumanMessage(content=message))
-
-    def add_ai_message(self, message: str) -> None:
-        self.append(AIMessage(content=message))
-
-    def append(self, message: BaseMessage) -> None:
+    def add_message(self, message: BaseMessage) -> None:
         """Append the message to the record in PostgreSQL"""
         from psycopg import sql
 

langchain/memory/chat_message_histories/redis.py

@@ -3,10 +3,8 @@
 from typing import List, Optional
 
 from langchain.schema import (
-    AIMessage,
     BaseChatMessageHistory,
     BaseMessage,
-    HumanMessage,
     _message_to_dict,
     messages_from_dict,
 )
@@ -52,13 +50,7 @@ def messages(self) -> List[BaseMessage]:  # type: ignore
         messages = messages_from_dict(items)
         return messages
 
-    def add_user_message(self, message: str) -> None:
-        self.append(HumanMessage(content=message))
-
-    def add_ai_message(self, message: str) -> None:
-        self.append(AIMessage(content=message))
-
-    def append(self, message: BaseMessage) -> None:
+    def add_message(self, message: BaseMessage) -> None:
         """Append the message to the record in Redis"""
         self.redis_client.lpush(self.key, json.dumps(_message_to_dict(message)))
         if self.ttl:

langchain/memory/chat_message_histories/sql.py

@@ -7,10 +7,8 @@
 from sqlalchemy.orm import sessionmaker
 
 from langchain.schema import (
-    AIMessage,
     BaseChatMessageHistory,
     BaseMessage,
-    HumanMessage,
     _message_to_dict,
     messages_from_dict,
 )
@@ -61,13 +59,7 @@ def messages(self) -> List[BaseMessage]:  # type: ignore
             messages = messages_from_dict(items)
             return messages
 
-    def add_user_message(self, message: str) -> None:
-        self.append(HumanMessage(content=message))
-
-    def add_ai_message(self, message: str) -> None:
-        self.append(AIMessage(content=message))
-
-    def append(self, message: BaseMessage) -> None:
+    def add_message(self, message: BaseMessage) -> None:
         """Append the message to the record in db"""
         with self.Session() as session:
             jsonstr = json.dumps(_message_to_dict(message))

langchain/memory/chat_message_histories/zep.py

@@ -116,13 +116,7 @@ def _get_memory(self) -> Optional[Memory]:
             return None
         return zep_memory
 
-    def add_user_message(self, message: str) -> None:
-        self.append(HumanMessage(content=message))
-
-    def add_ai_message(self, message: str) -> None:
-        self.append(AIMessage(content=message))
-
-    def append(self, message: BaseMessage) -> None:
+    def add_message(self, message: BaseMessage) -> None:
         """Append the message to the Zep memory history"""
         from zep_python import Memory, Message
 

langchain/schema.py

@@ -234,32 +234,29 @@ def messages(self):
                        messages = json.loads(f.read())
                     return messages_from_dict(messages)
 
-               def add_user_message(self, message: str):
-                   message_ = HumanMessage(content=message)
-                   messages = self.messages.append(_message_to_dict(_message))
+               def add_message(self, message: BaseMessage) -> None:
+                   messages = self.messages.append(_message_to_dict(message))
                    with open(os.path.join(storage_path, session_id), 'w') as f:
                        json.dump(f, messages)
-
-               def add_ai_message(self, message: str):
-                   message_ = AIMessage(content=message)
-                   messages = self.messages.append(_message_to_dict(_message))
-                   with open(os.path.join(storage_path, session_id), 'w') as f:
-                       json.dump(f, messages)
-
+
                def clear(self):
                    with open(os.path.join(storage_path, session_id), 'w') as f:
                        f.write("[]")
     """
 
     messages: List[BaseMessage]
 
-    @abstractmethod
     def add_user_message(self, message: str) -> None:
         """Add a user message to the store"""
+        self.add_message(HumanMessage(content=message))
 
-    @abstractmethod
     def add_ai_message(self, message: str) -> None:
         """Add an AI message to the store"""
+        self.add_message(AIMessage(content=message))
+
+    def add_message(self, message: BaseMessage) -> None:
+        """Add a self-created message to the store"""
+        raise NotImplementedError
 
     @abstractmethod
     def clear(self) -> None:

tests/unit_tests/memory/chat_message_histories/test_zep.py

@@ -60,7 +60,7 @@ def test_add_ai_message(mocker: MockerFixture, zep_chat: ZepChatMessageHistory)
 
 @pytest.mark.requires("zep_python")
 def test_append(mocker: MockerFixture, zep_chat: ZepChatMessageHistory) -> None:
-    zep_chat.append(AIMessage(content="test message"))
+    zep_chat.add_message(AIMessage(content="test message"))
     zep_chat.zep_client.add_memory.assert_called_once()  # type: ignore