LangChain ConversationBufferMemory — deprecated, migrate to RunnableWithMessageHistory

# ============================================================
# BEFORE (deprecated)
# ============================================================
# from langchain.chains import LLMChain
# from langchain.memory import ConversationBufferMemory
# from langchain_openai import ChatOpenAI
# from langchain_core.prompts import PromptTemplate
#
# memory = ConversationBufferMemory(memory_key="history", return_messages=True)
# prompt = PromptTemplate.from_template("History:\n{history}\n\nUser: {input}\nAssistant:")
# chain = LLMChain(llm=ChatOpenAI(model="gpt-4o-mini"), prompt=prompt, memory=memory)
#
# chain.run(input="Hello")
# chain.run(input="What did I say?")
# ============================================================

# AFTER (current, LCEL-native)
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.chat_history import InMemoryChatMessageHistory

# 1) Build chain with a MessagesPlaceholder where history was
prompt = ChatPromptTemplate.from_messages([
    MessagesPlaceholder("history"),
    ("human", "{input}"),
])
chain = prompt | ChatOpenAI(model="gpt-4o-mini") | StrOutputParser()

# 2) Set up session-to-history factory
store = {}
def get_session_history(session_id: str):
    if session_id not in store:
        store[session_id] = InMemoryChatMessageHistory()
    return store[session_id]

# 3) Wrap
chain_with_history = RunnableWithMessageHistory(
    chain,
    get_session_history,
    input_messages_key="input",
    history_messages_key="history",
)

# 4) Invoke with session config
config = {"configurable": {"session_id": "user_42"}}
chain_with_history.invoke({"input": "Hello"}, config=config)
chain_with_history.invoke({"input": "What did I say?"}, config=config)
# > You said "Hello".

# BEFORE (deprecated)
# from langchain.memory import ConversationBufferWindowMemory
# memory = ConversationBufferWindowMemory(k=10, return_messages=True)

# AFTER (current)
from langchain_core.messages import trim_messages
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

llm = ChatOpenAI(model="gpt-4o-mini")

# Trim to last N messages OR last N tokens
# Count-based (equivalent to k=10):
trim_by_count = trim_messages(
    max_tokens=10,          # despite the name, this is the count when token_counter=len
    strategy="last",
    token_counter=len,       # use message count instead of tokens
    include_system=True,
    allow_partial=False,
)

# Token-based (better for varying message sizes):
trim_by_tokens = trim_messages(
    max_tokens=4000,
    strategy="last",
    token_counter=llm,        # use the LLM's tokenizer
    include_system=True,
    allow_partial=False,
    start_on="human",
)

# Wire the trimmer INTO the chain — history flows through it before the prompt
prompt = ChatPromptTemplate.from_messages([
    ("system", "You are helpful."),
    MessagesPlaceholder("history"),
    ("human", "{input}"),
])

chain = (
    RunnablePassthrough.assign(history=lambda x: trim_by_tokens.invoke(x["history"]))
    | prompt
    | llm
    | StrOutputParser()
)

# Wrap with RunnableWithMessageHistory as before — the trim happens BEFORE the prompt sees it
chain_with_history = RunnableWithMessageHistory(chain, get_session_history,
                                                input_messages_key="input",
                                                history_messages_key="history")

# BEFORE (deprecated)
# from langchain.memory import ConversationSummaryMemory
# memory = ConversationSummaryMemory(llm=ChatOpenAI(), memory_key="history")

# AFTER (current) — explicit summarize-then-inject pattern
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage

summariser_llm = ChatOpenAI(model="gpt-4o-mini")
main_llm = ChatOpenAI(model="gpt-4o")

# Summarization sub-chain
summarise_prompt = ChatPromptTemplate.from_messages([
    ("system", "Summarise this conversation in 2-3 sentences, preserving key facts."),
    MessagesPlaceholder("messages"),
])
summarise_chain = summarise_prompt | summariser_llm | StrOutputParser()

def summarise_if_long(messages, threshold=20):
    if len(messages) <= threshold:
        return messages
    # Summarise everything except last N (keep recent context detailed)
    keep = messages[-6:]
    to_summarise = messages[:-6]
    summary = summarise_chain.invoke({"messages": to_summarise})
    return [SystemMessage(content=f"Earlier conversation summary: {summary}")] + keep

# Main chain that summarises history before using it
main_prompt = ChatPromptTemplate.from_messages([
    ("system", "You are helpful."),
    MessagesPlaceholder("history"),
    ("human", "{input}"),
])

chain = (
    RunnablePassthrough.assign(history=lambda x: summarise_if_long(x["history"]))
    | main_prompt
    | main_llm
    | StrOutputParser()
)

chain_with_history = RunnableWithMessageHistory(chain, get_session_history,
                                                input_messages_key="input",
                                                history_messages_key="history")