import gradio as gr
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from openai import OpenAI
from translate_utils import translate_ko_to_en

YOUR_OPENAI_API_KEY = "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"
# 임베딩 모델 로드
embedding_model_name = "snunlp/KR-SBERT-V40K-klueNLI-augSTS"
embedding_model = HuggingFaceEmbeddings(
    model_name=embedding_model_name,
    model_kwargs={"device": "cpu"},  # GPU 설정 시 "cuda"로 변경 가능
    encode_kwargs={"normalize_embeddings": True},
)

# vector DB 로드드
save_path = "./my_faiss_index"
vectorstore = FAISS.load_local(save_path, embedding_model,allow_dangerous_deserialization=True)




def chatbot(input_question, eng_trans = True, num_ref = 3):
    retriever = vectorstore.as_retriever(search_kwargs={"k": num_ref})
    if eng_trans == False:
        basic_docs = retriever.invoke(input_question)
    else:
        eng = translate_ko_to_en(input_question)

        basic_docs = retriever.invoke(input_question)
        eng_docs = retriever.invoke(eng)
        basic_docs = basic_docs + eng_docs

    context = "\n".join([doc.page_content for doc in basic_docs])  

    client = OpenAI(
    api_key=YOUR_OPENAI_API_KEY)

    # question과 context를 message에 삽입
    completion = client.chat.completions.create(
        model="gpt-4o-mini",  
        messages=[
            {
                "role": "user",
                "content": f"""당신은 반도체와 차세대반도체학과에 대해서 설명하는 Assistant입니다.
    차세대반도체학과는 한국 서울의 중앙대학교의 창의ICT공과대학에 설립된 학과입니다. 
    다음 맥락에 맞게 질문에 답하세요.

    맥락: {context}

    질문: {input_question}
    """
            }
        ]
    )



    return completion.choices[0].message.content




demo = gr.Interface(fn=chatbot, 
                    inputs=[gr.Textbox(label="질문 입력"),  # 텍스트 입력
                            gr.Checkbox(label="영어 번역 사용 여부", value=True),  # 체크박스 추가
                            gr.Slider(1,5, value=3, step=1),
    ], outputs="text")
demo.launch()