| from langchain_community.chat_message_histories import ChatMessageHistory | |
| from langchain_core.chat_history import BaseChatMessageHistory | |
| from langchain_core.runnables.history import RunnableWithMessageHistory | |
| class MemoryManager: | |
| """ | |
| Manages chat memory for the chatbot using LangChain's built-in memory features. | |
| """ | |
| def __init__(self): | |
| self.store = {} | |
| def get_session_history(self, session_id: str) -> BaseChatMessageHistory: | |
| """ | |
| Get or create a chat message history for a given session ID. | |
| :param session_id: Unique identifier for the chat session | |
| :return: ChatMessageHistory instance for the session | |
| """ | |
| if session_id not in self.store: | |
| self.store[session_id] = ChatMessageHistory() | |
| return self.store[session_id] | |
| def create_memory_enabled_model(self, model, session_id: str = "default"): | |
| """ | |
| Create a model with message history enabled. | |
| :param model: The LLM model to wrap with memory | |
| :param session_id: Session ID for this conversation | |
| :return: Model with message history | |
| """ | |
| with_message_history = RunnableWithMessageHistory( | |
| model, | |
| self.get_session_history | |
| ) | |
| config = {"configurable": {"session_id": session_id}} | |
| return with_message_history, config | |
| def clear_session(self, session_id: str): | |
| """ | |
| Clear the message history for a specific session. | |
| :param session_id: Session ID to clear | |
| """ | |
| if session_id in self.store: | |
| del self.store[session_id] | |
| def get_all_sessions(self): | |
| """ | |
| Get all active session IDs. | |
| :return: List of session IDs | |
| """ | |
| return list(self.store.keys()) | |