import gradio as gr
import shutil, openai, os

from langchain.chains import RetrievalQA
from langchain.chat_models import ChatOpenAI
from langchain.document_loaders.blob_loaders.youtube_audio import YoutubeAudioLoader
from langchain.document_loaders.generic import GenericLoader
from langchain.document_loaders.parsers import OpenAIWhisperParser
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.prompts import PromptTemplate
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma

from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv())

#openai.api_key = os.environ["OPENAI_API_KEY"]

template = """Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up 
              an answer. Keep the answer as concise as possible. Always say "🔥 Thanks for using the app, Bernd Straehle." at the end of the answer. 
              {context} Question: {question} Helpful Answer: """

QA_CHAIN_PROMPT = PromptTemplate(input_variables = ["context", "question"], template = template)

print(0)

qa_chain = None 

def invoke(openai_api_key, youtube_url, prompt):
    openai.api_key = openai_api_key
    if (os.path.isdir("docs/chroma/") == False):
        print(1)
        youtube_dir = "docs/youtube/"
        loader = GenericLoader(YoutubeAudioLoader([youtube_url], youtube_dir), OpenAIWhisperParser())
        docs = loader.load()
        text_splitter = RecursiveCharacterTextSplitter(chunk_size = 1500, chunk_overlap = 150)
        splits = text_splitter.split_documents(docs)
        chroma_dir = "docs/chroma/"
        vectordb = Chroma.from_documents(documents = splits, embedding = OpenAIEmbeddings(), persist_directory = chroma_dir)
        llm = ChatOpenAI(model_name = "gpt-4", temperature = 0)
        global qa_chain = RetrievalQA.from_chain_type(llm, retriever = vectordb.as_retriever(), return_source_documents = True, chain_type_kwargs = {"prompt": QA_CHAIN_PROMPT})
    print(2)
    result = global qa_chain({"query": prompt})
    shutil.rmtree(youtube_dir)
    #shutil.rmtree(chroma_dir)
    return result["result"]

description = """The app demonstrates how to use a <strong>Large Language Model</strong> (LLM) with <strong>Retrieval Augmented Generation</strong> (RAG) on external data. 
                 Enter an OpenAI API key, YouTube URL (external data), and prompt to perform semantic search, sentiment analysis, summarization, translation, etc.\n\n
                 Implementation: <a href='https://www.gradio.app/'>Gradio</a> UI using <a href='https://platform.openai.com/'>OpenAI</a> API 
                 via AI-first <a href='https://www.langchain.com/'>LangChain</a> toolkit with <a href='https://openai.com/research/whisper'>Whisper</a> (speech to text) 
                 and <a href='https://openai.com/research/gpt-4'>GPT-4</a> (LLM use cases) foundation models as well as AI-native 
                 <a href='https://www.trychroma.com/'>Chroma</a> embedding database."""

gr.close_all()
demo = gr.Interface(fn=invoke, 
                    inputs = [gr.Textbox(label = "OpenAI API Key", value = "sk-", lines = 1), gr.Textbox(label = "YouTube URL", value = "https://www.youtube.com/watch?v=--khbXchTeE", lines = 1), gr.Textbox(label = "Prompt", value = "GPT-4 human level performance", lines = 1)],
                    outputs = [gr.Textbox(label = "Completion", lines = 1)],
                    title = "Generative AI - LLM & RAG",
                    description = description)
demo.launch()