Update test.py

test.py

@@ -1,3 +1,4 @@
+# to-do: Enable downloading multiple patent PDFs via corresponding links 
 import sys
 import os
 import re
@@ -61,7 +62,16 @@ def load_docs(document_path):
         )
         documents = loader.load()
         text_splitter = NLTKTextSplitter(chunk_size=1000)
-        return text_splitter.split_documents(documents)
+        split_docs = text_splitter.split_documents(documents)
+        
+        # Filter metadata to only include str, int, float, or bool
+        for doc in split_docs:
+            if hasattr(doc, "metadata") and isinstance(doc.metadata, dict):
+                doc.metadata = {
+                    k: v for k, v in doc.metadata.items()
+                    if isinstance(v, (str, int, float, bool))
+                }
+        return split_docs
     except Exception as e:
         st.error(f"Failed to load and process PDF: {e}")
         st.stop()
@@ -128,59 +138,71 @@ if __name__ == "__main__":
     )
     st.header("📖 Patent Chat: Google Patents Chat Demo")
 
-    patent_link = st.text_input("Enter Google Patent Link:", key="PATENT_LINK")
-
-    if not patent_link:
-        st.warning("Please enter a Google patent link to proceed.")
-        st.stop()
-
-    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.stop()
-
-    st.write(f"Patent number: **{patent_number}**")
-
-    pdf_path = os.path.join(tempfile.gettempdir(), f"{patent_number}.pdf")
-    if os.path.isfile(pdf_path):
-        st.write("✅ File already downloaded.")
-    else:
-        st.write("📥 Downloading patent file...")
-        pdf_path = download_pdf(patent_number)
-        st.write(f"✅ File downloaded: {pdf_path}")
-
-    st.write("🔄 Loading document into the system...")
-    chain = load_chain(pdf_path)
-    st.success("🚀 Document successfully loaded! You can now start asking questions.")
-
+    # 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)
+
+    # Button to start processing
+    if st.button("Load and Process Patent"):
+        if not patent_link:
+            st.warning("Please enter a Google patent link to proceed.")
+            st.stop()
+
+        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.stop()
+
+        st.write(f"Patent number: **{patent_number}**")
+
+        pdf_path = os.path.join(tempfile.gettempdir(), f"{patent_number}.pdf")
+        if os.path.isfile(pdf_path):
+            st.write("✅ File already downloaded.")
+        else:
+            st.write("📥 Downloading patent file...")
+            pdf_path = download_pdf(patent_number)
+            st.write(f"✅ File downloaded: {pdf_path}")
+
+        st.write("🔄 Loading document into the system...")
+
+        # Persist the chain in session state to prevent reloading
+        if "chain" not in st.session_state or st.session_state.get("loaded_file") != pdf_path:
+            st.session_state.chain = load_chain(pdf_path)
+            st.session_state.loaded_file = 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.")
+
+    # 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?"}
-        ]
+        st.session_state.messages = [{"role": "assistant", "content": "Hello! How can I assist you with this patent?"}]
 
+    # Display previous chat messages
     for message in st.session_state.messages:
         with st.chat_message(message["role"]):
             st.markdown(message["content"])
 
-    if user_input := st.chat_input("What is your question?"):
-        st.session_state.messages.append({"role": "user", "content": user_input})
-        with st.chat_message("user"):
-            st.markdown(user_input)
-
-        with st.chat_message("assistant"):
-            message_placeholder = st.empty()
-            full_response = ""
-
-        with st.spinner("Generating response..."):
-            try:
-                assistant_response = chain({"question": user_input})
-                for chunk in assistant_response["answer"].split():
-                    full_response += chunk + " "
-                    time.sleep(0.05)
-                    message_placeholder.markdown(full_response + "▌")
-            except Exception as e:
-                full_response = f"An error occurred: {e}"
-            finally:
-                message_placeholder.markdown(full_response)
+    if "chain" in st.session_state:
+        if user_input := st.chat_input("What is your question?"):
+            st.session_state.messages.append({"role": "user", "content": user_input})
+            with st.chat_message("user"):
+                st.markdown(user_input)
+
+            with st.chat_message("assistant"):
+                message_placeholder = st.empty()
+                full_response = ""
 
-        st.session_state.messages.append({"role": "assistant", "content": full_response})
+                with st.spinner("Generating response..."):
+                    try:
+                        assistant_response = st.session_state.chain({"question": user_input})
+                        full_response = assistant_response["answer"]
+                    except Exception as e:
+                        full_response = f"An error occurred: {e}"
+
+                message_placeholder.markdown(full_response)
+                st.session_state.messages.append({"role": "assistant", "content": full_response})
+    else:
+        st.info("Press the 'Load and Process Patent' button to start processing.")