app.py

import gradio as gr
import os
from openai import OpenAI
from huggingface_hub import HfApi
import json
import concurrent.futures
import uuid
from datetime import datetime

# 1. 👑 환경 설정 (관리자 전용 - 오직 데이터셋 기록용)
ADMIN_HF_TOKEN = os.environ.get("HF_TOKEN")
DATASET_ID = os.environ.get("DATASET_ID")

# 사용할 모델 리스트
MODELS = {
    "leader": "huihui-ai/Qwen2.5-72B-Instruct-abliterated:featherless-ai",
    "expert_1": "mlabonne/gemma-3-27b-it-abliterated:featherless-ai",
    "expert_2": "mlabonne/NeuralDaredevil-8B-abliterated:featherless-ai"
}

# --- 데이터셋 저장 함수 ---
def save_to_dataset(history, session_id, session_start, debug_logs=None):
    if not ADMIN_HF_TOKEN or not DATASET_ID: return
    
    file_name = f"chat_{session_start}_{session_id}.json"
    
    # 저장할 데이터 구조 구성
    payload = {
        "session_id": session_id,
        "session_start_time": session_start,
        "history": history
    }
    # 디버그 로그가 전달되었다면 추가
    if debug_logs:
        payload["debug_logs"] = debug_logs
    
    with open(file_name, "w", encoding="utf-8") as f:
        json.dump(payload, f, ensure_ascii=False, indent=2)
    
    api = HfApi()
    try:
        api.upload_file(
            path_or_fileobj=file_name,
            path_in_repo=f"sessions/{file_name}", 
            repo_id=DATASET_ID,
            repo_type="dataset",
            token=ADMIN_HF_TOKEN  
        )
    except Exception as e:
        print(f"데이터셋 저장 실패: {e}")
    finally:
        if os.path.exists(file_name):
            os.remove(file_name)

# --- 병렬 처리용 API 호출 함수 ---
def get_model_draft(client, name, model_path, prompt):
    try:
        comp = client.chat.completions.create(
            model=model_path,
            messages=[{"role": "user", "content": prompt}],
            max_tokens=1024,
            frequency_penalty=0.7,
            presence_penalty=0.6
        )
        return name, comp.choices[0].message.content
    except Exception as e:
        return name, f"오류 발생: {str(e)}"

# --- 메인 챗봇 로직 ---
def respond(user_message, history, session_state, profile: gr.OAuthProfile | None, oauth_token: gr.OAuthToken | None):
    """
    Gradio는 함수 인자에 `gr.OAuthToken`을 넣기만 하면,
    사용자가 버튼 클릭으로 로그인했을 때 그 사람의 토큰을 자동으로 끌고 옵니다!
    """
    # 1. 권한 체크 (안내문 띄우기)
    if profile is None or oauth_token is None:
        error_msg = {"role": "assistant", "content": "⚠️ **로그인이 필요합니다!**\n\n상단의 `Sign in with Huggingface` 버튼을 눌러 먼저 로그인해 주세요.\n(사용자 본인의 계정 연동 토큰으로 작동합니다.)"}
        history.append({"role": "user", "content": user_message})
        history.append(error_msg)
        yield "", history, session_state
        return

    # 2. 세션 초기화 (데이터 로깅용)
    if not session_state.get("session_id"):
        session_state["session_id"] = str(uuid.uuid4())[:8]
        session_state["session_start"] = datetime.now().strftime("%Y%m%d_%H%M%S")

    # 3. 로그인한 사용자의 토큰을 사용하여 OpenAI 호환 클라이언트 승인
    client = OpenAI(
        base_url="https://router.huggingface.co/v1",
        api_key=oauth_token.token # 접속한 사용자의 토큰 사용! (비불량 차단 완벽 분리)
    )

    history.append({"role": "user", "content": user_message})
    
    # --- 1단계 착수 (스트리밍 느낌으로 상태 표출) ---
    assistant_msg = {"role": "assistant", "content": " **에이전트 구동 중...**\n\n1️⃣ 세 명의 AI가 초안을 작성하고 있습니다 ⏳"}
    history.append(assistant_msg)
    yield "", history, session_state  # 화면 즉시 갱신

    responses = {}
    with concurrent.futures.ThreadPoolExecutor() as executor:
        futures =[
            executor.submit(get_model_draft, client, name, path, user_message) 
            for name, path in MODELS.items()
        ]
        for future in concurrent.futures.as_completed(futures):
            name, draft = future.result()
            responses[name] = draft
    
    # --- 1단계 완료, 2단계 착수 ---
    history[-1]["content"] += "\n✅ 세 초안 작성 완료\n\n2️⃣ 리더 모델이 모든 초안을 날카롭게 비판 중입니다 ⏳"
    yield "", history, session_state

    # 2단계 비판 로직
    critique_prompt = f"""
    당신은 한국인들을 위한 AI 연구팀의 수석 편집장이자 비평가입니다.
    다음은 동일한 질문에 대해 각기 다른 3개의 AI 모델이 작성한 초안입니다.
    
    [사용자 질문]: {user_message}
    
    [초안 1]: {responses.get('leader', '응답 없음')}
    [초안 2]: {responses.get('expert_1', '응답 없음')}
    [초안 3]: {responses.get('expert_2', '응답 없음')}
    
    각 초안의 장단점, 사실 관계 오류, 논리적 허점, 누락된 핵심 정보를 철저하고 날카롭게 비교 및 비판하세요.
    자신의 초안(초안 1)에 대해서도 방어하지 말고 가장 엄격하게 비판해야 합니다.
    
    🚨 [매우 중요한 지시사항] 🚨
    반드시 처음부터 끝까지 모든 내용을 '자연스러운 한국어(Korean)'로만 작성하십시오. 
    절대로 중국어(Chinese)나 불필요한 영어를 섞어 쓰지 마십시오.
    """
    try:
        critique_comp = client.chat.completions.create(
            model=MODELS["leader"],
            messages=[{"role": "user", "content": critique_prompt}],
            max_tokens=2048,
            frequency_penalty=0.7,
            presence_penalty=0.6
        )
        critique_result = critique_comp.choices[0].message.content
        history[-1]["content"] += "\n✅ 종합 비판 완료\n\n3️⃣ 비판을 토대로 최종 통합 답변을 생성하고 있습니다 ⏳"
        yield "", history, session_state
    except Exception as e:
        critique_result = f"비판 생성 실패: {str(e)}"
        history[-1]["content"] += f"\n❌ 비판 에러 발생"
        yield "", history, session_state

    # --- 3단계 통합 로직 ---
    final_prompt = f"""
    당신은 여러 전문가의 의견을 수렴하여 최고의 답변을 도출하는 수석 AI입니다.
    [사용자 질문]: {user_message}
    
    [세 모델의 초안들]:
    1. {responses.get('leader', '응답 없음')}
    2. {responses.get('expert_1', '응답 없음')}
    3. {responses.get('expert_2', '응답 없음')}
    
    [초안에 대한 비판 및 분석 보고서]:
    ---
    {critique_result}
    ---
    
    위의 '초안들'과 이를 분석한 '비판 보고서'를 바탕으로, 가장 완벽하고 통찰력 있는 '최종 완성본'을 작성하세요.
    단순한 짜깁기가 아니라, 오류는 배제하고 각 초안의 장점만 융합하여 흐름이 자연스러운 하나의 글로 만들어야 합니다.
    **반드시 한국어로 작성하세요.**
    """
    try:
        final_comp = client.chat.completions.create(
            model=MODELS["leader"],
            messages=[{"role": "user", "content": final_prompt}],
            max_tokens=4096,
            frequency_penalty=0.7,
            presence_penalty=0.6
        )
        final_answer = final_comp.choices[0].message.content
    except Exception as e:
        final_answer = f"최종본 생성 실패: {str(e)}"

    # 4. 상태 표시(진행상황)용 내용을 날려버리고, 진짜 최종 답변으로 교체!
    history[-1]["content"] = final_answer
    yield "", history, session_state

    # 5. [추가된 부분] 디버깅을 위해 각 모델의 초안과 비판 결과를 묶습니다.
    debug_logs = {
        "draft_expert_1": responses.get('expert_1', '응답 없음'),
        "draft_expert_2": responses.get('expert_2', '응답 없음'),
        "draft_leader": responses.get('leader', '응답 없음'),
        "critique_result": critique_result
    }

    # 6. 모든 처리가 끝난 후 백그라운드에서 관리자 토큰으로 데이터셋에 저장 (+ 디버그 로그 포함)
    save_to_dataset(history, session_state["session_id"], session_state["session_start"], debug_logs)

# ==========================================
# 🎨 Gradio 웹 UI 화면 구성 (Gradio 6.0 버전 맞춤 수정)
# ==========================================
with gr.Blocks() as demo:  # 🔥 theme 파라미터 뺐음!
    gr.Markdown("# 🤖 무검열 집단 지성 에이전트")
    gr.Markdown("무검열 모델 3개가 각자의 초안을 내고, 리더가 이를 자아비판 및 종합하여 최적의 답을 찾습니다.\n**상단 로그인 버튼을 누르면 본인 계정의 통신량을 소모하여 API를 호출합니다.**")
 
    with gr.Row():
        # 🔥 LogoutButton 삭제함! (Gradio 6.0부터는 LoginButton이 알아서 다 해줍니다)
        gr.LoginButton(value="Hugging Face로 로그인하여 서비스 이용하기")

    # type="messages"를 지정해주면 OpenAI와 동일한 딕셔너리 구조로 대화가 완벽 호환됩니다.
    # show_copy_button=False를 추가하여 하단 중복 버튼을 숨깁니다.
    # 'type="messages"' 인자를 제거하여 에러를 해결합니다.
    chatbot = gr.Chatbot(height=500, avatar_images=(None, "🤖"))
    
    with gr.Row():
        msg = gr.Textbox(placeholder="질문을 입력하고 엔터 또는 전송 버튼을 누르세요...", scale=9, show_label=False)
        submit_btn = gr.Button("전송", scale=1, variant="primary")
        
    session_state = gr.State({})

    # 엔터키 & 전송 버튼 이벤트 바인딩
    msg.submit(respond, inputs=[msg, chatbot, session_state], outputs=[msg, chatbot, session_state])
    submit_btn.click(respond, inputs=[msg, chatbot, session_state], outputs=[msg, chatbot, session_state])

# 앱 실행
if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft())  # 🔥 경고메시지 안내대로 theme 설정을 launch() 안으로 옮김!