import streamlit as st
from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_aws import ChatBedrock

llm = ChatBedrock(
    model="anthropic.claude-3-sonnet-20240229-v1:0",
    model_kwargs=dict(temperature=0),
    # other params...
)
store = {}  # memory is maintained outside the chain

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



chain = RunnableWithMessageHistory(llm, get_session_history)

# Streamlit app starts here
st.title("Chat with AI")
session_id = st.text_input("Enter your session ID:", "1")
user_input = st.text_area("You:", height=100)

# Initialize session state for messages if it doesn't exist
if 'messages' not in st.session_state:
    st.session_state.messages = []

if st.button("Send"):
    response = chain.invoke(user_input, config={"configurable": {"session_id": session_id}})
    
    # Display assistant response in chat format
    with st.chat_message("assistant"):
        st.markdown(response)
    
    # Add assistant response to chat history
    st.session_state.messages.append({"role": "assistant", "content": response})