partners: add custom options to MongoDBChatMessageHistory (#22944)

**Description:** Adds options for configuring MongoDBChatMessageHistory
(no breaking changes):
- session_id_key: name of the field that stores the session id
- history_key: name of the field that stores the chat history
- create_index: whether to create an index on the session id field
- index_kwargs: additional keyword arguments to pass to the index
creation

**Discussion:**
https://github.com/langchain-ai/langchain/discussions/22918
**Twitter handle:** @userlerueda

---------

Co-authored-by: Jib <Jibzade@gmail.com>
Co-authored-by: Eugene Yurtsev <eugene@langchain.dev>
pull/20260/head
Luis Rueda 3 months ago committed by GitHub
parent 1e750f12f6
commit 168e9ed3a5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

@ -1,6 +1,6 @@
import json
import logging
from typing import List
from typing import Dict, List, Optional
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.messages import (
@ -14,6 +14,8 @@ logger = logging.getLogger(__name__)
DEFAULT_DBNAME = "chat_history"
DEFAULT_COLLECTION_NAME = "message_store"
DEFAULT_SESSION_ID_KEY = "SessionId"
DEFAULT_HISTORY_KEY = "History"
class MongoDBChatMessageHistory(BaseChatMessageHistory):
@ -25,6 +27,10 @@ class MongoDBChatMessageHistory(BaseChatMessageHistory):
of a single chat session.
database_name: name of the database to use
collection_name: name of the collection to use
session_id_key: name of the field that stores the session id
history_key: name of the field that stores the chat history
create_index: whether to create an index on the session id field
index_kwargs: additional keyword arguments to pass to the index creation
"""
def __init__(
@ -33,11 +39,18 @@ class MongoDBChatMessageHistory(BaseChatMessageHistory):
session_id: str,
database_name: str = DEFAULT_DBNAME,
collection_name: str = DEFAULT_COLLECTION_NAME,
*,
session_id_key: str = DEFAULT_SESSION_ID_KEY,
history_key: str = DEFAULT_HISTORY_KEY,
create_index: bool = True,
index_kwargs: Optional[Dict] = None,
):
self.connection_string = connection_string
self.session_id = session_id
self.database_name = database_name
self.collection_name = collection_name
self.session_id_key = session_id_key
self.history_key = history_key
try:
self.client: MongoClient = MongoClient(connection_string)
@ -46,18 +59,21 @@ class MongoDBChatMessageHistory(BaseChatMessageHistory):
self.db = self.client[database_name]
self.collection = self.db[collection_name]
self.collection.create_index("SessionId")
if create_index:
index_kwargs = index_kwargs or {}
self.collection.create_index(self.session_id_key, **index_kwargs)
@property
def messages(self) -> List[BaseMessage]: # type: ignore
"""Retrieve the messages from MongoDB"""
try:
cursor = self.collection.find({"SessionId": self.session_id})
cursor = self.collection.find({self.session_id_key: self.session_id})
except errors.OperationFailure as error:
logger.error(error)
if cursor:
items = [json.loads(document["History"]) for document in cursor]
items = [json.loads(document[self.history_key]) for document in cursor]
else:
items = []
@ -69,8 +85,8 @@ class MongoDBChatMessageHistory(BaseChatMessageHistory):
try:
self.collection.insert_one(
{
"SessionId": self.session_id,
"History": json.dumps(message_to_dict(message)),
self.session_id_key: self.session_id,
self.history_key: json.dumps(message_to_dict(message)),
}
)
except errors.WriteError as err:
@ -79,6 +95,6 @@ class MongoDBChatMessageHistory(BaseChatMessageHistory):
def clear(self) -> None:
"""Clear session memory from MongoDB"""
try:
self.collection.delete_many({"SessionId": self.session_id})
self.collection.delete_many({self.session_id_key: self.session_id})
except errors.WriteError as err:
logger.error(err)

@ -13,6 +13,8 @@ class PatchedMongoDBChatMessageHistory(MongoDBChatMessageHistory):
self.database_name = "test-database"
self.collection_name = "test-collection"
self.collection = MockCollection()
self.session_id_key = "SessionId"
self.history_key = "History"
def test_memory_with_message_store() -> None:

Loading…
Cancel
Save