forked from Archives/langchain
mongodb support for chat history (#4266)
parent
f510940bde
commit
8a338412fa
@ -0,0 +1,98 @@
|
||||
import json
|
||||
import logging
|
||||
from typing import List
|
||||
|
||||
from langchain.schema import (
|
||||
AIMessage,
|
||||
BaseChatMessageHistory,
|
||||
BaseMessage,
|
||||
HumanMessage,
|
||||
_message_to_dict,
|
||||
messages_from_dict,
|
||||
)
|
||||
|
||||
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,
|
||||
):
|
||||
from pymongo import MongoClient, errors
|
||||
|
||||
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]
|
||||
|
||||
@property
|
||||
def messages(self) -> List[BaseMessage]: # type: ignore
|
||||
"""Retrieve the messages from MongoDB"""
|
||||
from pymongo import errors
|
||||
|
||||
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_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:
|
||||
"""Append the message to the record in MongoDB"""
|
||||
from pymongo import errors
|
||||
|
||||
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"""
|
||||
from pymongo import errors
|
||||
|
||||
try:
|
||||
self.collection.delete_many({"SessionId": self.session_id})
|
||||
except errors.WriteError as err:
|
||||
logger.error(err)
|
@ -0,0 +1,36 @@
|
||||
import json
|
||||
import os
|
||||
|
||||
from langchain.memory import ConversationBufferMemory
|
||||
from langchain.memory.chat_message_histories import MongoDBChatMessageHistory
|
||||
from langchain.schema import _message_to_dict
|
||||
|
||||
# Replace these with your mongodb connection string
|
||||
connection_string = os.environ["MONGODB_CONNECTION_STRING"]
|
||||
|
||||
|
||||
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 == []
|
Loading…
Reference in New Issue