core/memory.py

from langchain.memory import ConversationBufferWindowMemory
from typing import Dict

conversation_memories: Dict[str, ConversationBufferWindowMemory] = {}

def get_memory(session_id: str, max_history: int = 10) -> ConversationBufferWindowMemory:
    """Get or create conversation memory for a session."""
    if session_id not in conversation_memories:
        conversation_memories[session_id] = ConversationBufferWindowMemory(
            k=max_history,
            return_messages=True,
            input_key="input",
            output_key="output"
        )
    return conversation_memories[session_id]