import gradio as gr
from src.rag import RAG

import bs4
from langchain_community.document_loaders import WebBaseLoader
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter

from dotenv import load_dotenv

# set the required env variables
load_dotenv(".env")

def rag_handler(web_paths, model_name, temperature, question):
    print(web_paths)
    web_paths = web_paths.split(',')
    print(web_paths)
    loader = WebBaseLoader(
        web_paths=web_paths,
        bs_kwargs=dict(
            parse_only=bs4.SoupStrainer(
                class_=("post-content", "post-title", "post-header")
            )
        ),
    )

    llm = ChatOpenAI(model_name=model_name, temperature=temperature)
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) # TODO: Parameterize this
    
    rag_pipeline = RAG(llm, loader, text_splitter, OpenAIEmbeddings)

    return rag_pipeline.invoke(question)

def create_rag_interface():
    return gr.Interface(
        fn=rag_handler,
        inputs=[
            gr.Textbox(value="https://lilianweng.github.io/posts/2023-06-23-agent/"),
            gr.Dropdown(["gpt-3.5-turbo"], type="value"),
            gr.Slider(0, 1, step=0.1),
            'text'
        ],
        outputs="text"
    )


if __name__ == '__main__':
    interface_list = []
    interface_list.append(create_rag_interface())

    demo = gr.TabbedInterface(interface_list)

    demo.launch()