from fastapi import FastAPI
from utils import get_us_speeches
from config import UPDATE_SPEECHES

from haystack.document_stores import ElasticsearchDocumentStore
from haystack.nodes import ElasticsearchRetriever
from haystack.nodes import FARMReader
from haystack.pipelines import ExtractiveQAPipeline

import gradio as gr


document_store = ElasticsearchDocumentStore(
    host='fgm-v2.es.eastus2.azure.elastic-cloud.com',
    username='elastic',
    password='cxjWqZfmhcfhzpWmfX57ylJc',
    scheme='https',
    port=9243,
    index='us-speeches'

)

if UPDATE_SPEECHES:
    us_speeches = get_us_speeches()
    document_store.write_documents(us_speeches)

retriever = ElasticsearchRetriever(
    document_store=document_store
)

reader = FARMReader(model_name_or_path="deepset/roberta-base-squad2", use_gpu=False)

pipeline = ExtractiveQAPipeline(reader=reader, retriever=retriever)

app = FastAPI()


async def run_query(query: str):
    return pipeline.run(query=query)


gr.Interface(run_query, "textbox", ["label", "label"]).launch()