|
import streamlit as st |
|
from langchain.agents import initialize_agent, AgentType, Tool |
|
from langchain.chains import LLMMathChain, RetrievalQA |
|
from langchain.chat_models import ChatOpenAI |
|
from langchain.embeddings import OpenAIEmbeddings |
|
from langchain.text_splitter import CharacterTextSplitter |
|
from langchain.utilities import SQLDatabase |
|
from langchain.vectorstores import FAISS |
|
from langchain_experimental.sql import SQLDatabaseChain |
|
from langchain.document_loaders import PyPDFLoader |
|
from langchain.schema import ChatMessage |
|
from langchain.callbacks.manager import CallbackManager |
|
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler |
|
|
|
|
|
|
|
USERS = { |
|
"johndoe": {"password": "Streamlit@12", "username": "John Doe"}, |
|
"janesmith": {"password": "Streamlit@12", "username": "Jane Smith"} |
|
} |
|
|
|
|
|
llm = ChatOpenAI(temperature=0, model="gpt-3.5-turbo-0613") |
|
llm_math_chain = LLMMathChain.from_llm(llm=llm, verbose=True) |
|
db = SQLDatabase.from_uri("sqlite:///utilities_monthly.db") |
|
db_chain = SQLDatabaseChain.from_llm(llm, db, verbose=True) |
|
pdf_loader = PyPDFLoader("sewa.pdf") |
|
pages = pdf_loader.load_and_split() |
|
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0) |
|
docs = text_splitter.split_documents(pages) |
|
embeddings = OpenAIEmbeddings() |
|
retriever = FAISS.from_documents(docs, embeddings).as_retriever() |
|
|
|
tools = [ |
|
Tool( |
|
name="FAQ", |
|
func=RetrievalQA.from_chain_type(llm=llm, retriever=retriever), |
|
description="Useful for answering generic questions related to SEWA, bill, usage, etc." |
|
), |
|
Tool( |
|
name="Electricity-DB", |
|
func=db_chain.run, |
|
description="Useful for answering questions about Customer Energy Usage.Plase use this username given in the prompt Input should be in the form of a question containing full context" |
|
) |
|
] |
|
callback_manager = CallbackManager([StreamingStdOutCallbackHandler()]) |
|
agent_executor = initialize_agent(tools, llm, agent=AgentType.OPENAI_FUNCTIONS, verbose=True, |
|
callback_manager=callback_manager) |
|
|
|
|
|
st.markdown(""" |
|
<style> |
|
.stApp > div > div > div > div:first-child { |
|
background-color: blue; |
|
} |
|
</style> |
|
""", unsafe_allow_html=True) |
|
|
|
|
|
|
|
with st.sidebar: |
|
if "logged_in_user" not in st.session_state: |
|
st.session_state["logged_in_user"] = None |
|
|
|
if st.session_state["logged_in_user"]: |
|
user_data = st.session_state["logged_in_user"] |
|
st.title(f"Welcome {user_data['username']}!") |
|
|
|
st.write("You can ask about energy usage, bill, due date, etc..") |
|
|
|
|
|
if st.button("Log Out"): |
|
st.session_state["logged_in_user"] = None |
|
st.session_state["messages"] = [ChatMessage(role="assistant", content="How can I help you?")] |
|
st.experimental_rerun() |
|
else: |
|
st.title("SEWA - At your service") |
|
user_id = st.text_input("User ID") |
|
password = st.text_input("Password", type="password") |
|
if st.button("Login"): |
|
if user_id in USERS and USERS[user_id]['password'] == password: |
|
st.session_state["logged_in_user"] = USERS[user_id] |
|
st.success(f"Logged in as {USERS[user_id]['username']}!") |
|
st.experimental_rerun() |
|
else: |
|
st.error("Invalid user ID or password") |
|
|
|
if "messages" not in st.session_state: |
|
st.session_state["messages"] = [ChatMessage(role="assistant", content="How can I help you?")] |
|
|
|
for msg in st.session_state.messages: |
|
st.chat_message(msg.role).write(msg.content) |
|
|
|
if prompt := st.chat_input(): |
|
|
|
|
|
|
|
st.session_state.messages.append(ChatMessage(role="user", content=prompt)) |
|
st.chat_message("user").write(prompt) |
|
if st.session_state["logged_in_user"]: |
|
user_name = st.session_state["logged_in_user"]["username"] |
|
prompt = prompt.replace("my", user_name) |
|
response = agent_executor.invoke({"input": prompt}) |
|
st.session_state.messages.append(ChatMessage(role="assistant", content=response['output'])) |
|
st.chat_message("assistant").write(response['output']) |
|
|