# app.py
import os
import gradio as gr
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
from langchain_community.document_transformers import LongContextReorder
from config import LLM_MODEL, STREAMING
from embeddings import get_embeddings
from retrievers import load_retrievers
from llm import get_llm
from prompt import get_prompt


def create_rag_chain(chat_history):
    embeddings = get_embeddings()
    retriever = load_retrievers(embeddings)
    llm = get_llm(streaming=STREAMING)
    prompt = get_prompt(chat_history)

    return (
        {
            "context": retriever
            | RunnableLambda(LongContextReorder().transform_documents),
            "question": RunnablePassthrough(),
        }
        | prompt
        | llm.with_config(configurable={"llm": LLM_MODEL})
        | StrOutputParser()
    )


def respond_stream(message, history):
    rag_chain = create_rag_chain(history)
    response = ""
    for chunk in rag_chain.stream(message):
        response += chunk
        yield response


def respond(message, history):
    rag_chain = create_rag_chain(history)
    return rag_chain.invoke(message)


demo = gr.ChatInterface(
    respond_stream if STREAMING else respond,
    title="판례에 대해서 물어보세요!",
    description="안녕하세요!\n저는 판례에 대한 인공지능 QA봇입니다. 판례에 대해 깊은 지식을 가지고 있어요. 판례에 관한 도움이 필요하시면 언제든지 질문해주세요!",
)

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