persistent outputs

app.py

@@ -5,6 +5,14 @@ import base64
 from PIL import Image
 from io import BytesIO
 
+# Initialize session state variables
+if "uploaded_file_names" not in st.session_state:
+    st.session_state["uploaded_file_names"] = set()
+if "search_results" not in st.session_state:
+    st.session_state["search_results"] = []
+if "collections_list" not in st.session_state:
+    st.session_state["collections_list"] = ""
+
 
 def validate_api_key(api_key):
     # check if the api key is valid
@@ -16,7 +24,7 @@ def validate_api_key(api_key):
         return False
 
 
-# ask for api key from user in professional way
+# Ask for API key from user
 st.title("Colivara Demo")
 st.write("Welcome to the Colivara Demo.")
 st.subheader("API Key")
@@ -28,6 +36,7 @@ if not api_key or not validate_api_key(api_key):
 
 client = Colivara(base_url="https://api.colivara.com", api_key=api_key)
 
+# Create Collection Section
 st.subheader("Create Collection (Optional)")
 collection_name = st.text_input(
     "Collection Name",
@@ -41,58 +50,67 @@ if st.button("Create Collection"):
             client.create_collection(collection_name)
             st.success(f"Collection '{collection_name}' successfully created.")
 
-
+# List Collections Section
 st.subheader("List Collections")
 if st.button("List Collections"):
     with st.spinner("Fetching collections..."):
         collections = client.list_collections()
         collections_text = "\n".join([collection.name for collection in collections])
-        st.text_area(
-            "Existing Collections:",
-            value=collections_text,
-            height=200,
-            help="List of existing collections.",
-        )
-
+        st.session_state["collections_list"] = collections_text
+
+if st.session_state["collections_list"]:
+    st.text_area(
+        "Existing Collections:",
+        value=st.session_state["collections_list"],
+        height=200,
+        help="List of existing collections.",
+        key="collections_text_area",
+    )
+
+# Upsert Documents Section
 st.subheader("Upsert Documents")
 uploaded_files = st.file_uploader(
-    "Upload Docoments",
+    "Upload Documents",
     accept_multiple_files=True,
     help="Upload your documents here, most extensions are supported (pdf, docx, png, jpg, xlsx, etc.)",
 )
 selected_collection = st.text_input(
     "Collection Name", help="Enter the collection name to store the documents in."
 )
-if not uploaded_files or not selected_collection:
-    st.write("Please upload your documents and enter the collection name to proceed.")
-    st.stop()
-
-uploaded_file_names = set()
 
 if st.button("Upsert Files"):
-    new_files = [
-        file for file in uploaded_files if file.name not in uploaded_file_names
-    ]
-    if not new_files:
-        st.info("No new files to upsert. Please upload new files to proceed.")
+    if not uploaded_files or not selected_collection:
+        st.error(
+            "Please upload your documents and enter the collection name to proceed."
+        )
     else:
-        with st.spinner("Upserting Files..."):
-            progress_bar = st.progress(0)
-            total_files = len(new_files)
-            for idx, file in enumerate(new_files):
-                uploaded_file_names.add(file.name)
-                encoded_file = base64.b64encode(file.read()).decode("utf-8")
-                client.upsert_document(
-                    name=file.name,
-                    document_base64=encoded_file,
-                    collection_name=selected_collection,
-                    wait=True,
+        # Filter new files only
+        new_files = [
+            file
+            for file in uploaded_files
+            if file.name not in st.session_state["uploaded_file_names"]
+        ]
+        if not new_files:
+            st.info("No new files to upsert. Please upload new files to proceed.")
+        else:
+            with st.spinner("Upserting Files..."):
+                progress_bar = st.progress(0)
+                total_files = len(new_files)
+                for idx, file in enumerate(new_files):
+                    st.session_state["uploaded_file_names"].add(file.name)
+                    encoded_file = base64.b64encode(file.read()).decode("utf-8")
+                    client.upsert_document(
+                        name=file.name,
+                        document_base64=encoded_file,
+                        collection_name=selected_collection,
+                        wait=True,
+                    )
+                    progress_bar.progress((idx + 1) / total_files)
+                st.success(
+                    f"{len(new_files)} new files successfully upserted to the collection '{selected_collection}'."
                 )
-                progress_bar.progress((idx + 1) / total_files)
-            st.success(
-                f"{len(new_files)} new files successfully upserted to the collection '{selected_collection}'."
-            )
 
+# Search and Retrieve Documents Section
 st.subheader("Search and Retrieve Documents")
 query = st.text_input(
     "Search Query", help="Enter the search query to retrieve documents."
@@ -106,7 +124,6 @@ top_k = st.slider(
     help="Select the number of top documents to retrieve.",
 )
 
-
 if st.button("Search"):
     if not query:
         st.error("Please enter a search query to retrieve documents.")
@@ -115,22 +132,24 @@ if st.button("Search"):
             results = client.search(
                 query=query, collection_name=selected_collection, top_k=top_k
             )
-            results = results.results
-
-            if not results:
-                st.info("No documents found for the given query.")
-
-            else:
-                for idx, result in enumerate(results):
-                    img_base64 = result.img_base64
-                    page_number = result.page_number
-                    document_name = result.document_name
-                    img = Image.open(BytesIO(base64.b64decode(img_base64)))
-                    st.image(
-                        img, caption=f"Document: {document_name}, Page: {page_number}"
-                    )
-                    st.success("Search completed.")
-
+            st.session_state["search_results"] = results.results
+
+# Display Search Results
+if st.session_state["search_results"]:
+    st.write("Search Results:")
+    cols = st.columns(2)
+    for idx, result in enumerate(st.session_state["search_results"]):
+        img_base64 = result.img_base64
+        page_number = result.page_number
+        document_name = result.document_name
+        img = Image.open(BytesIO(base64.b64decode(img_base64)))
+        cols[idx % 2].image(
+            img,
+            caption=f"Document: {document_name}, Page: {page_number}",
+            use_container_width=False,
+            width=300,
+        )
+    st.success("Search completed.")
 
 st.markdown("----")
 st.markdown("Developed by Abdulhaleem from TJM Labs")