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