"""runnable_with_message_history_demo.py — RunnableWithMessageHistory + InMemoryChatMessageHistory"""

from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory

load_dotenv()

model = ChatOpenAI(model="gpt-4o-mini", temperature=0)

prompt = ChatPromptTemplate.from_messages([
    ("system", "Keep each reply to one short sentence."),
    MessagesPlaceholder(variable_name="history"),
    ("human", "{input}"),
])

chain = prompt | model

store = {}

def get_session_history(session_id: str):
    if session_id not in store:
        store[session_id] = InMemoryChatMessageHistory()
    return store[session_id]

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

config = {"configurable": {"session_id": "demo-1"}}

turns = [
    "What does the HTML <a> tag do?",
    "What attribute opens the link in a new tab?",
]

for question in turns:
    reply = chain_with_history.invoke({"input": question}, config=config)
    print(f"Human: {question}")
    print(f"AI:   {reply.content}\n")