Update app.py

app.py

@@ -20,7 +20,7 @@ if "id" not in st.session_state:
 session_id = st.session_state.id
 client = None
 
-# Initialize Cerebras LLM
+# Initialize Cerebras LLM (ensure it is available across the app)
 def load_llm():
     # Ensure you have the API Key set in your environment or via input
     api_key = os.getenv("CEREBRAS_API_KEY")
@@ -33,6 +33,10 @@ def load_llm():
         st.error("API Key is required.")
         return None
 
+# Load llm at the beginning of the session
+if "llm" not in st.session_state:
+    st.session_state.llm = load_llm()
+
 def reset_chat():
     st.session_state.messages = []
     st.session_state.context = None
@@ -64,7 +68,6 @@ with st.sidebar:
                 st.write("Indexing your document...")
 
                 if file_key not in st.session_state.get('file_cache', {}):
-
                     if os.path.exists(temp_dir):
                         reader = DoclingReader()
                         loader = SimpleDirectoryReader(
@@ -78,7 +81,7 @@ with st.sidebar:
                     docs = loader.load_data()
 
                     # setup llm & embedding model
-                    llm = load_llm()  # Load the Cerebras model
+                    llm = st.session_state.llm  # Load the Cerebras model from session state
                     if not llm:
                         st.stop()  # Stop execution if model initialization failed
                     embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-large-en-v1.5", trust_remote_code=True)
@@ -111,7 +114,7 @@ with st.sidebar:
                 else:
                     query_engine = st.session_state.file_cache[file_key]
 
-                # Inform the user that the file is processed and Display the PDF uploaded
+                # Inform the user that the file is processed and Display the Excel uploaded
                 st.success("Ready to Chat!")
                 display_excel(uploaded_file)
 
@@ -150,16 +153,13 @@ if prompt := st.chat_input("What's up?"):
         full_response = ""
 
         # Ensure llm is loaded
-        if 'file_cache' in st.session_state and len(st.session_state.file_cache) > 0:
-            query_engine = list(st.session_state.file_cache.values())[0]  # Get the first query engine
-
+        if st.session_state.llm:
             # Using Cerebras stream_chat for streaming response
             messages = [
                 ChatMessage(role="user", content=prompt)
             ]
 
-            response = query_engine.query(prompt)
-            st.write(response)  # Display raw query response for debugging
+            response = st.session_state.llm.stream_chat(messages)
             for r in response:
                 full_response += r.delta
                 message_placeholder.markdown(full_response + "▌")