| 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()) |
|
|
