bot_with_temporary_memory.py

import streamlit as st
import os
from dotenv import load_dotenv
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory

# Load environment variables
load_dotenv()

# Initialize store for chat history
store = {}

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

# Initialize LLM and chat chain
@st.cache_resource
def init_chat_chain():
    llm = ChatGroq(
        model="gemma2-9b-it",
        groq_api_key=os.getenv("GROQ_API_KEY")
    )
    
    prompt = ChatPromptTemplate.from_messages([
        ("system", "You are a helpful assistant please answer the question."),
        ("human", "{input}")
    ])
    
    chain = prompt | llm
    
    return RunnableWithMessageHistory(
        chain,
        get_session_history,
    )

def main():
    st.title("💬 Chatbot with Memory")
    
    # Initialize session state for messages
    if "messages" not in st.session_state:
        st.session_state.messages = []
        st.session_state.session_id = str(hash(str(st.session_state)))
    
    # Initialize chat chain
    chat_chain = init_chat_chain()
    
    # Display chat messages
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])
    
    # Chat input
    if prompt := st.chat_input("What would you like to know?"):
        # Add user message to chat history
        st.session_state.messages.append({"role": "user", "content": prompt})
        
        # Display user message
        with st.chat_message("user"):
            st.markdown(prompt)
        
        # Get bot response
        with st.chat_message("assistant"):
            config = {"configurable": {"session_id": st.session_state.session_id}}
            with st.spinner("Thinking..."):
                response = chat_chain.invoke(
                    {"input": prompt},
                    config=config,
                )
            
            # Display bot response
            st.markdown(response.content)
            
            # Add bot response to chat history
            st.session_state.messages.append(
                {"role": "assistant", "content": response.content}
            )

if __name__ == "__main__":
    main()