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 (ensure it is available across the app)
+# Initialize Cerebras LLM
 def load_llm():
     # Ensure you have the API Key set in your environment or via input
     api_key = os.getenv("CEREBRAS_API_KEY")
@@ -33,10 +33,6 @@ 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
@@ -68,6 +64,7 @@ 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(
@@ -81,7 +78,7 @@ with st.sidebar:
                     docs = loader.load_data()
 
                     # setup llm & embedding model
-                    llm = st.session_state.llm  # Load the Cerebras model from session state
+                    llm = load_llm()  # Load the Cerebras model
                     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)
@@ -114,7 +111,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 Excel uploaded
+                # Inform the user that the file is processed and Display the PDF uploaded
                 st.success("Ready to Chat!")
                 display_excel(uploaded_file)
 
@@ -125,7 +122,7 @@ with st.sidebar:
 col1, col2 = st.columns([6, 1])
 
 with col1:
-    st.header(f"RAG over Excel using DuckLink 🐥 &  Llama-3.3 70B")
+    st.header(f"RAG over Excel using Dockling 🐥 &  Llama-3.3 70B")
 
 with col2:
     st.button("Clear ↺", on_click=reset_chat)
@@ -153,13 +150,16 @@ if prompt := st.chat_input("What's up?"):
         full_response = ""
 
         # Ensure llm is loaded
-        if st.session_state.llm:
+        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
+
             # Using Cerebras stream_chat for streaming response
             messages = [
                 ChatMessage(role="user", content=prompt)
             ]
 
-            response = st.session_state.llm.stream_chat(messages)
+            response = query_engine.query(prompt)
+            st.write(response)  # Display raw query response for debugging
             for r in response:
                 full_response += r.delta
                 message_placeholder.markdown(full_response + "▌")
@@ -170,4 +170,4 @@ if prompt := st.chat_input("What's up?"):
             st.error("LLM model is not initialized correctly.")
 
     # Add assistant response to chat history
-    st.session_state.messages.append({"role": "assistant", "content": full_response})
+    st.session_state.messages.append({"role": "assistant", "content": full_response})