Update app.py

app.py

@@ -1,81 +1,64 @@
-from langchain_community.chat_message_histories import StreamlitChatMessageHistory
-from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
-from langchain_core.runnables.history import RunnableWithMessageHistory
-from langchain_google_genai import GoogleGenerativeAIimport
-import os
-from dotenv import load_dotenv
-
-
 import streamlit as st
 
+from llama_index.llms.gemini import Gemini
+from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
+import os 
 
-load_dotenv()
-api_key=os.getenv("GOOGLE_API_KEY")
-# Configure l'API de Gemini
-#llm = GoogleGenerativeAI(model="models/text-bison-001", google_api_key=api_key)
-#genai.configure(api_key=api_key)
-
-
-st.set_page_config(page_title="StreamlitChatMessageHistory", page_icon="📖")
-st.title("📖 StreamlitChatMessageHistory")
-
-"""
-A basic example of using StreamlitChatMessageHistory to help LLMChain remember messages in a conversation.
-The messages are stored in Session State across re-runs automatically. You can view the contents of Session State
-in the expander below. View the
-[source code for this app](https://github.com/langchain-ai/streamlit-agent/blob/main/streamlit_agent/basic_memory.py).
-"""
-
-# Set up memory
-msgs = StreamlitChatMessageHistory(key="langchain_messages")
-if len(msgs.messages) == 0:
-    msgs.add_ai_message("How can I help you?")
-
-view_messages = st.expander("View the message contents in session state")
-
+#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="📃")
 
-# Set up the LangChain, passing in Message History
-
-prompt = ChatPromptTemplate.from_messages(
-    [
-        ("system", "You are an AI chatbot having a conversation with a human."),
-        MessagesPlaceholder(variable_name="history"),
-        ("human", "{question}"),
+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!",
+        }
     ]
-)
-
-chain = prompt | GoogleGenerativeAI(model="models/gemini-2.0-flash-exp", google_api_key=api_key)
-
-
-
-chain_with_history = RunnableWithMessageHistory(
-    chain,
-    lambda session_id: msgs,
-    input_messages_key="question",
-    history_messages_key="history",
-)
-
-# Render current messages from StreamlitChatMessageHistory
-for msg in msgs.messages:
-    st.chat_message(msg.type).write(msg.content)
-
-# If user inputs a new prompt, generate and draw a new response
-if prompt := st.chat_input():
-    st.chat_message("human").write(prompt)
-    # Note: new messages are saved to history automatically by Langchain during run
-    config = {"configurable": {"session_id": "any"}}
-    response = chain_with_history.invoke({"question": prompt}, config)
-    st.chat_message("ai").write(response.content)
-
-# Draw the messages at the end, so newly generated ones show up immediately
-with view_messages:
-    """
-    Message History initialized with:
-    ```python
-    msgs = StreamlitChatMessageHistory(key="langchain_messages")
-    ```
 
-    Contents of `st.session_state.langchain_messages`:
-    """
-    view_messages.json(st.session_state.langchain_messages)
+@st.cache_resource(show_spinner=False)
+def load_data():
+    reader = SimpleDirectoryReader(input_dir="./data", recursive=True)
+    docs = reader.load_data()
+    Settings.llm = Gemini(
+        model="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.""",
+    )
+    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)