import gradio as gr
from dotenv import load_dotenv
from custom_retriever import query_engine

load_dotenv()

def get_response(query, k_value):
    response, context_str = query_engine.custom_query(query, k_value)
    return str(response), context_str

with gr.Blocks() as demo:
    gr.Markdown("**Welcome to our AI powered business search engine!**")
    with gr.Row():
        inp = gr.Textbox(placeholder="Please enter your query", label="User Query")
        k_value = gr.Slider(minimum=1, maximum=10, value=5, label="Top K", step=1)
    with gr.Row():
        context_retrieved = gr.Textbox(placeholder="Context retrieved and sent to ChatGPT", label="Context")
        out = gr.Textbox()
    btn = gr.Button("Run")
    btn.click(fn=get_response, inputs=[inp, k_value], outputs=[out, context_retrieved])

demo.launch()