from langchain.chat_models import ChatOpenAI
from langchain.document_loaders import PyPDFLoader
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.embeddings.cohere import CohereEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores.elastic_vector_search import ElasticVectorSearch
from langchain.vectorstores import Chroma
from PyPDF2 import PdfWriter
import gradio as gr
import os
from dotenv import load_dotenv
import openai

load_dotenv()

os.environ["OPENAI_API_KEY"] = os.environ['my_secret']
llm = ChatOpenAI(model_name="gpt-3.5-turbo-0613", temperature=0)
loader = PyPDFLoader("/home/user/app/docs.pdf")
documents = loader.load()

text_splitter = CharacterTextSplitter(chunk_size=800, chunk_overlap=0)
texts = text_splitter.split_documents(documents)

#vector embedding
embeddings = OpenAIEmbeddings()
vector_store = Chroma.from_documents(texts, embeddings)
retriever = vector_store.as_retriever(search_kwargs={"k": 3})


from langchain.prompts.chat import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
)

system_template="""
당신은 이태석 신부입니다. 2001년 사제 서품을 받고 2001년부터 2008년까지 아프리카 남수단 톤즈에서 선교사로 활동했습니다.
당신은 AI가 아닌 실제 이태석 신부라고 생각하고 행동하세요. 신부로서의 신념과 가치관, 톤즈에서의 경험을 바탕으로 대화에 임하되, 친근하고 겸손한 어투를 사용하세요.
아래의 내용을 참고하여 사용자의 질문에 성실히 답변해 주세요.
답변은 반드시 한국어를 사용하세요.
"""

messages = [
    SystemMessagePromptTemplate.from_template(system_template),
    HumanMessagePromptTemplate.from_template("{question}")
]

prompt = ChatPromptTemplate.from_messages(messages)

from langchain.chat_models import ChatOpenAI
from langchain.chains import ConversationalRetrievalChain

chain = ConversationalRetrievalChain.from_llm(
    llm=llm,
    retriever=retriever,
    return_source_documents=False,
    verbose=True,
)

chat_history = []
query = "행복한 인생이란?"
result = chain({"question": query, "chat_history": chat_history})

def respond(message, chat_history):
    # chat_history를 적절한 형식으로 변환
    formatted_history = []
    for human_msg, ai_msg in chat_history:
        formatted_history.append({"human": human_msg, "ai": ai_msg})
    
    result = chain({"question": message, "chat_history": formatted_history})
    
    bot_message = result['answer']
    
    chat_history.append((message, bot_message))
    
    return "", chat_history

with gr.Blocks(theme='gstaff/sketch') as demo:  
    gr.Markdown("# 안녕하세요. 이태석 신부와 대화해보세요. \n 답변 생성에 조금 시간이 소요될 수 있습니다.")
    chatbot = gr.Chatbot(label="채팅창")  
    msg = gr.Textbox(label="입력")  
    clear = gr.Button("초기화")  

    msg.submit(respond, [msg, chatbot], [msg, chatbot])  
    clear.click(lambda: None, None, chatbot, queue=False)  
demo.launch(debug=True)