import gradio as gr
from langchain.chains import RetrievalQA
from langchain.embeddings import OpenAIEmbeddings
from langchain.llms import OpenAI
from langchain.vectorstores import Chroma


PERSIST_DIR_NAME = "podcast-75"


def get_retrieval_qa() -> RetrievalQA:
    embeddings = OpenAIEmbeddings()
    db = Chroma(persist_directory=PERSIST_DIR_NAME, embedding_function=embeddings)
    retriever = db.as_retriever()
    return RetrievalQA.from_chain_type(
        llm=OpenAI(), chain_type="stuff", retriever=retriever
    )


def main(query: str):
    qa = get_retrieval_qa()
    answer = qa(query)
    return answer["result"]


pyhack_qa = gr.Interface(
    fn=main,
    inputs=[gr.Textbox(label="query")],
    outputs="text",
)
pyhack_qa.launch()