Spaces:
Running
Running
| import uuid | |
| from pymongo import MongoClient | |
| from langchain_mongodb.chat_message_histories import MongoDBChatMessageHistory | |
| class ChatManagement: | |
| def __init__(self, cluster_url, database_name, collection_name): | |
| """ | |
| Initializes the ChatManagement class with MongoDB connection details. | |
| Args: | |
| cluster_url (str): MongoDB cluster URL. | |
| database_name (str): Name of the database. | |
| collection_name (str): Name of the collection. | |
| """ | |
| self.connection_string = cluster_url | |
| self.database_name = database_name | |
| self.collection_name = collection_name | |
| self.chat_sessions = {} # Dictionary to store chat history objects for each session | |
| def create_new_chat(self): | |
| """ | |
| Creates a new chat session by initializing a MongoDBChatMessageHistory object. | |
| Returns: | |
| str: The unique chat ID. | |
| """ | |
| # Generate a unique chat ID | |
| chat_id = str(uuid.uuid4()) | |
| # Initialize MongoDBChatMessageHistory for the chat session | |
| chat_message_history = MongoDBChatMessageHistory( | |
| session_id=chat_id, | |
| connection_string=self.connection_string, | |
| database_name=self.database_name, | |
| collection_name=self.collection_name, | |
| ) | |
| # Store the chat_message_history object in the session dictionary | |
| self.chat_sessions[chat_id] = chat_message_history | |
| return chat_id | |
| def get_chat_history(self, chat_id): | |
| """ | |
| Retrieves the MongoDBChatMessageHistory object for a given chat session by its chat ID. | |
| Args: | |
| chat_id (str): The unique ID of the chat session. | |
| Returns: | |
| MongoDBChatMessageHistory or None: The chat history object of the chat session, or None if not found. | |
| """ | |
| # Check if the chat session is already in memory | |
| if chat_id in self.chat_sessions: | |
| return self.chat_sessions[chat_id] | |
| # If not in memory, try to fetch from the database | |
| chat_message_history = MongoDBChatMessageHistory( | |
| session_id=chat_id, | |
| connection_string=self.connection_string, | |
| database_name=self.database_name, | |
| collection_name=self.collection_name, | |
| ) | |
| if chat_message_history.messages: # Check if the session exists in the database | |
| self.chat_sessions[chat_id] = chat_message_history | |
| return chat_message_history | |
| return None # Chat session not found | |
| def initialize_chat_history(self, chat_id): | |
| """ | |
| Initializes a new chat history for the given chat ID if it does not already exist. | |
| Args: | |
| chat_id (str): The unique ID of the chat session. | |
| Returns: | |
| MongoDBChatMessageHistory: The initialized chat history object. | |
| """ | |
| # If the chat history already exists, return it | |
| if chat_id in self.chat_sessions: | |
| return self.chat_sessions[chat_id] | |
| # Otherwise, create a new chat history | |
| chat_message_history = MongoDBChatMessageHistory( | |
| session_id=chat_id, | |
| connection_string=self.connection_string, | |
| database_name=self.database_name, | |
| collection_name=self.collection_name, | |
| ) | |
| # Save the new chat session to the session dictionary | |
| self.chat_sessions[chat_id] = chat_message_history | |
| return chat_message_history | |