import gradio as gr
import os
os.environ["OPENAI_API_KEY"] = os.getenv('OPENAI_API_KEY')
from langchain.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.embeddings import OpenAIEmbeddings
from langchain.chat_models import ChatOpenAI
from langchain.retrievers.multi_query import MultiQueryRetriever
from langchain.chains import RetrievalQA

#Loader
loader = PyPDFLoader("2020_emaster_keynote.pdf")
pages = loader.load_and_split()

#Split
text_splitter = RecursiveCharacterTextSplitter(
    # Set a really small chunk size, just to show.
    chunk_size = 300,
    chunk_overlap  = 20,
    length_function = len,
    is_separator_regex = False,
)
texts = text_splitter.split_documents(pages)

#Embedding
embeddings_model = OpenAIEmbeddings()

# load it into Chroma
db = Chroma.from_documents(texts, embeddings_model)

# question = "데이터베이스의 정의?"
llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0.5)
qa_chain = RetrievalQA.from_chain_type(llm,retriever=db.as_retriever())


# 인터페이스를 생성.
with gr.Blocks() as demo:
    gr.Image('images/emaster.png')

    gr.Text('''
                이런 질문 어떠세요? 
                데이터베이스의 정의
            ''')    
    chatbot = gr.Chatbot(label="정보처리산업기사챗봇") # 청년정책챗봇 레이블을 좌측 상단에 구성
    msg = gr.Textbox(label="질문해주세요!")  # 하단의 채팅창의 레이블
    clear = gr.Button("대화 초기화")  # 대화 초기화 버튼

    # 챗봇의 답변을 처리하는 함수
    def respond(message, chat_history):
      result = qa_chain({"query": message})
    #   result = qa_chain(message)
      bot_message = result['result']

      # 채팅 기록에 사용자의 메시지와 봇의 응답을 추가.
      chat_history.append((message, bot_message))
      return "", chat_history

    # 사용자의 입력을 제출(submit)하면 respond 함수가 호출.
    msg.submit(respond, [msg, chatbot], [msg, chatbot])

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

# 인터페이스 실행.
demo.launch(debug=True, share=True)