Surendradjh's picture
Update app.py
0db1638 verified
import streamlit as st
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.output_parsers import StrOutputParser
from langchain_community.chat_message_histories import SQLChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory
import os
API_KEY = os.getenv("Google_api_key")
template = ChatPromptTemplate(
messages=[
("system", "You're a helpful data science AI chatbot. Answer only questions related to data science and what he told you within a 300-word limit."),
MessagesPlaceholder(variable_name="chat_history"),
("human", "{input}")
]
)
model = ChatGoogleGenerativeAI(api_key=API_KEY, model="gemini-1.5-pro")
output = StrOutputParser()
chain = template | model | output
def messages_history(session_id):
return SQLChatMessageHistory(session_id=session_id, connection="sqlite:///sqlite.db")
conversation_chain = RunnableWithMessageHistory(
chain, messages_history, input_message_key="input", history_messages_key="chat_history"
)
with st.sidebar:
st.title("🤖 AI Data Science Chatbot")
st.header("User Login")
user_id = st.text_input("Enter your User ID:", key="user_id_input")
# st.header(API_KEY)
if not user_id:
st.warning("Please enter a User ID to start chatting.")
st.stop()
if "last_user_id" not in st.session_state or st.session_state.last_user_id != user_id:
st.session_state.chat_history = []
st.session_state.last_user_id = user_id
chat_history = messages_history(user_id).messages
st.session_state.chat_history = [(msg.type, msg.content) for msg in chat_history]
st.write("Welcome! Start chatting below:")
# st.write(API_KEY)
for role, message in st.session_state.chat_history:
st.chat_message(role).write(message)
user_input = st.chat_input("Type your message...")
if user_input:
st.session_state.chat_history.append(("user", user_input))
st.chat_message("user").write(user_input)
config = {"configurable": {"session_id": user_id}}
input_prompt = {"input": user_input}
response = conversation_chain.invoke(input_prompt, config=config)
st.session_state.chat_history.append(("assistant", response))
st.chat_message("assistant").write(response)