import fn
import gradio as gr
from refresh import create_refresh_button

def fn_search(dir, query):
    args = {
        'dir': dir,
        'query': query,
        'k': 1,
    }
    vector_store = fn.vector_load(args)
    result, detail = fn.search(vector_store, args)

    return result

with gr.Blocks() as demo:
    title = gr.Markdown('# FAISS(Naive RAG)')
    info = gr.Markdown()
    upload_dir = gr.Dropdown(
        choices=fn.load_dirs(),
        label='name',
        show_label=True,
        interactive=True,
        allow_custom_value=True,
    )

    with gr.Tab('Upload'):
        chunk_size = gr.Textbox(
            value=0,
            label='(optional) chunk_size if split',
            show_label=True,
            interactive=True,
        )
        rag_zip = gr.UploadButton(
            label='Zip Upload and save rag_dir',
            interactive=True,
        )
        create_refresh_button(gr, upload_dir, lambda: None, lambda: {'choices': fn.load_dirs()}, 'refresh-button', interactive=True)

    with gr.Tab('Search'):
        query = gr.Textbox(
            lines=1,
            label='query',
            show_label=True,
            interactive=True,
            show_copy_button=True,
            )
        search_button = gr.Button(value='search')
        result = gr.Textbox(label='result', show_label=True, show_copy_button=True)

    search_button.click(
        fn=fn_search,
        inputs=[upload_dir, query],
        outputs=[result],
        )
        
    rag_zip.upload(
        fn=fn.upload,
        inputs=[upload_dir, chunk_size, rag_zip],
        outputs=[info],
        )

if __name__ == '__main__':
    demo.launch()