import os
import pandas as pd
import gradio as gr

from langchain.vectorstores import Chroma
from langchain.embeddings import OpenAIEmbeddings
from langchain.llms import OpenAI
from langchain.chat_models import ChatOpenAI
from langchain.chains import RetrievalQA

def Loading():
    return "데이터 로딩 중..."

    
def LoadData(openai_key):
        
    if openai_key is not None:
        
        os.environ["OPENAI_API_KEY"] = openai_key
        
        persist_directory = 'realdb_LLM'

        embedding = OpenAIEmbeddings()


        vectordb = Chroma(
            persist_directory=persist_directory, 
            embedding_function=embedding
        )

        global retriever
        retriever = vectordb.as_retriever(search_kwargs={"k": 1})
        
        return "준비 완료"
    else:
        return "사용하시는 API Key를 입력하여 주시기 바랍니다."


# 챗봇의 답변을 처리하는 함수
def respond(message, chat_history, temperature, top_p):
    try:
        
        print(temperature)

        qa_chain = RetrievalQA.from_chain_type(
            llm=OpenAI(temperature=temperature, top_p=top_p),
            # llm=OpenAI(temperature=0.4), 
            # llm=ChatOpenAI(temperature=0),
            chain_type="stuff", 
            retriever=retriever
        )

        result = qa_chain(message)
        
        bot_message = result['result']
        
        # 채팅 기록에 사용자의 메시지와 봇의 응답을 추가.
        chat_history.append((message, bot_message))
        
        return "", chat_history
    except:
        chat_history.append(("", "API Key 입력 요망"))
        
        return " ", chat_history
    

# 챗봇 설명
title = """
<div style="text-align: center; max-width: 500px; margin: 0 auto;">
    <div>
        <h1>Pretraining Chatbot V2 Real</h1>
    </div>
    <p style="margin-bottom: 10px; font-size: 94%">
        OpenAI LLM를 이용한 Chatbot (Similarity)
    </p>
</div>
"""

# 꾸미기
css="""
#col-container {max-width: 700px; margin-left: auto; margin-right: auto;}
"""
with gr.Blocks(css=css) as UnivChatbot:
    with gr.Column(elem_id="col-container"):
        gr.HTML(title)

    with gr.Row():
        with gr.Column(scale=3):
            openai_key = gr.Textbox(label="You OpenAI API key", type="password", placeholder="OpenAI Key Type", elem_id="InputKey", show_label=False, container=False)
        with gr.Column(scale=1):
            langchain_status = gr.Textbox(placeholder="Status", interactive=False, show_label=False, container=False)
    
    with gr.Row():
        with gr.Column(scale=4):
            temperature = gr.Slider(
            label="Temperature",
            minimum=0,
            maximum=2.0,
            step=0.01,
            value=0.7,
        )
        with gr.Column(scale=4):
            top_p = gr.Slider(
            label="Top_p",
            minimum=0,
            maximum=1,
            step=0.01,
            value=0.5,
        )
        with gr.Column(scale=1):
            chk_key = gr.Button("확인", variant="primary")
                    
    chatbot = gr.Chatbot(label="대학 챗봇시스템(OpenAI LLM)", elem_id="chatbot") # 상단 좌측 

    with gr.Row():
        with gr.Column(scale=9):
            msg = gr.Textbox(label="입력", placeholder="궁금하신 내역을 입력하여 주세요.", elem_id="InputQuery", show_label=False, container=False)
        
    with gr.Row():
        with gr.Column(scale=1):
            submit = gr.Button("전송", variant="primary")
        with gr.Column(scale=1):
            clear = gr.Button("초기화", variant="stop")

    #chk_key.click(Loading, None, langchain_status, queue=False)  
    chk_key.click(
        fn=LoadData, 
        inputs=[openai_key], 
        outputs=[langchain_status], 
        queue=False
    )
    
    # 사용자의 입력을 제출(submit)하면 respond 함수가 호출.
    msg.submit(
        fn=respond, 
        inputs=[msg, chatbot, temperature, top_p], 
        outputs=[msg, chatbot]
    )

    submit.click(respond, [msg, chatbot, temperature, top_p], [msg, chatbot])

    # '초기화' 버튼을 클릭하면 채팅 기록을 초기화.
    clear.click(lambda: None, None, chatbot, queue=False)

  
UnivChatbot.launch()