Update app.py

app.py

@@ -110,7 +110,18 @@ if "documents" in st.session_state:
         "Chat with paper"
     ])
 
-    if st.button("🚀 Run Agent"):
+    # Handle Chat with paper separately
+    if task == "Chat with paper":
+        query = st.text_input("💬 Ask a question about the paper:")
+        if query and st.button("🚀 Ask Question"):
+            with st.spinner("Searching paper for answer..."):
+                retriever = st.session_state.vectorstore.as_retriever()
+                qa_chain = RetrievalQA.from_chain_type(llm=llm, retriever=retriever)
+                output = qa_chain.run(query)
+                st.session_state["last_agent_output"] = output
+    
+    # Handle other tasks
+    elif st.button("🚀 Run Agent"):
         with st.spinner("Running agents..."):
             docs = st.session_state.documents[:10]
             output = ""
@@ -143,13 +154,6 @@ if "documents" in st.session_state:
                 citation_chain = create_stuff_documents_chain(llm, citation_prompt)
                 output = run_chain(citation_chain, {"context": docs})
 
-            elif task == "Chat with paper":
-                query = st.text_input("💬 Ask a question about the paper:")
-                if query:
-                    retriever = st.session_state.vectorstore.as_retriever()
-                    qa_chain = RetrievalQA.from_chain_type(llm=llm, retriever=retriever)
-                    output = qa_chain.run(query)
-
             if output:
                 st.session_state["last_agent_output"] = output
 
@@ -184,4 +188,4 @@ if "last_agent_output" in st.session_state:
             })
 
             st.markdown(f"### 🌐 Translated Response ({user_language})")
-            st.write(translated)
+            st.write(translated)