I implemented this use case from langchain to add chat history to my RAG to contextualize questions that may be incomplete. Everything works fine but now I would like to use only the last x messages instead of the whole chat history. How can I do that?

Currently the code I’m using is the same from the tutorial:

store = {}


def get_session_history(session_id: str) -> BaseChatMessageHistory:
    if session_id not in store:
        store[session_id] = ChatMessageHistory()
    return store[session_id]


conversational_rag_chain = RunnableWithMessageHistory(
    rag_chain,
    get_session_history,
    input_messages_key="input",
    history_messages_key="chat_history",
    output_messages_key="answer",
)

I saw that there is a class called ConversationBufferWindowMemory that might be what I’m looking for but I can’t wrap my head around on how to swap ChatMessageHistory from the code above with ConversationBufferWindowMemory, if possible.

This is what I tried:

### Statefully manage chat history
store = {}

def get_session_history(session_id: str):
    if session_id not in store:
        store[session_id] = ConversationBufferWindowMemory(memory_key="chat_history", k=2)
    return store[session_id]


conversational_rag_chain = ConversationChain(
    rag_chain,
    get_session_history,
    verbose=True, 
    memory="chat_history"
)

Obiviously this was unsuccessul since I’m swapping a runnable with a chain. The error that I get is:

TypeError: _init_() takes 1 positional argument but 3 were given

I’m not very familiar with langchain in general since I just started working with it. Is it actually possible to do what I’m trying to do? or maybe is there some different approaches that are better suited for my problem?

Thanks in advance.