File size: 1,851 Bytes
0cbebee |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
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())
|