import gradio as gr
from main import (
    answer_query, 
    set_keys,
    handle_file
    )

from pydantic import ConfigDict
model_config = ConfigDict(protected_namespaces=())

setting_keys = gr.Interface(
    fn=set_keys,
    inputs=[
        gr.Textbox(label="Enter your CO_API_KEY"),
        gr.Textbox(label="Enter your LLAMA_CLOUD_API_KEY"),
    ],
    outputs=gr.Textbox(label="Status")
)

uploading_files = gr.Interface(
    fn=handle_file,
    inputs=gr.File(
        label="Upload a file",
        file_count="single",
        file_types=["text", ".pdf"],
    ),
    outputs=gr.Textbox(label="Status")
)
    
qa = gr.Interface(
    fn=answer_query,
    inputs=gr.Textbox(label="Enter your question"),
    outputs=[
        gr.Textbox(label="Answer"), 
        gr.Textbox(label="Relevant Nodes"),
    ],
    title="Document Q&A System"
)

demo = gr.TabbedInterface(
    interface_list=[setting_keys, uploading_files, qa],
    tab_names=["Settings",  "Upload File", "Q&A System"]
    )

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