convert colivara demo only search feature

app.py

@@ -4,146 +4,104 @@ from colivara_py import Colivara
 import base64
 from PIL import Image
 from io import BytesIO
+import os
+
+API_KEY = os.getenv("COLIVARA_API_KEY")
 
-# 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
-    try:
-        client = Colivara(base_url="https://api.colivara.com", api_key=api_key)
-        client.list_collections()
-        return True
-    except Exception:
-        return False
+client = Colivara(base_url="https://api.colivara.com", api_key=API_KEY)
+# Add badges at the top of the page
+st.markdown(
+    """
+    <div style="display: flex; gap: 10px; margin-bottom: 20px;">
+        <a href="https://discord.gg/DtGRxWuj8y" target="_blank">
+            <img src="https://dcbadge.limes.pink/api/server/https://discord.gg/DtGRxWuj8y" alt="Discord">
+        </a>
+        <a href="https://www.colivara.com" target="_blank">
+            <img src="https://img.shields.io/badge/Website-0078D7?style=for-the-badge&logo=internetexplorer&logoColor=white" alt="Website">
+        </a>
+        <a href="https://docs.colivara.com" target="_blank">
+            <img src="https://img.shields.io/badge/Docs-217346?style=for-the-badge&logo=readthedocs&logoColor=white" alt="Docs">
+        </a>
+        <a href="https://github.com/tjmlabs/ColiVara" target="_blank">
+            <img src="https://img.shields.io/badge/GitHub-181717?style=for-the-badge&logo=github&logoColor=white" alt="GitHub">
+        </a>
+    </div>
+    """,
+    unsafe_allow_html=True,
+)
 
 
 # Ask for API key from user
 st.title("Colivara Demo")
-st.write("Welcome to the Colivara Demo.")
-st.subheader("API Key")
-api_key = st.text_input("API Key", type="password")
-if not api_key or not validate_api_key(api_key):
-    st.write("Please enter your API key to proceed.")
-    st.markdown("If you don't have one, get yours [here](https://colivara.com).")
-    st.stop()
 
-client = Colivara(base_url="https://api.colivara.com", api_key=api_key)
-st.divider()
-st.subheader("Features Overview")
-st.write("Each section below is independent, and you can use them separately.")
+st.markdown(
+    """
+    ### Welcome to the Colivara Search Demo
+    This demo allows you to search through a sample NVIDIA quarterly report (Form 10-Q).
 
+    - Enter a search query in the input box below.
+    - Retrieve the most relevant pages from the document.
 
-st.divider()
-# Create Collection Section
-st.subheader("Create Collection")
-collection_name = st.text_input(
-    "Collection Name",
-    help="Enter the name of the collection to create",
-)
-if st.button("Create Collection"):
-    if not collection_name:
-        st.error("Please enter a collection name to create the collection.")
-    else:
-        with st.spinner("Creating collection..."):
-            client.create_collection(collection_name)
-            st.success(f"Collection '{collection_name}' successfully created.")
-st.divider()
-# 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.session_state["collections_list"] = collections_text
+    **Example Queries:**
+    ```
+    1. What were Nvidia's revenues for the most recent quarter?
+    2. What are the main factors that are driving Nvidia's growth?
+    3. What are the risks and challenges that Nvidia is facing?
+    ```
 
-if st.session_state["collections_list"]:
-    st.code(st.session_state["collections_list"])
-st.divider()
-# Upsert Documents Section
-st.subheader("Upsert Documents")
-uploaded_files = st.file_uploader(
-    "Upload Documents",
-    accept_multiple_files=True,
-    help="Upload your documents here, most extensions are supported (pdf, docx, png, jpg, xlsx, etc.)",
+    Try it out and see how the AI retrieves the information you need efficiently.
+    """
 )
-selected_collection = st.text_input(
-    "Collection Name", help="Enter the collection name to store the documents in."
+
+# Display placeholder for document details with a badge for the PDF
+st.markdown(
+    """
+    #### Document Overview:
+    **Title:** NVIDIA Corporation Quarterly Report (Form 10-Q, October 27, 2024)  
+    **Type:** PDF  
+    **Pages:** 50+  
+    <a href="https://d18rn0p25nwr6d.cloudfront.net/CIK-0001045810/ed2a395c-5e9b-4411-8b4a-a718d192155a.pdf" target="_blank">
+        <img src="https://img.shields.io/badge/View%20PDF-EF8D21?style=for-the-badge&logo=adobeacrobatreader&logoColor=white" alt="View PDF">
+    </a>
+    """,
+    unsafe_allow_html=True,
 )
 
-if st.button("Upsert Files"):
-    if not uploaded_files or not selected_collection:
-        st.error(
-            "Please upload your documents and enter the collection name to proceed."
-        )
-    else:
-        # 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}'."
-                )
-st.divider()
-# Search and Retrieve Documents Section
+# Search and retrieve documents
 st.subheader("Search and Retrieve Documents")
 
-# Add collection name input for search
-search_collection_name = st.text_input(
-    "Collection Name for Search", help="Enter the collection name to search in."
-)
-
+# User inputs for search
 query = st.text_input(
-    "Search Query", help="Enter the search query to retrieve documents."
+    "Enter your search query:", help="Type a query to search through the NVIDIA report."
 )
-
 top_k = st.slider(
     "Number of Top Documents to Retrieve",
     min_value=1,
     max_value=10,
     value=3,
-    help="Select the number of top documents to retrieve.",
+    help="Select the number of top pages to retrieve.",
 )
 
+collection_name = "Demo-Final"
 if st.button("Search"):
-    if not search_collection_name:
-        st.error("Please enter the collection name for the search.")
-    elif not query:
-        st.error("Please enter a search query to retrieve documents.")
+    if not query:
+        st.error("Please enter a search query to retrieve results.")
     else:
         with st.spinner("Searching..."):
-            results = client.search(
-                query=query, collection_name=search_collection_name, top_k=top_k
-            )
-            st.session_state["search_results"] = results.results
+            try:
+                results = client.search(
+                    query=query, collection_name=collection_name, top_k=top_k
+                )
+                st.session_state["search_results"] = results.results
+            except Exception as e:
+                st.error(f"Error during search: {str(e)}")
 
-# Display Search Results
-if st.session_state["search_results"]:
-    st.write("Search Results:")
+# Display search results
+if "search_results" in st.session_state and 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
@@ -154,7 +112,7 @@ if st.session_state["search_results"]:
             img,
             caption=f"Document: {document_name}, Page: {page_number}",
         )
-    st.success("Search completed.")
+    st.success("Search completed successfully.")
 
 st.markdown("----")
-st.markdown("Developed by Abdulhaleem from TJM Labs")
+st.markdown("Developed by Abdulhaleem from TJM Labs.")