app.py

from typing import List
from dataclasses import asdict
import pandas as pd
import gradio as gr

from SmartSearch.database.chromadb import ChromaDB
from SmartSearch.providers.SentenceTransformerEmbedding import SentenceTransformerEmbedding
from utils import combine_metadata_with_distance
st_chroma = ChromaDB(
    embedding_function=SentenceTransformerEmbedding(model_name='all-mpnet-base-v2'),
    collection_name="novel_mockup_collection"
)

# Function to search for products
def search_novels(query, k):
    result = st_chroma.search(query_text=query, n_results=k)
    
    result = combine_metadata_with_distance(result['metadatas'], result['distances'])
    result = pd.DataFrame(result)
    return result

with gr.Blocks() as demo:
    with gr.Row():
        query = gr.Textbox(label="Search Query", placeholder="write a query to find the novels")
    with gr.Row():
        # search_type = gr.Dropdown(label="Search Type", choices=['semantic', 'keyword', 'hybrid'], value='hybrid')
        k = gr.Number(label="Items Count", value=10)
        # rerank = gr.Checkbox(value=True, label="Rerank")
    results = gr.Dataframe(label="Search Results")
    
    search_button = gr.Button("Search", variant='primary')
    search_button.click(fn=search_novels, inputs=[query, k], outputs=results)

demo.launch()