langchain_chatbot_with_in_memory_session_history.py

1from langchain_openai import ChatOpenAI
2from langchain_core.messages import HumanMessage, AIMessage
3from langchain_community.chat_message_histories import ChatMessageHistory
4from langchain_core.chat_history import BaseChatMessageHistory
5from langchain_core.runnables.history import RunnableWithMessageHistory
6
7# Initialize the model
8model = ChatOpenAI(model="gpt-3.5-turbo")
9
10# Dictionary to store session histories
11store = {}
12
13def get_session_history(session_id: str) -> BaseChatMessageHistory:
14    if session_id not in store:
15        store[session_id] = ChatMessageHistory()
16    return store[session_id]
17
18# Wrap the model with message history logic
19with_message_history = RunnableWithMessageHistory(model, get_session_history)
20
21# Configuration for the session
22config = {"configurable": {"session_id": "abc2"}}
23
24# First interaction
25response = with_message_history.invoke(
26    [HumanMessage(content="Hi! I'm Bob")],
27    config=config,
28)
29print(f"AI: {response.content}")
30
31# Second interaction (testing memory)
32response = with_message_history.invoke(
33    [HumanMessage(content="What's my name?")],
34    config=config,
35)
36print(f"AI: {response.content}")