-
Notifications
You must be signed in to change notification settings - Fork 315
/
streamlit_app.py
61 lines (52 loc) · 2.5 KB
/
streamlit_app.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
import streamlit as st
import openai
from llama_index.llms.openai import OpenAI
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
st.set_page_config(page_title="Chat with the Streamlit docs, powered by LlamaIndex", page_icon="🦙", layout="centered", initial_sidebar_state="auto", menu_items=None)
openai.api_key = st.secrets.openai_key
st.title("Chat with the Streamlit docs, powered by LlamaIndex 💬🦙")
st.info("Check out the full tutorial to build this app in our [blog post](https://blog.streamlit.io/build-a-chatbot-with-custom-data-sources-powered-by-llamaindex/)", icon="📃")
if "messages" not in st.session_state.keys(): # Initialize the chat messages history
st.session_state.messages = [
{
"role": "assistant",
"content": "Ask me a question about Streamlit's open-source Python library!",
}
]
@st.cache_resource(show_spinner=False)
def load_data():
reader = SimpleDirectoryReader(input_dir="./data", recursive=True)
docs = reader.load_data()
Settings.llm = OpenAI(
model="gpt-3.5-turbo",
temperature=0.2,
system_prompt="""You are an expert on
the Streamlit Python library and your
job is to answer technical questions.
Assume that all questions are related
to the Streamlit Python library. Keep
your answers technical and based on
facts – do not hallucinate features.""",
)
index = VectorStoreIndex.from_documents(docs)
return index
index = load_data()
if "chat_engine" not in st.session_state.keys(): # Initialize the chat engine
st.session_state.chat_engine = index.as_chat_engine(
chat_mode="condense_question", verbose=True, streaming=True
)
if prompt := st.chat_input(
"Ask a question"
): # Prompt for user input and save to chat history
st.session_state.messages.append({"role": "user", "content": prompt})
for message in st.session_state.messages: # Write message history to UI
with st.chat_message(message["role"]):
st.write(message["content"])
# If last message is not from assistant, generate a new response
if st.session_state.messages[-1]["role"] != "assistant":
with st.chat_message("assistant"):
response_stream = st.session_state.chat_engine.stream_chat(prompt)
st.write_stream(response_stream.response_gen)
message = {"role": "assistant", "content": response_stream.response}
# Add response to message history
st.session_state.messages.append(message)