# @title set API key
import os
import getpass
from pprint import pprint
import warnings

warnings.filterwarnings("ignore")

if "UPSTAGE_API_KEY" not in os.environ:    
    os.environ["UPSTAGE_API_KEY"] = getpass.getpass("Enter your Upstage API key: ")


from pydantic import BaseModel
from langchain_upstage import ChatUpstage, UpstageEmbeddings

from langchain_chroma import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders.csv_loader import CSVLoader
from langchain_community.document_loaders import JSONLoader
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_upstage import ChatUpstage
from langchain import hub

from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

import pandas as pd



import gradio as gr
from langchain_core.messages import AIMessage, HumanMessage

rag_with_history_prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            """
You are an intelligent assistant helping the members of the Korean National Assembly with questions related to law and policy. As you will respond to the Korean National Assembly, you must answer politely. Read the given questions carefully and give the answer in Korean ONLY using the following pieces of the context.

Do not try to make up an answer:
 - If the answer to the question cannot be determined from the context alone, say "I cannot determine the answer to that."
 - If the context is empty, just say "I do not know the answer to that."

Answer the question chronologically by issue.

Context: {context}

Question: {input} 구체적으로 시간순으로 답변해줘.

Answer:
            """,
        ),
        MessagesPlaceholder(variable_name="history"),
        ("human", "{input}"),
    ]
)

llm = ChatUpstage()
chain = rag_with_history_prompt | llm | StrOutputParser()


DB_PATH = './chroma_db/chroma_db'
db = Chroma(persist_directory=DB_PATH, embedding_function=UpstageEmbeddings(model="solar-embedding-1-large"))
retriever = db.as_retriever(search_kwargs={"k": 3})


history = []

def chatbot_response(input_text, history):
    result_docs = retriever.invoke(input_text)  
    response = chain.invoke({"history": history, "context": result_docs, "input": input_text}) 

    history.append((input_text, response))
    
    return history, history  

with gr.Blocks() as demo:
    gr.Markdown("## 국회 회의록 기반 의정활동 지원 및 대국민 알권리 보장 챗봇")
    
    chatbot = gr.Chatbot(label="챗봇")
    txt = gr.Textbox(label="질문을 입력하세요")
    submit_btn = gr.Button("질문하기")
    
    submit_btn.click(chatbot_response, inputs=[txt, chatbot], outputs=[chatbot, chatbot])

# 앱 실행
demo.launch(share=True)