Update app.py

app.py

@@ -152,73 +152,85 @@ if __name__ == "__main__":
     )
     st.header("📖 Patent Chat: Google Patents Chat Demo")
 
-    # Fetch query parameters safely
-    query_params = st.query_params
-    default_patent_link = query_params.get("patent_link", "https://patents.google.com/patent/US8676427B1/en")
-
     # Input for Google Patent Link
-    patent_link = st.text_area("Enter Google Patent Link:", value=default_patent_link, height=100)
+    patent_link = st.text_area(
+        "Enter Google Patent Link:", 
+        value="https://patents.google.com/patent/US8676427B1/en", 
+        height=100
+    )
 
-    # Button to start processing
+    # Initialize session state
+    if "LOADED_PATENT" not in st.session_state:
+        st.session_state.LOADED_PATENT = None
+    if "pdf_preview" not in st.session_state:
+        st.session_state.pdf_preview = None
+    if "loaded_pdf_path" not in st.session_state:
+        st.session_state.loaded_pdf_path = None
+    if "chain" not in st.session_state:
+        st.session_state.chain = None
+    if "messages" not in st.session_state:
+        st.session_state.messages = [{"role": "assistant", "content": "Hello! How can I assist you with this patent?"}]
+
+    # Load and process patent
     if st.button("Load and Process Patent"):
         if not patent_link:
-            st.warning("Please enter a Google patent link to proceed.")
+            st.warning("Please enter a valid Google patent link.")
             st.stop()
 
         # Extract patent number
         patent_number = extract_patent_number(patent_link)
         if not patent_number:
-            st.error("Invalid patent link format. Please provide a valid Google patent link.")
+            st.error("Invalid patent link format.")
             st.stop()
 
         st.write(f"Patent number: **{patent_number}**")
 
-        # File download handling
+        # File handling
         pdf_path = os.path.join(tempfile.gettempdir(), f"{patent_number}.pdf")
-        if os.path.isfile(pdf_path):
-            st.write("✅ File already downloaded.")
-        else:
+        if not os.path.isfile(pdf_path):
             st.write("📥 Downloading patent file...")
             pdf_path = download_pdf(patent_number)
             st.write(f"✅ File downloaded: {pdf_path}")
-
-        # Generate and display PDF preview
-        st.write("🖼️ Generating PDF preview...")
-        preview_image_path = preview_pdf(pdf_path)
-
-        if preview_image_path:
-            st.image(preview_image_path, caption="First Page Preview", use_column_width=True)
         else:
-            st.warning("Failed to generate a preview for this PDF.")
-
-        # Load the document into the system
-        st.write("🔄 Loading document into the system...")
+            st.write("✅ File already downloaded.")
 
-        # Check if the patent is already loaded in session state
-        if "LOADED_PATENT" not in st.session_state or st.session_state.get("LOADED_PATENT") != patent_number:
+        # Generate preview if not already stored
+        if st.session_state.LOADED_PATENT != patent_number:
+            st.write("🖼️ Generating PDF preview...")
+            preview_image_path = preview_pdf(pdf_path)
+            if preview_image_path:
+                st.session_state.pdf_preview = preview_image_path
+            else:
+                st.warning("Failed to generate PDF preview.")
+                st.session_state.pdf_preview = None
+
+            # Load the chain and store metadata
             st.session_state.chain = load_chain(pdf_path)
-            st.session_state.LOADED_PATENT = patent_number  # Store the current patent number
+            st.session_state.LOADED_PATENT = patent_number
+            st.session_state.loaded_pdf_path = pdf_path
             st.session_state.messages = [{"role": "assistant", "content": "Hello! How can I assist you with this patent?"}]
-            st.success("🚀 Document successfully loaded! You can now start asking questions.")
+            st.success("🚀 Document successfully loaded!")
         else:
-            st.success("✅ This patent is already loaded. You can continue asking questions.")
+            st.info("✅ This patent is already loaded.")
 
-    # Initialize messages if not already done
-    if "messages" not in st.session_state:
-        st.session_state.messages = [{"role": "assistant", "content": "Hello! How can I assist you with this patent?"}]
+    # Display the PDF preview if available
+    if st.session_state.pdf_preview:
+        st.image(st.session_state.pdf_preview, caption="First Page Preview", use_container_width=True)
 
     # Display previous chat messages
     for message in st.session_state.messages:
         with st.chat_message(message["role"]):
             st.markdown(message["content"])
 
-    # User input and chatbot response
-    if "chain" in st.session_state:
+    # User input for questions
+    if st.session_state.chain:
         if user_input := st.chat_input("What is your question?"):
+            # User message
             st.session_state.messages.append({"role": "user", "content": user_input})
             with st.chat_message("user"):
                 st.markdown(user_input)
 
+            # Assistant response
             with st.chat_message("assistant"):
                 message_placeholder = st.empty()
                 full_response = ""
@@ -226,7 +238,7 @@ if __name__ == "__main__":
                 with st.spinner("Generating response..."):
                     try:
                         assistant_response = st.session_state.chain({"question": user_input})
-                        full_response = assistant_response["answer"]
+                        full_response = assistant_response.get("answer", "I couldn't process that question.")
                     except Exception as e:
                         full_response = f"An error occurred: {e}"