File size: 2,785 Bytes
8c09efe 10be51f 11fcf8f 10be51f 8c09efe 10be51f fb3c0e4 10be51f 8c09efe 10be51f 8c09efe 10be51f 82e3417 10be51f 11fcf8f 10be51f |
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 66 67 68 |
import streamlit as st
from llama_index.llms.gemini import Gemini
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
import os
from llama_index.embeddings.gemini import GeminiEmbedding
#os.environ["GOOGLE_API_KEY"] = GOOGLE_API_KEY
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 = Gemini(
model="models/gemini-2.0-flash-exp",
temperature=1,
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.""",
)
Settings.embed_model = GeminiEmbedding(
model="models/embedding-001", embed_batch_size=100
)
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) |