import gradio as gr
import shutil, openai, os

from langchain.document_loaders.generic import GenericLoader
from langchain.document_loaders.parsers import OpenAIWhisperParser
from langchain.document_loaders.blob_loaders.youtube_audio import YoutubeAudioLoader

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

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

def invoke(openai_api_key, youtube_url, prompt):
    openai.api_key = openai_api_key
    url = youtube_url
    save_dir = "docs/youtube/"
    loader = GenericLoader(
        YoutubeAudioLoader([url], save_dir),
        OpenAIWhisperParser()
    )
    docs = loader.load()
    shutil.rmtree(save_dir)
    retrieval = docs[0].page_content
    ###
    return retrieval

description = """The app demonstrates how to use a Large Language Model (LLM) with Retrieval Augmented Generation (RAG) on external data.\n\n
                 Enter an OpenAI API key, YouTube URL, and prompt to search the video, summarize it, translate it, etc.\n\n
                 <a href='https://www.gradio.app/'>Gradio</a> UI using <a href='https://platform.openai.com/'>OpenAI</a> API 
                 with <a href='https://openai.com/research/whisper'>Whisper</a> foundation model (for speech to text generation) 
                 and <a href='https://openai.com/research/gpt-4'>GPT-4</a> foundation model (for LLM use cases like search, summarization, translation, etc.) 
                 via AI-first toolkit <a href='https://www.langchain.com/'>LangChain</a>."""

gr.close_all()
demo = gr.Interface(fn=invoke, 
                    inputs = [gr.Textbox(label = "OpenAI API Key", lines = 1), gr.Textbox(label = "YouTube URL", lines = 1), gr.Textbox(label = "Prompt", lines = 1)],
                    outputs = [gr.Textbox(label = "Completion", lines = 1)],
                    title = "Generative AI - RAG",
                    description = description)
demo.launch()