mirror of
https://github.com/hwchase17/langchain.git
synced 2025-07-16 09:48:04 +00:00
ccb6238de1
# Implemented appending arbitrary messages to the base chat message history, the in-memory and cosmos ones. <!-- Thank you for contributing to LangChain! Your PR will appear in our next release under the title you set. Please make sure it highlights your valuable contribution. Replace this with a description of the change, the issue it fixes (if applicable), and relevant context. List any dependencies required for this change. After you're done, someone will review your PR. They may suggest improvements. If no one reviews your PR within a few days, feel free to @-mention the same people again, as notifications can get lost. --> As discussed this is the alternative way instead of #4480, with a add_message method added that takes a BaseMessage as input, so that the user can control what is in the base message like kwargs. <!-- Remove if not applicable --> Fixes # (issue) ## Before submitting <!-- If you're adding a new integration, include an integration test and an example notebook showing its use! --> ## Who can review? Community members can review the PR once tests pass. Tag maintainers/contributors who might be interested: @hwchase17 --------- Co-authored-by: Harrison Chase <hw.chase.17@gmail.com>
91 lines
2.6 KiB
Python
91 lines
2.6 KiB
Python
import json
|
|
import logging
|
|
from typing import List
|
|
|
|
from langchain.schema import (
|
|
BaseChatMessageHistory,
|
|
BaseMessage,
|
|
_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_message(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)
|