app.py

# import libraries
import streamlit as st
from colivara_py import Colivara
import base64
from PIL import Image
from io import BytesIO
import os

API_KEY = os.getenv("COLIVARA_API_KEY")

if "search_results" not in st.session_state:
    st.session_state["search_results"] = []

client = Colivara(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://img.shields.io/badge/Join%20Discord-7289DA?style=for-the-badge&logo=discord&logoColor=white" 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.markdown(
    """
    ### Welcome to the Colivara Search Demo
    This demo allows you to search through the Competitive Programmer's Handbook by Antti Laaksonen.
    We included 12 chapters from the handbook in this demo to help you understand how Colivara works.
    **Chapters Covered:**
    `Introduction` `Time Complexity` `Sorting` `Data Structures` `Complete Search` `Greedy Algorithms` `Dynamic Programming` `Amortized Analysis` `Range Queries` `Bit Manipulation` `Basics of Graphs` `Graph Traversal`

    - Enter a search query in the input box below.
    - Retrieve the most relevant pages from the document.

    **Example Queries:**
    ```
    1. What is set theory?
    2. How to analyze complexity using Big O notation?
    3. Explain Binary Search Algorithm and provide implementation of lower bound.
    4. How to use Stack, Queue and Deque?
    5. Explain generating permutations and subsets in Complete Search.
    6. What makes an algorithm greedy?
    7. How to solve the coin change problem using Dynamic Programming?
    8. What is amortized analysis?
    9. How to implement prefix sum array?
    10. Explain basic bitwise operations.
    11. Show graph representation using matrix and adjacency list.
    12. How to implement DFS and BFS traversal?
    ```

    Try it out and see how the AI retrieves the information you need efficiently.
    """
)

# Display placeholder for document details with a badge for the PDF
st.markdown(
    """
    #### Document Overview:
    **Title:** Competitive Programmer's Handbook  
    **Author:** Antti Laaksonen  
    **Type:** PDF  
    **Included Pages:** 100+  
    **Included Chapters Count:** 12  
    **Full Book Chapters Count:** 30

    <a href="https://cses.fi/book/book.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,
)
# Search and retrieve documents
st.subheader("Search and Retrieve Documents")

# User inputs for search
query = st.text_input(
    "Enter your search query:",
    help="Type a query to search through 12 chapters of the Competitive Programmer's Handbook.",
)
top_k = st.slider(
    "Number of Top Documents to Retrieve",
    min_value=1,
    max_value=10,
    value=3,
    help="Select the number of top pages to retrieve.",
)

collection_name = "Algorithm"
if st.button("Search"):
    if not query:
        st.error("Please enter a search query to retrieve results.")
    else:
        with st.spinner("Searching..."):
            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 "search_results" in st.session_state and st.session_state["search_results"]:
    st.write("### Search Results")
    # Remove the columns and display images one by one
    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)))

        # Create a container for each result
        with st.container():
            # Display the image with increased width
            st.image(
                img,
                caption=f"Document: {document_name}, Page: {page_number}",
                use_container_width=True,  # This makes the image use full column width
            )
            # Add some spacing between images
            st.markdown("---")

    st.success("Search completed successfully.")

st.markdown("----")
st.markdown(
    """
    <div style="text-align: center; padding: 10px;">
        Developed by <a href="mailto:abdulhaleem@tjmlabs.com" style="color: #ff4b4b; text-decoration: none;">Abdulhaleem</a> @ 
        <a href="https://tjmlabs.com" style="color: #ff4b4b; text-decoration: none;">TJM Labs</a>
    </div>
    """,
    unsafe_allow_html=True,
)