import gradio as gr
from autollm import AutoQueryEngine
import os
import spaces

@spaces.GPU
def query_engine(llm_model, document, query):
    api_key = os.getenv("HUGGINGFACE_API_KEY")
    llm_api_base = "https://api-inference.huggingface.co/models/"

    content = str(document)

    query_engine = AutoQueryEngine.from_defaults(
        documents=[content],
        llm_model=llm_model,
        llm_api_base=llm_api_base,
    )
    response = query_engine.query(query)
    return response

interface = gr.Blocks()
with interface:
    gr.Markdown("# AutoQueryEngine Interface")
    llm_model = gr.Textbox(label="LLM Model", value="mistralai/Mixtral-8x7B-Instruct-v0.1")
    document = gr.File(label="Documents")
    query = gr.Textbox(label="Query")
    output = gr.Textbox(label="Output")
    query_btn = gr.Button("Query")
    query_btn.click(fn=query_engine, inputs=[llm_model, document, query], outputs=output)
interface.launch()