Spaces:
Sleeping
Sleeping
File size: 1,751 Bytes
f202266 |
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 |
from typing import Dict, Optional
from langchain.memory import ConversationBufferMemory
class MemoryManager:
def __init__(self):
self.user_memories: Dict[str, Dict[str, ConversationBufferMemory]] = {}
def get_memory(self, user_id: str, context_id: Optional[str] = "general") -> ConversationBufferMemory:
"""Get or create a memory instance for a specific user and context.
Args:
user_id: The ID of the user
context_id: The context identifier (e.g., "general" for course chat, or "lecture_{id}" for lecture chat)
"""
if user_id not in self.user_memories:
self.user_memories[user_id] = {}
if context_id not in self.user_memories[user_id]:
self.user_memories[user_id][context_id] = ConversationBufferMemory(
memory_key="chat_history",
return_messages=True,
output_key="answer"
)
return self.user_memories[user_id][context_id]
def clear_memory(self, user_id: str, context_id: Optional[str] = None):
"""Clear memory for a user.
Args:
user_id: The ID of the user
context_id: If provided, only clear this specific context. Otherwise clear all contexts.
"""
if user_id in self.user_memories:
if context_id:
if context_id in self.user_memories[user_id]:
self.user_memories[user_id][context_id].clear()
else:
# Clear all contexts for this user
for mem in self.user_memories[user_id].values():
mem.clear()
# Global memory manager instance
memory_manager = MemoryManager() |