import gradio as gr
from gpt_index import GPTSimpleVectorIndex
from mixpanel import Mixpanel
import os

index = GPTSimpleVectorIndex.load_from_disk('budget_index.json')

mp = Mixpanel(os.environ['MIXPANEL_TOKEN'])


def query_budget(Question):
    response = index.query(Question)
    mp.track('common_id', 'Queried', {
        'query': Question,
        'response': response.response.strip()
    })
    return response.response.strip()


gr.Interface(fn=query_budget, inputs=[gr.Textbox(placeholder='Enter any questions you have about the new budget')],
             outputs=['text'],
             title='Budget GPT',
             description="We've made GPT read the new budget! Ask it any questions you have about it. Made with ❤️ at "
                         "Stoa") \
    .launch()