forked from langchain-ai/langchain
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
mongodb[patch]: Migrate MongoDBChatMessageHistory (langchain-ai#18590)
## **Description** Migrate the `MongoDBChatMessageHistory` to the managed `langchain-mongodb` partner-package ## **Dependencies** None ## **Twitter handle** @mongodb ## **tests and docs** - [x] Migrate existing integration test - [x ]~ Convert existing integration test to a unit test~ Creation is out of scope for this ticket - [x ] ~Considering delaying work until langchain-ai#17470 merges to leverage the `MockCollection` object. ~ - [x] **Lint and test**: Run `make format`, `make lint` and `make test` from the root of the package(s) you've modified. See contribution guidelines for more: https://python.langchain.com/docs/contributing/ --------- Co-authored-by: Erick Friis <erick@langchain.dev>
- Loading branch information
1 parent
c87f267
commit 324c59d
Showing
8 changed files
with
858 additions
and
8 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,7 +1,9 @@ | ||
from langchain_mongodb.chat_message_histories import MongoDBChatMessageHistory | ||
from langchain_mongodb.vectorstores import ( | ||
MongoDBAtlasVectorSearch, | ||
) | ||
|
||
__all__ = [ | ||
"MongoDBAtlasVectorSearch", | ||
"MongoDBChatMessageHistory", | ||
] |
84 changes: 84 additions & 0 deletions
84
libs/partners/mongodb/langchain_mongodb/chat_message_histories.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,84 @@ | ||
import json | ||
import logging | ||
from typing import List | ||
|
||
from langchain_core.chat_history import BaseChatMessageHistory | ||
from langchain_core.messages import ( | ||
BaseMessage, | ||
message_to_dict, | ||
messages_from_dict, | ||
) | ||
from pymongo import MongoClient, errors | ||
|
||
logger = logging.getLogger(__name__) | ||
|
||
DEFAULT_DBNAME = "chat_history" | ||
DEFAULT_COLLECTION_NAME = "message_store" | ||
|
||
|
||
class MongoDBChatMessageHistory(BaseChatMessageHistory): | ||
"""Chat message history that stores history in MongoDB. | ||
Args: | ||
connection_string: connection string to connect to MongoDB | ||
session_id: arbitrary key that is used to store the messages | ||
of a single chat session. | ||
database_name: name of the database to use | ||
collection_name: name of the collection to use | ||
""" | ||
|
||
def __init__( | ||
self, | ||
connection_string: str, | ||
session_id: str, | ||
database_name: str = DEFAULT_DBNAME, | ||
collection_name: str = DEFAULT_COLLECTION_NAME, | ||
): | ||
self.connection_string = connection_string | ||
self.session_id = session_id | ||
self.database_name = database_name | ||
self.collection_name = collection_name | ||
|
||
try: | ||
self.client: MongoClient = MongoClient(connection_string) | ||
except errors.ConnectionFailure as error: | ||
logger.error(error) | ||
|
||
self.db = self.client[database_name] | ||
self.collection = self.db[collection_name] | ||
self.collection.create_index("SessionId") | ||
|
||
@property | ||
def messages(self) -> List[BaseMessage]: # type: ignore | ||
"""Retrieve the messages from MongoDB""" | ||
try: | ||
cursor = self.collection.find({"SessionId": self.session_id}) | ||
except errors.OperationFailure as error: | ||
logger.error(error) | ||
|
||
if cursor: | ||
items = [json.loads(document["History"]) for document in cursor] | ||
else: | ||
items = [] | ||
|
||
messages = messages_from_dict(items) | ||
return messages | ||
|
||
def add_message(self, message: BaseMessage) -> None: | ||
"""Append the message to the record in MongoDB""" | ||
try: | ||
self.collection.insert_one( | ||
{ | ||
"SessionId": self.session_id, | ||
"History": json.dumps(message_to_dict(message)), | ||
} | ||
) | ||
except errors.WriteError as err: | ||
logger.error(err) | ||
|
||
def clear(self) -> None: | ||
"""Clear session memory from MongoDB""" | ||
try: | ||
self.collection.delete_many({"SessionId": self.session_id}) | ||
except errors.WriteError as err: | ||
logger.error(err) |
Large diffs are not rendered by default.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
37 changes: 37 additions & 0 deletions
37
libs/partners/mongodb/tests/integration_tests/test_chat_message_histories.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,37 @@ | ||
import json | ||
import os | ||
|
||
from langchain.memory import ConversationBufferMemory | ||
from langchain_core.messages import message_to_dict | ||
|
||
from langchain_mongodb.chat_message_histories import MongoDBChatMessageHistory | ||
|
||
# Replace these with your mongodb connection string | ||
connection_string = os.environ.get("MONGODB_ATLAS_URI", "") | ||
|
||
|
||
def test_memory_with_message_store() -> None: | ||
"""Test the memory with a message store.""" | ||
# setup MongoDB as a message store | ||
message_history = MongoDBChatMessageHistory( | ||
connection_string=connection_string, session_id="test-session" | ||
) | ||
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 Azure Cosmos DB, so the next test run won't pick it up | ||
memory.chat_memory.clear() | ||
|
||
assert memory.chat_memory.messages == [] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters