forked from hwchase17/notion-qa
-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
65 lines (48 loc) · 1.91 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
"""Python file to serve as the frontend"""
import streamlit as st
from streamlit_chat import message
import faiss
from langchain import OpenAI
from langchain.chains import VectorDBQAWithSourcesChain
import pickle
from dotenv import load_dotenv
import os
import openai
from langchain.callbacks import get_openai_callback
import time
## ENV ##
load_dotenv()
openai.api_key = os.getenv("OPENAI_API_KEY")
# Load the LangChain.
index = faiss.read_index("docs.index")
with open("faiss_store.pkl", "rb") as f:
store = pickle.load(f)
store.index = index
chain = VectorDBQAWithSourcesChain.from_llm(llm=OpenAI(temperature=0), vectorstore=store)
# From here down is all the StreamLit UI.
st.set_page_config(page_title="Blendle Notion QA Bot", page_icon=":robot:")
st.header("Blendle Notion QA Bot")
if "generated" not in st.session_state:
st.session_state["generated"] = []
if "past" not in st.session_state:
st.session_state["past"] = []
def get_text():
input_text = st.text_input("You: ", "Hello, how are you?", key="input")
return input_text
user_input = get_text()
if user_input:
start_time = time.time()
result = chain({"question": user_input})
with get_openai_callback() as cb: # check openai sending
st.info('🧠 Processing the relevant documents with gpt...')
response = chain({"question": user_input})
output = f"Answer: {result['answer']}\nSources: {result['sources']}"
st.session_state.past.append(user_input)
st.session_state.generated.append(output)
with st.expander('💰 View the response token usage'):
st.code(cb)
st.text(f"Time taken: {time.time() - start_time} seconds")
if st.session_state["generated"]:
for i in range(len(st.session_state["generated"]) - 1, -1, -1):
message(st.session_state["generated"][i], key=str(i))
message(st.session_state["past"][i], is_user=True, key=str(i) + "_user")