import os
from dotenv import load_dotenv
load_dotenv()

import gradio as gr
from ui.gradio_ui import get_ui
from style import custom_css

# Import RAG pipeline
from modules.lawbot.rag_with_langchain import (
    load_documents,
    split_documents,
    create_vectorstore,
    setup_rag_chain
)

def initialize_rag():
    documents = load_documents()
    docs = split_documents(documents)
    vectorstore = create_vectorstore(docs)
    return setup_rag_chain(vectorstore)

qa_chain = initialize_rag()

demo = get_ui(qa_chain)

demo.css = custom_css

if __name__ == "__main__":
    demo.launch()