app.py

import time
import gradio as gr
import random
import os
import threading
import pandas as pd
from datasets import load_dataset

# CSV 파일 경로와 동시 접근을 위한 Lock 선언
DATA_FILE = "global_data.csv"
data_lock = threading.Lock()

def initialize_global_data():
    """
    DATA_FILE이 존재하지 않으면, gaeunseo/Taskmaster_sample_data 데이터셋을 로드하여 DataFrame으로 변환한 후 CSV 파일로 저장합니다.
    이미 파일이 있으면 파일에서 데이터를 읽어 DataFrame을 반환합니다.
    """
    if not os.path.exists(DATA_FILE):
        ds = load_dataset("gaeunseo/Taskmaster_sample_data", split="train")
        data = ds.to_pandas()
        # 필요한 컬럼이 없으면 추가합니다.
        if "used" not in data.columns:
            data["used"] = False
        if "overlapping" not in data.columns:
            data["overlapping"] = ""
        if "text" not in data.columns:
            data["text"] = ""
        data.to_csv(DATA_FILE, index=False)
        return data
    else:
        with data_lock:
            df = pd.read_csv(DATA_FILE)
        return df

def load_global_data():
    """CSV 파일에서 global_data DataFrame을 읽어옵니다."""
    with data_lock:
        df = pd.read_csv(DATA_FILE)
    return df

def save_global_data(df):
    """DataFrame을 CSV 파일에 저장합니다."""
    with data_lock:
        df.to_csv(DATA_FILE, index=False)

# CSV 파일에 저장된 global_data 초기화
global_data = initialize_global_data()

def get_random_row_from_dataset():
    """
    CSV 파일에 저장된 global_data에서,
      1. conversation_id별로 그룹화하고,
      2. 각 그룹에서 모든 행의 used 컬럼이 False이며, 그룹 내에 overlapping 컬럼이 "TT"인 행이 존재하는 그룹만 valid로 간주합니다.
    valid한 그룹들 중 랜덤하게 하나의 그룹을 선택한 후,
      - 해당 그룹의 모든 행의 used 값을 True로 업데이트(즉, 전체 그룹을 할당)하고 CSV 파일에 저장합니다.
      - 선택된 그룹 내에서 overlapping 컬럼이 "TT"인 행(여러 개라면 첫 번째)을 반환합니다.
    """
    global global_data
    global_data = load_global_data()  # 최신 데이터 로드
    groups = global_data.groupby('conversation_id')
    valid_groups = []
    for cid, group in groups:
        # 모든 행의 used 값이 False이고, 그룹 내에 overlapping 값이 "TT"인 행이 있는 그룹 필터링
        if group['used'].apply(lambda x: bool(x) == False).all() and (group['overlapping'] == "TT").any():
            valid_groups.append((cid, group))
    if not valid_groups:
        return None
    chosen_cid, chosen_group = random.choice(valid_groups)
    # 전체 그룹의 used 값을 True로 업데이트
    global_data.loc[global_data['conversation_id'] == chosen_cid, 'used'] = True
    save_global_data(global_data)
    # 선택된 그룹 내에서 overlapping 컬럼이 "TT"인 행을 반환 (여러 개라면 첫 번째)
    chosen_rows = chosen_group[chosen_group['overlapping'] == "TT"]
    if chosen_rows.empty:
        return None
    chosen_row = chosen_rows.iloc[0]
    return chosen_row.to_dict()

# --- 초기 대화 불러오기 ---
# 데이터셋의 text 컬럼은 "[turn]"을 기준으로 대화가 구분되어 있다고 가정합니다.
row = get_random_row_from_dataset()
if row is None:
    human_message = "No valid conversation available."
    ai_message = "No valid conversation available."
else:
    raw_text = row['text']
    human_message = raw_text.split("[turn]")[0].strip()
    ai_message = raw_text.split("[turn]")[1].strip()

#############################################
# 채팅 인터페이스 관련 함수 (말풍선, 타이핑 효과, 편집 기능)
#############################################

def get_initial_human_html():
    """
    페이지 로드 시, 빈 Human 말풍선과 오른쪽 🧑 이모티콘을 포함한 초기 HTML 반환
    """
    wrapper_start = (
        """<div class="human-wrapper" style="display: flex; align-items: flex-end; justify-content: flex-end; gap: 5px; width: 100%;">"""
    )
    bubble_start = """<div id="human_message" class="speech-bubble human">"""
    bubble_end = "</div>"
    emoji_html = "<div class='emoji'>🧑</div>"
    wrapper_end = "</div>"
    return wrapper_start + bubble_start + bubble_end + emoji_html + wrapper_end

def stream_human_message():
    """
    Start Typing 버튼 클릭 시, 전역 변수 human_message의 내용을 한 글자씩 타이핑 효과로 출력합니다.
    이전 상태(✂️ 아이콘, 회색 처리 등)는 리셋됩니다.
    """
    bubble_content = ""
    wrapper_start = (
        """<div class="human-wrapper" style="display: flex; align-items: flex-end; justify-content: flex-end; gap: 5px; width: 100%;">"""
    )
    bubble_start = """<div id="human_message" class="speech-bubble human">"""
    bubble_end = "</div>"
    emoji_html = "<div class='emoji'>🧑</div>"
    wrapper_end = "</div>"

    # 초기 상태: 빈 말풍선과 이모티콘
    yield wrapper_start + bubble_start + bubble_end + emoji_html + wrapper_end

    # human_message를 한 글자씩 추가 (타이핑 효과)
    for i, ch in enumerate(human_message):
         bubble_content += f"<span data-index='{i}'>{ch}</span>"
         current_html = wrapper_start + bubble_start + bubble_content + bubble_end + emoji_html + wrapper_end
         yield current_html
         time.sleep(0.05)

def submit_edit(edited_text):
    """
    Submit 버튼 클릭 시 호출되는 함수.
    1. 편집된 human 메시지(✂️ 앞부분)를 새 행으로 global_data에 추가합니다.
    2. get_random_row_from_dataset()을 통해 새로운 대화를 가져오고, 전역 변수 human_message와 ai_message를 업데이트합니다.
    3. 초기 상태의 human 말풍선과 ai 말풍선 HTML을 반환하여 인터페이스를 리셋합니다.
    """
    global human_message, ai_message
    data = load_global_data()
    new_row = {
         "conversation_id": "edited_" + str(random.randint(1000, 9999)),
         "used": False,
         "overlapping": "",
         "text": edited_text,
         "human_message": edited_text,
         "ai_message": ""
    }
    new_df = pd.DataFrame([new_row])
    data = pd.concat([data, new_df], ignore_index=True)
    save_global_data(data)
    
    new_row_data = get_random_row_from_dataset()
    if new_row_data is None:
         human_message = "No valid conversation available."
         ai_message = "No valid conversation available."
    else:
         raw_text = new_row_data['text']
         human_message = raw_text.split("[turn]")[0].strip()
         ai_message = raw_text.split("[turn]")[1].strip()
    
    new_human_html = get_initial_human_html()
    new_ai_html = f"""
        <div class="ai-wrapper" style="display: flex; align-items: flex-end; justify-content: flex-start; gap: 5px; width: 100%;">
          <div class="emoji">🤖</div>
          <div id="ai_message" class="speech-bubble ai">{ai_message}</div>
        </div>
    """
    return new_human_html, new_ai_html

#############################################
# Gradio 인터페이스 구성
#############################################

with gr.Blocks() as demo:
    # (A) 페이지 상단 스크립트: Human 말풍선 내의 각 <span data-index="...">를 클릭하면,
    #     해당 위치에 ✂️ 아이콘이 삽입되고, 그 이후 텍스트가 회색으로 변경됩니다.
    gr.HTML(
        """
        <script>
        document.addEventListener("click", function(event) {
            if (event.target && event.target.matches("div.speech-bubble.human span[data-index]")) {
                var span = event.target;
                var container = span.closest("div.speech-bubble.human");
                var oldScissors = container.querySelectorAll("span.scissor");
                oldScissors.forEach(function(s) { s.remove(); });
                var spans = container.querySelectorAll("span[data-index]");
                spans.forEach(function(s) { s.style.color = ''; });
                var scissor = document.createElement('span');
                scissor.textContent = '✂️';
                scissor.classList.add("scissor");
                container.insertBefore(scissor, span.nextSibling);
                var cutIndex = parseInt(span.getAttribute("data-index"));
                spans.forEach(function(s) {
                    var idx = parseInt(s.getAttribute("data-index"));
                    if (idx > cutIndex) {
                        s.style.color = "grey";
                    }
                });
            }
        });
        </script>
        """
    )
    
    # (B) 추가 스크립트: Submit 버튼 클릭 시, human_message div의 innerText에서 "✂️"를 기준으로 편집된 텍스트(앞부분)를 숨김 텍스트박스에 업데이트
    gr.HTML(
        """
        <script>
        document.addEventListener("DOMContentLoaded", function() {
            var submitBtn = document.getElementById("submit_btn");
            if(submitBtn){
                submitBtn.addEventListener("click", function(){
                    var humanDiv = document.getElementById("human_message");
                    if(humanDiv){
                        var edited_text = humanDiv.innerText.split("✂️")[0];
                        document.getElementById("edited_text_input").value = edited_text;
                    }
                });
            }
        });
        </script>
        """
    )
    
    # (C) CSS 스타일
    gr.HTML(
        """
        <style>
        .chat-container {
          display: flex;
          flex-direction: column;
          gap: 10px;
          width: 100%;
        }
        .speech-bubble {
          position: relative;
          padding: 10px 15px;
          border-radius: 15px;
          max-width: 70%;
          font-family: sans-serif;
          font-size: 16px;
          line-height: 1.4;
        }
        .human {
          background: #d0f0d0;
          margin-right: 10px;
        }
        .human::after {
          content: "";
          position: absolute;
          right: -10px;
          top: 10px;
          border-width: 10px 0 10px 10px;
          border-style: solid;
          border-color: transparent transparent transparent #d0f0d0;
        }
        .ai {
          background: #e0e0e0;
          margin-left: 10px;
        }
        .ai::after {
          content: "";
          position: absolute;
          left: -10px;
          top: 10px;
          border-width: 10px 10px 10px 0;
          border-style: solid;
          border-color: transparent #e0e0e0 transparent transparent;
        }
        .emoji {
          font-size: 24px;
          line-height: 1;
        }
        </style>
        """
    )
    
    gr.Markdown("## Chat Interface")
    
    with gr.Column(elem_classes="chat-container"):
        # Human 말풍선 (초기: 빈 메시지 + 🧑 이모티콘)
        human_bubble = gr.HTML(get_initial_human_html())
        # AI 말풍선 (왼쪽: 🤖 이모티콘 + 메시지)
        ai_html = f"""
        <div class="ai-wrapper" style="display: flex; align-items: flex-end; justify-content: flex-start; gap: 5px; width: 100%;">
          <div class="emoji">🤖</div>
          <div id="ai_message" class="speech-bubble ai">{ai_message}</div>
        </div>
        """
        ai_bubble = gr.HTML(ai_html)
    
    # 숨김 텍스트박스 (편집된 텍스트 저장용)
    edited_text_input = gr.Textbox(visible=False, elem_id="edited_text_input")
    
    # 버튼 영역: Start Typing과 Submit 버튼을 같은 행에 배치
    with gr.Row():
        start_button = gr.Button("Start Typing")
        submit_button = gr.Button("Submit", elem_id="submit_btn")
    
    # 버튼 이벤트 연결
    start_button.click(fn=stream_human_message, outputs=human_bubble)
    submit_button.click(fn=submit_edit, inputs=edited_text_input, outputs=[human_bubble, ai_bubble])

demo.launch()