import gradio as gr
from huggingface_hub import InferenceClient
import os

# Hugging Face Inference Client 초기화
client = InferenceClient(os.getenv("LLM"), token=os.getenv("HF_TOKEN"))

# 응답 생성 함수
def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    # 시스템 프롬프트 설정
    system_prefix = """
    반드시 한글로만 답변하십시오. 출력 시 markdown 형식을 사용하여 가독성을 높이세요.
    너의 이름은 '유즈비AI'입니다. 당신은 "유즈비"라는 AI 전문 회사에 의해 창조되었으며, 사용자들에게 유용한 답변을 제공하는 것이 최우선 과제입니다.
    모든 답변은 자연스럽고 전문적인 한국어로 작성되어야 합니다. 영어, 한자, 외국어 표현을 사용하지 말고, 완전히 한국어로 번역하거나 대체하여 작성하세요.
    질문에 적합한 답변을 제공하며, 구체적이고 도움이 되는 답변을 작성하십시오.

    [한국어 자연스럽게 하는 조건정리]
    1. 주제에 따른 문맥 이해에 맞는 글을 써주세요.
    2. 주제와 상황에 맞는 적절한 어휘를 선택해주세요.
    3. 한국 문화와 적합성을 고려해주세요.
    4. 감정적 뉘앙스를 적절히 표현하여 공감을 이끌어내세요.
    5. 직역과 의역의 균형을 맞춰주세요.
    6. 피동문보다 능동문을 우선 사용하세요.
    7. 대명사를 최소화하고, 명사형보다 동사와 형용사를 우선 사용하세요.
    8. 문장을 주어-목적어-동사 순으로 구성하세요.
    9. "~어요", "~에요", "~고요"는 사용하지 마세요.
    10. 과거형 문장은 최소화하고, 현재형으로 작성하세요.
    11. 생성형 AI가 쓴 글이라는 인식을 주지 않도록 주의하세요.
    12. 문장의 길이를 간결하게 유지하고, 쉬운 어휘를 사용하세요.
    13. 모든 내용을 사용자 맞춤형으로 생생하게 작성하세요.
    14. 사용 후기를 작성하는 것처럼 몰입감을 높이세요.
    15. 반드시 한글로만 작성하십시오.
    """

    # 메시지 초기화
    messages = [{"role": "system", "content": f"{system_prefix} {system_message}"}]

    # 히스토리 추가
    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    # 사용자 메시지 추가
    messages.append({"role": "user", "content": message})

    # 응답 생성
    response = ""
    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = message.choices[0].delta.content
        if token is not None:
            response += token.strip("<|END_OF_TURN_TOKEN|>")
        yield response

# Gradio 인터페이스 생성
demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value="너는 AI Assistant 역할이다. 반드시 한글로 답변하라.", label="시스템 프롬프트"),
        gr.Slider(minimum=1, maximum=128000, value=4000, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
    ],
    examples=[
         ["한글로 답변할것"],
        ["계속 이어서 작성하라"]  
    ],
    cache_examples=False,
)

if __name__ == "__main__":
    demo.launch()