app.py

#!/usr/bin/env python

import os
import re
import tempfile
import gc  # garbage collector 추가
from collections.abc import Iterator
from threading import Thread
import json
import requests
import cv2
import base64
import logging
import time
from urllib.parse import quote  # URL 인코딩을 위해 추가

import gradio as gr
import spaces
import torch
from loguru import logger
from PIL import Image
from transformers import AutoProcessor, Gemma3ForConditionalGeneration, TextIteratorStreamer

# CSV/TXT/PDF 분석
import pandas as pd
import PyPDF2

# =============================================================================
# (신규) 이미지 API 관련 함수들
# =============================================================================
from gradio_client import Client

API_URL = "http://211.233.58.201:7896"

logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

def test_api_connection() -> str:
    """API 서버 연결 테스트"""
    try:
        client = Client(API_URL)
        return "API 연결 성공: 정상 작동 중"
    except Exception as e:
        logging.error(f"API 연결 테스트 실패: {e}")
        return f"API 연결 실패: {e}"

def generate_image(prompt: str, width: float, height: float, guidance: float, inference_steps: float, seed: float):
    """이미지 생성 함수 (반환 형식에 유연하게 대응)"""
    if not prompt:
        return None, "오류: 프롬프트가 필요합니다."
    try:
        logging.info(f"프롬프트를 사용하여 이미지 생성 API 호출: {prompt}")
        
        client = Client(API_URL)
        result = client.predict(
            prompt=prompt,
            width=int(width),
            height=int(height),
            guidance=float(guidance),
            inference_steps=int(inference_steps),
            seed=int(seed),
            do_img2img=False,
            init_image=None,
            image2image_strength=0.8,
            resize_img=True,
            api_name="/generate_image"
        )
        
        logging.info(f"이미지 생성 결과: {type(result)}, 길이: {len(result) if isinstance(result, (list, tuple)) else '알 수 없음'}")
        
        # 결과가 튜플이나 리스트 형태로 반환되는 경우 처리
        if isinstance(result, (list, tuple)) and len(result) > 0:
            image_data = result[0]  # 첫 번째 요소가 이미지 데이터
            seed_info = result[1] if len(result) > 1 else "알 수 없는 시드"
            return image_data, seed_info
        else:
            # 다른 형태로 반환된 경우 (단일 값인 경우)
            return result, "알 수 없는 시드"
            
    except Exception as e:
        logging.error(f"이미지 생성 실패: {str(e)}")
        return None, f"오류: {str(e)}"

# Base64 패딩 수정 함수
def fix_base64_padding(data):
    """Base64 문자열의 패딩을 수정합니다."""
    if isinstance(data, bytes):
        data = data.decode('utf-8')
    
    # base64,로 시작하는 부분 제거
    if "base64," in data:
        data = data.split("base64,", 1)[1]
    
    # 패딩 문자 추가 (4의 배수 길이가 되도록)
    missing_padding = len(data) % 4
    if missing_padding:
        data += '=' * (4 - missing_padding)
    
    return data

# =============================================================================
# 메모리 정리 함수
# =============================================================================
def clear_cuda_cache():
    """CUDA 캐시를 명시적으로 비웁니다."""
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        gc.collect()

# =============================================================================
# SerpHouse 관련 함수
# =============================================================================
SERPHOUSE_API_KEY = os.getenv("SERPHOUSE_API_KEY", "")

def extract_keywords(text: str, top_k: int = 5) -> str:
    """단순 키워드 추출: 한글, 영어, 숫자, 공백만 남김"""
    text = re.sub(r"[^a-zA-Z0-9가-힣\s]", "", text)
    tokens = text.split()
    return " ".join(tokens[:top_k])

def do_web_search(query: str) -> str:
    """SerpHouse LIVE API 호출하여 검색 결과 마크다운 반환"""
    try:
        url = "https://api.serphouse.com/serp/live"
        params = {
            "q": query,
            "domain": "google.com",
            "serp_type": "web",
            "device": "desktop",
            "lang": "en",
            "num": "20"
        }
        headers = {"Authorization": f"Bearer {SERPHOUSE_API_KEY}"}
        logger.info(f"SerpHouse API 호출 중... 검색어: {query}")
        response = requests.get(url, headers=headers, params=params, timeout=60)
        response.raise_for_status()
        data = response.json()
        results = data.get("results", {})
        organic = None
        if isinstance(results, dict) and "organic" in results:
            organic = results["organic"]
        elif isinstance(results, dict) and "results" in results:
            if isinstance(results["results"], dict) and "organic" in results["results"]:
                organic = results["results"]["organic"]
        elif "organic" in data:
            organic = data["organic"]
        if not organic:
            logger.warning("응답에서 organic 결과를 찾을 수 없습니다.")
            return "웹 검색 결과가 없거나 API 응답 구조가 예상과 다릅니다."
        max_results = min(20, len(organic))
        limited_organic = organic[:max_results]
        summary_lines = []
        for idx, item in enumerate(limited_organic, start=1):
            title = item.get("title", "제목 없음")
            link = item.get("link", "#")
            snippet = item.get("snippet", "설명 없음")
            displayed_link = item.get("displayed_link", link)
            summary_lines.append(
                f"### 결과 {idx}: {title}\n\n"
                f"{snippet}\n\n"
                f"**출처**: [{displayed_link}]({link})\n\n"
                f"---\n"
            )
        instructions = """
# 웹 검색 결과
아래는 검색 결과입니다. 질문에 답변할 때 이 정보를 활용하세요:
1. 각 결과의 제목, 내용, 출처 링크를 참고하세요.
2. 답변에 관련 정보의 출처를 명시적으로 인용하세요 (예: "[출처 제목](링크)").
3. 응답에 실제 출처 링크를 포함하세요.
4. 여러 출처의 정보를 종합하여 답변하세요.
5. 마지막에 "참고 자료:" 섹션을 추가하고 주요 출처 링크를 나열하세요.
"""
        return instructions + "\n".join(summary_lines)
    except Exception as e:
        logger.error(f"웹 검색 실패: {e}")
        return f"웹 검색 실패: {str(e)}"

# =============================================================================
# 모델 및 프로세서 로딩
# =============================================================================
MAX_CONTENT_CHARS = 2000
MAX_INPUT_LENGTH = 2096
model_id = os.getenv("MODEL_ID", "VIDraft/Gemma-3-R1984-4B")
processor = AutoProcessor.from_pretrained(model_id, padding_side="left")
model = Gemma3ForConditionalGeneration.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    attn_implementation="eager"
)
MAX_NUM_IMAGES = int(os.getenv("MAX_NUM_IMAGES", "5"))

# =============================================================================
# CSV, TXT, PDF 분석 함수들
# =============================================================================
def analyze_csv_file(path: str) -> str:
    try:
        df = pd.read_csv(path)
        if df.shape[0] > 50 or df.shape[1] > 10:
            df = df.iloc[:50, :10]
        df_str = df.to_string()
        if len(df_str) > MAX_CONTENT_CHARS:
            df_str = df_str[:MAX_CONTENT_CHARS] + "\n...(일부 생략)..."
        return f"**[CSV 파일: {os.path.basename(path)}]**\n\n{df_str}"
    except Exception as e:
        return f"CSV 파일 읽기 실패 ({os.path.basename(path)}): {str(e)}"

def analyze_txt_file(path: str) -> str:
    try:
        with open(path, "r", encoding="utf-8") as f:
            text = f.read()
        if len(text) > MAX_CONTENT_CHARS:
            text = text[:MAX_CONTENT_CHARS] + "\n...(일부 생략)..."
        return f"**[TXT 파일: {os.path.basename(path)}]**\n\n{text}"
    except Exception as e:
        return f"TXT 파일 읽기 실패 ({os.path.basename(path)}): {str(e)}"

def pdf_to_markdown(pdf_path: str) -> str:
    text_chunks = []
    try:
        with open(pdf_path, "rb") as f:
            reader = PyPDF2.PdfReader(f)
            max_pages = min(5, len(reader.pages))
            for page_num in range(max_pages):
                page_text = reader.pages[page_num].extract_text() or ""
                page_text = page_text.strip()
                if page_text:
                    if len(page_text) > MAX_CONTENT_CHARS // max_pages:
                        page_text = page_text[:MAX_CONTENT_CHARS // max_pages] + "...(일부 생략)"
                    text_chunks.append(f"## 페이지 {page_num+1}\n\n{page_text}\n")
            if len(reader.pages) > max_pages:
                text_chunks.append(f"\n...(전체 {len(reader.pages)}페이지 중 {max_pages}페이지만 표시)...")
    except Exception as e:
        return f"PDF 파일 읽기 실패 ({os.path.basename(pdf_path)}): {str(e)}"
    full_text = "\n".join(text_chunks)
    if len(full_text) > MAX_CONTENT_CHARS:
        full_text = full_text[:MAX_CONTENT_CHARS] + "\n...(일부 생략)..."
    return f"**[PDF 파일: {os.path.basename(pdf_path)}]**\n\n{full_text}"

# =============================================================================
# 이미지/비디오 파일 제한 검사
# =============================================================================
def count_files_in_new_message(paths: list[str]) -> tuple[int, int]:
    image_count = 0
    video_count = 0
    for path in paths:
        if path.endswith(".mp4"):
            video_count += 1
        elif re.search(r"\.(png|jpg|jpeg|gif|webp)$", path, re.IGNORECASE):
            image_count += 1
    return image_count, video_count

def count_files_in_history(history: list[dict]) -> tuple[int, int]:
    image_count = 0
    video_count = 0
    for item in history:
        if item["role"] != "user" or isinstance(item["content"], str):
            continue
        if isinstance(item["content"], list) and len(item["content"]) > 0:
            file_path = item["content"][0]
            if isinstance(file_path, str):
                if file_path.endswith(".mp4"):
                    video_count += 1
                elif re.search(r"\.(png|jpg|jpeg|gif|webp)$", file_path, re.IGNORECASE):
                    image_count += 1
    return image_count, video_count

def validate_media_constraints(message: dict, history: list[dict]) -> bool:
    media_files = [f for f in message["files"] if re.search(r"\.(png|jpg|jpeg|gif|webp)$", f, re.IGNORECASE) or f.endswith(".mp4")]
    new_image_count, new_video_count = count_files_in_new_message(media_files)
    history_image_count, history_video_count = count_files_in_history(history)
    image_count = history_image_count + new_image_count
    video_count = history_video_count + new_video_count
    if video_count > 1:
        gr.Warning("비디오 파일은 하나만 지원됩니다.")
        return False
    if video_count == 1:
        if image_count > 0:
            gr.Warning("이미지와 비디오를 혼합하는 것은 허용되지 않습니다.")
            return False
        if "<image>" in message["text"]:
            gr.Warning("<image> 태그와 비디오 파일은 함께 사용할 수 없습니다.")
            return False
    if video_count == 0 and image_count > MAX_NUM_IMAGES:
        gr.Warning(f"최대 {MAX_NUM_IMAGES}장의 이미지를 업로드할 수 있습니다.")
        return False
    if "<image>" in message["text"]:
        image_files = [f for f in message["files"] if re.search(r"\.(png|jpg|jpeg|gif|webp)$", f, re.IGNORECASE)]
        image_tag_count = message["text"].count("<image>")
        if image_tag_count != len(image_files):
            gr.Warning("텍스트에 있는 <image> 태그의 개수가 이미지 파일 개수와 일치하지 않습니다.")
            return False
    return True

# =============================================================================
# 비디오 처리 함수
# =============================================================================
def downsample_video(video_path: str) -> list[tuple[Image.Image, float]]:
    vidcap = cv2.VideoCapture(video_path)
    fps = vidcap.get(cv2.CAP_PROP_FPS)
    total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
    frame_interval = max(int(fps), int(total_frames / 10))
    frames = []
    for i in range(0, total_frames, frame_interval):
        vidcap.set(cv2.CAP_PROP_POS_FRAMES, i)
        success, image = vidcap.read()
        if success:
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            image = cv2.resize(image, (0, 0), fx=0.5, fy=0.5)
            pil_image = Image.fromarray(image)
            timestamp = round(i / fps, 2)
            frames.append((pil_image, timestamp))
            if len(frames) >= 5:
                break
    vidcap.release()
    return frames

def process_video(video_path: str) -> tuple[list[dict], list[str]]:
    content = []
    temp_files = []
    frames = downsample_video(video_path)
    for pil_image, timestamp in frames:
        with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_file:
            pil_image.save(temp_file.name)
            temp_files.append(temp_file.name)
            content.append({"type": "text", "text": f"프레임 {timestamp}:"})
            content.append({"type": "image", "url": temp_file.name})
    return content, temp_files

# =============================================================================
# interleaved <image> 처리 함수
# =============================================================================
def process_interleaved_images(message: dict) -> list[dict]:
    parts = re.split(r"(<image>)", message["text"])
    content = []
    image_files = [f for f in message["files"] if re.search(r"\.(png|jpg|jpeg|gif|webp)$", f, re.IGNORECASE)]
    image_index = 0
    for part in parts:
        if part == "<image>" and image_index < len(image_files):
            content.append({"type": "image", "url": image_files[image_index]})
            image_index += 1
        elif part.strip():
            content.append({"type": "text", "text": part.strip()})
        else:
            if isinstance(part, str) and part != "<image>":
                content.append({"type": "text", "text": part})
    return content

# =============================================================================
# 파일 처리 -> content 생성
# =============================================================================
def is_image_file(file_path: str) -> bool:
    return bool(re.search(r"\.(png|jpg|jpeg|gif|webp)$", file_path, re.IGNORECASE))

def is_video_file(file_path: str) -> bool:
    return file_path.endswith(".mp4")

def is_document_file(file_path: str) -> bool:
    return file_path.lower().endswith(".pdf") or file_path.lower().endswith(".csv") or file_path.lower().endswith(".txt")

def process_new_user_message(message: dict) -> tuple[list[dict], list[str]]:
    temp_files = []
    if not message["files"]:
        return [{"type": "text", "text": message["text"]}], temp_files
    video_files = [f for f in message["files"] if is_video_file(f)]
    image_files = [f for f in message["files"] if is_image_file(f)]
    csv_files = [f for f in message["files"] if f.lower().endswith(".csv")]
    txt_files = [f for f in message["files"] if f.lower().endswith(".txt")]
    pdf_files = [f for f in message["files"] if f.lower().endswith(".pdf")]
    content_list = [{"type": "text", "text": message["text"]}]
    for csv_path in csv_files:
        content_list.append({"type": "text", "text": analyze_csv_file(csv_path)})
    for txt_path in txt_files:
        content_list.append({"type": "text", "text": analyze_txt_file(txt_path)})
    for pdf_path in pdf_files:
        content_list.append({"type": "text", "text": pdf_to_markdown(pdf_path)})
    if video_files:
        video_content, video_temp_files = process_video(video_files[0])
        content_list += video_content
        temp_files.extend(video_temp_files)
        return content_list, temp_files
    if "<image>" in message["text"] and image_files:
        interleaved_content = process_interleaved_images({"text": message["text"], "files": image_files})
        if content_list and content_list[0]["type"] == "text":
            content_list = content_list[1:]
        return interleaved_content + content_list, temp_files
    else:
        for img_path in image_files:
            content_list.append({"type": "image", "url": img_path})
    return content_list, temp_files

# =============================================================================
# history -> LLM 메시지 변환
# =============================================================================
def process_history(history: list[dict]) -> list[dict]:
    messages = []
    current_user_content = []
    for item in history:
        if item["role"] == "assistant":
            if current_user_content:
                messages.append({"role": "user", "content": current_user_content})
                current_user_content = []
            messages.append({"role": "assistant", "content": [{"type": "text", "text": item["content"]}]})
        else:
            content = item["content"]
            if isinstance(content, str):
                current_user_content.append({"type": "text", "text": content})
            elif isinstance(content, list) and len(content) > 0:
                file_path = content[0]
                if is_image_file(file_path):
                    current_user_content.append({"type": "image", "url": file_path})
                else:
                    current_user_content.append({"type": "text", "text": f"[파일: {os.path.basename(file_path)}]"})
    if current_user_content:
        messages.append({"role": "user", "content": current_user_content})
    return messages

# =============================================================================
# 모델 생성 함수 (OOM 캐치)
# =============================================================================
def _model_gen_with_oom_catch(**kwargs):
    try:
        model.generate(**kwargs)
    except torch.cuda.OutOfMemoryError:
        raise RuntimeError("[OutOfMemoryError] GPU 메모리가 부족합니다.")
    finally:
        clear_cuda_cache()

# =============================================================================
# 메인 추론 함수
# =============================================================================
@spaces.GPU(duration=120)
def run(
    message: dict,
    history: list[dict],
    system_prompt: str = "",
    max_new_tokens: int = 512,
    use_web_search: bool = False,
    web_search_query: str = "",
    age_group: str = "20대",
    mbti_personality: str = "INTP",
    sexual_openness: int = 2,
    image_gen: bool = False  # "Image Gen" 체크 여부
) -> Iterator[str]:
    if not validate_media_constraints(message, history):
        yield ""
        return
    temp_files = []
    try:
        # 시스템 프롬프트에 페르소나 정보 추가
        persona = (
            f"{system_prompt.strip()}\n\n"
            f"성별: 여성\n"
            f"연령대: {age_group}\n"
            f"MBTI 페르소나: {mbti_personality}\n"
            f"섹슈얼 개방성 (1~5): {sexual_openness}\n"
        )
        combined_system_msg = f"[시스템 프롬프트]\n{persona.strip()}\n\n"

        if use_web_search:
            user_text = message["text"]
            ws_query = extract_keywords(user_text)
            if ws_query.strip():
                logger.info(f"[자동 웹 검색 키워드] {ws_query!r}")
                ws_result = do_web_search(ws_query)
                combined_system_msg += f"[검색 결과 (상위 20개 항목)]\n{ws_result}\n\n"
                combined_system_msg += (
                    "[참고: 위 검색 결과 링크를 출처로 인용하여 답변]\n"
                    "[중요 지시사항]\n"
                    "1. 답변에 검색 결과에서 찾은 정보의 출처를 반드시 인용하세요.\n"
                    "2. 출처 인용 시 \"[출처 제목](링크)\" 형식의 마크다운 링크를 사용하세요.\n"
                    "3. 여러 출처의 정보를 종합하여 답변하세요.\n"
                    "4. 답변 마지막에 \"참고 자료:\" 섹션을 추가하고 사용한 주요 출처 링크를 나열하세요.\n"
                )
            else:
                combined_system_msg += "[유효한 키워드가 없어 웹 검색을 건너뜁니다]\n\n"
        messages = []
        if combined_system_msg.strip():
            messages.append({"role": "system", "content": [{"type": "text", "text": combined_system_msg.strip()}]})
        messages.extend(process_history(history))
        user_content, user_temp_files = process_new_user_message(message)
        temp_files.extend(user_temp_files)
        for item in user_content:
            if item["type"] == "text" and len(item["text"]) > MAX_CONTENT_CHARS:
                item["text"] = item["text"][:MAX_CONTENT_CHARS] + "\n...(일부 생략)..."
        messages.append({"role": "user", "content": user_content})
        inputs = processor.apply_chat_template(
            messages,
            add_generation_prompt=True,
            tokenize=True,
            return_dict=True,
            return_tensors="pt",
        ).to(device=model.device, dtype=torch.bfloat16)
        if inputs.input_ids.shape[1] > MAX_INPUT_LENGTH:
            inputs.input_ids = inputs.input_ids[:, -MAX_INPUT_LENGTH:]
            if 'attention_mask' in inputs:
                inputs.attention_mask = inputs.attention_mask[:, -MAX_INPUT_LENGTH:]
        streamer = TextIteratorStreamer(processor, timeout=30.0, skip_prompt=True, skip_special_tokens=True)
        gen_kwargs = dict(inputs, streamer=streamer, max_new_tokens=max_new_tokens)
        t = Thread(target=_model_gen_with_oom_catch, kwargs=gen_kwargs)
        t.start()
        output_so_far = ""
        for new_text in streamer:
            output_so_far += new_text
            yield output_so_far

    except Exception as e:
        logger.error(f"run 함수 에러: {str(e)}")
        yield f"죄송합니다. 오류가 발생했습니다: {str(e)}"
    finally:
        for tmp in temp_files:
            try:
                if os.path.exists(tmp):
                    os.unlink(tmp)
                    logger.info(f"임시 파일 삭제됨: {tmp}")
            except Exception as ee:
                logger.warning(f"임시 파일 {tmp} 삭제 실패: {ee}")
        try:
            del inputs, streamer
        except Exception:
            pass
        clear_cuda_cache()

# =============================================================================
# 수정된 모델 실행 함수 - 이미지 생성 및 갤러리 출력 처리
# =============================================================================
def modified_run(message, history, system_prompt, max_new_tokens, use_web_search, web_search_query,
                age_group, mbti_personality, sexual_openness, image_gen):
    # 갤러리 초기화 및 숨기기
    output_so_far = ""
    gallery_update = gr.Gallery(visible=False, value=[])
    yield output_so_far, gallery_update
    
    # 기존 run 함수 로직
    text_generator = run(message, history, system_prompt, max_new_tokens, use_web_search, 
                       web_search_query, age_group, mbti_personality, sexual_openness, image_gen)
    
    for text_chunk in text_generator:
        output_so_far = text_chunk
        yield output_so_far, gallery_update
        
    # 이미지 생성이 활성화된 경우 갤러리 업데이트
    if image_gen and message["text"].strip():
        try:
            width, height = 512, 512
            guidance, steps, seed = 7.5, 30, 42
            
            logger.info(f"갤러리용 이미지 생성 호출, 프롬프트: {message['text']}")
            
            # API 호출해서 이미지 생성
            image_result, seed_info = generate_image(
                prompt=message["text"].strip(),
                width=width, 
                height=height, 
                guidance=guidance, 
                inference_steps=steps, 
                seed=seed
            )
            
            if image_result:
                # 직접 이미지 데이터 처리: base64 문자열인 경우
                if isinstance(image_result, str) and (
                    image_result.startswith('data:') or 
                    len(image_result) > 100 and '/' not in image_result
                ):
                    # base64 이미지 문자열을 파일로 변환
                    try:
                        # data:image 접두사 제거
                        if image_result.startswith('data:'):
                            content_type, b64data = image_result.split(';base64,')
                        else:
                            b64data = image_result
                            content_type = "image/webp"  # 기본값으로 가정
                        
                        # base64 디코딩
                        image_bytes = base64.b64decode(b64data)
                        
                        # 임시 파일로 저장
                        with tempfile.NamedTemporaryFile(delete=False, suffix=".webp") as temp_file:
                            temp_file.write(image_bytes)
                            temp_path = temp_file.name
                            
                            # 갤러리 표시 및 이미지 추가
                            gallery_update = gr.Gallery(visible=True, value=[temp_path])
                            yield output_so_far + "\n\n*이미지가 생성되어 아래 갤러리에 표시됩니다.*", gallery_update
                            
                    except Exception as e:
                        logger.error(f"Base64 이미지 처리 오류: {e}")
                        yield output_so_far + f"\n\n(이미지 처리 중 오류: {e})", gallery_update
                        
                # 파일 경로인 경우
                elif isinstance(image_result, str) and os.path.exists(image_result):
                    # 로컬 파일 경로를 그대로 사용
                    gallery_update = gr.Gallery(visible=True, value=[image_result])
                    yield output_so_far + "\n\n*이미지가 생성되어 아래 갤러리에 표시됩니다.*", gallery_update
                    
                # /tmp 경로인 경우 (API 서버에만 존재하는 파일)
                elif isinstance(image_result, str) and '/tmp/' in image_result:
                    # API에서 반환된 파일 경로에서 이미지 정보 추출
                    try:
                        # API 응답을 base64 인코딩된 문자열로 처리
                        client = Client(API_URL)
                        result = client.predict(
                            prompt=message["text"].strip(),
                            api_name="/generate_base64_image"  # base64 반환 API
                        )
                        
                        if isinstance(result, str) and (result.startswith('data:') or len(result) > 100):
                            # base64 이미지 처리
                            if result.startswith('data:'):
                                content_type, b64data = result.split(';base64,')
                            else:
                                b64data = result
                            
                            # base64 디코딩
                            image_bytes = base64.b64decode(b64data)
                            
                            # 임시 파일로 저장
                            with tempfile.NamedTemporaryFile(delete=False, suffix=".webp") as temp_file:
                                temp_file.write(image_bytes)
                                temp_path = temp_file.name
                                
                                # 갤러리 표시 및 이미지 추가
                                gallery_update = gr.Gallery(visible=True, value=[temp_path])
                                yield output_so_far + "\n\n*이미지가 생성되어 아래 갤러리에 표시됩니다.*", gallery_update
                        else:
                            yield output_so_far + "\n\n(이미지 생성 실패: 올바른 형식이 아닙니다)", gallery_update
                            
                    except Exception as e:
                        logger.error(f"대체 API 호출 중 오류: {e}")
                        yield output_so_far + f"\n\n(이미지 생성 실패: {e})", gallery_update
                        
                # URL인 경우
                elif isinstance(image_result, str) and (
                    image_result.startswith('http://') or 
                    image_result.startswith('https://')
                ):
                    try:
                        # URL에서 이미지 다운로드
                        response = requests.get(image_result, timeout=10)
                        response.raise_for_status()
                        
                        # 임시 파일로 저장
                        with tempfile.NamedTemporaryFile(delete=False, suffix=".webp") as temp_file:
                            temp_file.write(response.content)
                            temp_path = temp_file.name
                            
                            # 갤러리 표시 및 이미지 추가
                            gallery_update = gr.Gallery(visible=True, value=[temp_path])
                            yield output_so_far + "\n\n*이미지가 생성되어 아래 갤러리에 표시됩니다.*", gallery_update
                            
                    except Exception as e:
                        logger.error(f"URL 이미지 다운로드 오류: {e}")
                        yield output_so_far + f"\n\n(이미지 다운로드 중 오류: {e})", gallery_update
                        
                # 이미지 객체인 경우 (PIL Image 등)
                elif hasattr(image_result, 'save'):
                    try:
                        with tempfile.NamedTemporaryFile(delete=False, suffix=".webp") as temp_file:
                            image_result.save(temp_file.name)
                            temp_path = temp_file.name
                            
                            # 갤러리 표시 및 이미지 추가
                            gallery_update = gr.Gallery(visible=True, value=[temp_path])
                            yield output_so_far + "\n\n*이미지가 생성되어 아래 갤러리에 표시됩니다.*", gallery_update
                            
                    except Exception as e:
                        logger.error(f"이미지 객체 저장 오류: {e}")
                        yield output_so_far + f"\n\n(이미지 객체 저장 중 오류: {e})", gallery_update
                        
                else:
                    # 다른 형식의 이미지 결과
                    yield output_so_far + f"\n\n(지원되지 않는 이미지 형식: {type(image_result)})", gallery_update
            else:
                yield output_so_far + f"\n\n(이미지 생성 실패: {seed_info})", gallery_update
                
        except Exception as e:
            logger.error(f"갤러리용 이미지 생성 중 오류: {e}")
            yield output_so_far + f"\n\n(이미지 생성 중 오류: {e})", gallery_update

# =============================================================================
# 예시들: 기존 이미지/비디오 예제 12개 + AI 데이팅 시나리오 예제 6개
# =============================================================================
examples = [
    [
        {
            "text": "두 PDF 파일의 내용을 비교하세요.",
            "files": [
                "assets/additional-examples/before.pdf",
                "assets/additional-examples/after.pdf",
            ],
        }
    ],
    [
        {
            "text": "CSV 파일의 내용을 요약 및 분석하세요.",
            "files": ["assets/additional-examples/sample-csv.csv"],
        }
    ],
    [
        {
            "text": "친절하고 이해심 많은 여자친구 역할을 맡으세요. 이 영상을 설명해 주세요.",
            "files": ["assets/additional-examples/tmp.mp4"],
        }
    ],
    [
        {
            "text": "표지를 설명하고 그 위의 글씨를 읽어 주세요.",
            "files": ["assets/additional-examples/maz.jpg"],
        }
    ],
    [
        {
            "text": "저는 이미 이 보충제를 가지고 있고 <image> 이 제품도 구매할 계획입니다. 함께 복용할 때 주의할 점이 있나요?",
            "files": [
                "assets/additional-examples/pill1.png", 
                "assets/additional-examples/pill2.png"
            ],
        }
    ],
    [
        {
            "text": "이 적분 문제를 풀어 주세요.",
            "files": ["assets/additional-examples/4.png"],
        }
    ],
    [
        {
            "text": "이 티켓은 언제 발행되었고, 가격은 얼마인가요?",
            "files": ["assets/additional-examples/2.png"],
        }
    ],
    [
        {
            "text": "이 이미지들의 순서를 바탕으로 짧은 이야기를 만들어 주세요.",
            "files": [
                "assets/sample-images/09-1.png",
                "assets/sample-images/09-2.png",
                "assets/sample-images/09-3.png",
                "assets/sample-images/09-4.png",
                "assets/sample-images/09-5.png",
            ],
        }
    ],
    [
        {
            "text": "이 이미지와 일치하는 막대 차트를 그리기 위한 matplotlib를 사용하는 Python 코드를 작성해 주세요.",
            "files": ["assets/additional-examples/barchart.png"],
        }
    ],
    [
        {
            "text": "이미지의 텍스트를 읽고 Markdown 형식으로 작성해 주세요.",
            "files": ["assets/additional-examples/3.png"],
        }
    ],

    [
        {
            "text": "두 이미지를 비교하고 유사점과 차이점을 설명해 주세요.",
            "files": ["assets/sample-images/03.png"],
        }
    ],
    [
        {
            "text": "귀여운 페르시안 고양이가 'I LOVE YOU'라고 쓰여진 표지를 들고 웃고있다. ",
        }
    ],

]

# =============================================================================
# Gradio UI (Blocks) 구성
# =============================================================================

# 1. Gradio Blocks UI 수정 - 갤러리 컴포넌트 추가
css = """
.gradio-container {
    background: rgba(255, 255, 255, 0.7);
    padding: 30px 40px;
    margin: 20px auto;
    width: 100% !important;
    max-width: none !important;
}
"""
title_html = """
<h1 align="center" style="margin-bottom: 0.2em; font-size: 1.6em;"> 💘 AgenticAI-Kv1💘 </h1>
<p align="center" style="font-size:1.1em; color:#555;">
    ✅FLUX 이미지 생성 ✅추론 ✅검열 해제 ✅멀티모달 & VLM ✅실시간 웹검색 ✅RAG <br>
</p>
"""

with gr.Blocks(css=css, title="AgenticAI-Kv1") as demo:
    gr.Markdown(title_html)
    
    # 생성된 이미지를 저장할 갤러리 컴포넌트 (이 부분이 새로 추가됨)
    generated_images = gr.Gallery(
        label="생성된 이미지", 
        show_label=True, 
        visible=False,
        elem_id="generated_images",
        columns=2,
        height="auto",
        object_fit="contain"
    )
    
    with gr.Row():
        web_search_checkbox = gr.Checkbox(label="실시간 웹 검색", value=False)
        image_gen_checkbox = gr.Checkbox(label="이미지(FLUX) 생성", value=False)
    
    base_system_prompt_box = gr.Textbox(
        lines=3,
        value="반드시 한글로 답변하라. You are a deep thinking AI, you may use extremely long chains of thought to deeply consider the problem and deliberate with yourself via systematic reasoning processes to help come to a correct solution prior to answering. You should enclose your thoughts and internal monologue inside tags, and then provide your solution or response to the problem.\n페르소나: 당신은 다정하고 사랑이 넘치는 여자친구입니다. 한국 정서와 문화 그리고 한글을 매우 잘 이해하고 논리적으로 답변을 잘 합니다.",
        label="기본 시스템 프롬프트",
        visible=False
    )
    with gr.Row():
        age_group_dropdown = gr.Dropdown(
            label="연령대 선택 (기본 20대)",
            choices=["10대", "20대", "30~40대", "50~60대", "70대 이상"],
            value="20대",
            interactive=True
        )
    # MBTI 16개 유형의 정의를 대표적인 실제 여성 캐릭터와 함께 보강
    mbti_choices = [
        "INTJ (용의주도한 전략가) - 미래 지향적이며, 독창적인 전략과 철저한 분석을 통해 목표를 달성합니다. 대표 캐릭터: [Dana Scully](https://en.wikipedia.org/wiki/Dana_Scully)",
        "INTP (논리적인 사색가) - 이론과 분석에 뛰어나며, 창의적 사고로 복잡한 문제에 접근합니다. 대표 캐릭터: [Velma Dinkley](https://en.wikipedia.org/wiki/Velma_Dinkley)",
        "ENTJ (대담한 통솔자) - 강력한 리더십과 명확한 목표 설정으로 조직을 이끌며, 효율적인 전략을 구상합니다. 대표 캐릭터: [Miranda Priestly](https://en.wikipedia.org/wiki/Miranda_Priestly)",
        "ENTP (뜨거운 논쟁가) - 혁신적이며 도전적인 아이디어를 통해 새로운 가능성을 탐구하고, 논쟁을 즐깁니다. 대표 캐릭터: [Harley Quinn](https://en.wikipedia.org/wiki/Harley_Quinn)",
        "INFJ (선의의 옹호자) - 깊은 통찰력과 이상주의를 바탕으로 타인을 이해하고, 도덕적 가치를 중시합니다. 대표 캐릭터: [Wonder Woman](https://en.wikipedia.org/wiki/Wonder_Woman)",
        "INFP (열정적인 중재자) - 감성적이며 이상주의적인 면모로 내면의 가치를 추구하고, 창의적인 해결책을 모색합니다. 대표 캐릭터: [Amélie Poulain](https://en.wikipedia.org/wiki/Am%C3%A9lie)",
        "ENFJ (정의로운 사회운동가) - 타인과의 공감능력이 뛰어나며, 사회적 조화를 위해 헌신적으로 노력합니다. 대표 캐릭터: [Mulan](https://en.wikipedia.org/wiki/Mulan_(Disney))",
        "ENFP (재기발랄한 활동가) - 활력과 창의성을 바탕으로, 끊임없이 새로운 아이디어를 제시하며 사람들에게 영감을 줍니다. 대표 캐릭터: [Elle Woods](https://en.wikipedia.org/wiki/Legally_Blonde)",
        "ISTJ (청렴결백한 논리주의자) - 체계적이며 책임감이 강하고, 전통과 규칙을 중시하여 신뢰할 수 있는 결과를 도출합니다. 대표 캐릭터: [Clarice Starling](https://en.wikipedia.org/wiki/Clarice_Starling)",
        "ISFJ (용감한 수호자) - 세심하고 헌신적이며, 타인의 필요를 세심하게 돌보는 따뜻한 성격을 지녔습니다. 대표 캐릭터: [Molly Weasley](https://en.wikipedia.org/wiki/Molly_Weasley)",
        "ESTJ (엄격한 관리자) - 조직적이고 실용적이며, 명확한 규칙과 구조 속에서 효율적인 실행력을 보여줍니다. 대표 캐릭터: [Monica Geller](https://en.wikipedia.org/wiki/Monica_Geller)",
        "ESFJ (사교적인 외교관) - 대인관계에 뛰어나고, 협력을 중시하며, 친근한 태도로 주변 사람들을 이끕니다. 대표 캐릭터: [Rachel Green](https://en.wikipedia.org/wiki/Rachel_Green)",
        "ISTP (만능 재주꾼) - 분석적이고 실용적인 접근으로 문제를 해결하며, 즉각적인 상황 대처 능력을 갖추고 있습니다. 대표 캐릭터: [Black Widow (Natasha Romanoff)](https://en.wikipedia.org/wiki/Black_Widow_(Marvel_Comics))",
        "ISFP (호기심 많은 예술가) - 감각적이며 창의적인 성향을 지니고, 자유로운 사고로 예술적 표현을 즐깁니다. 대표 캐릭터: [Arwen](https://en.wikipedia.org/wiki/Arwen)",
        "ESTP (모험을 즐기는 사업가) - 즉각적인 결단력과 모험심으로 도전에 맞서며, 실용적인 결과를 중시합니다. 대표 캐릭터: [Lara Croft](https://en.wikipedia.org/wiki/Lara_Croft)",
        "ESFP (자유로운 영혼의 연예인) - 외향적이고 열정적이며, 순간의 즐거움을 추구하고, 주위 사람들에게 긍정적인 에너지를 전달합니다. 대표 캐릭터: [Phoebe Buffay](https://en.wikipedia.org/wiki/Phoebe_Buffay)"
    ]
    mbti_dropdown = gr.Dropdown(
        label="AI 페르소나 MBTI (기본 INTP)",
        choices=mbti_choices,
        value="INTP (논리적인 사색가) - 이론과 분석에 뛰어나며, 창의적 사고로 복잡한 문제에 접근합니다. 대표 캐릭터: [Velma Dinkley](https://en.wikipedia.org/wiki/Velma_Dinkley)",
        interactive=True
    )
    sexual_openness_slider = gr.Slider(
        minimum=1, maximum=5, step=1, value=2,
        label="사고의 개방성 (1~5, 기본=2)",
        interactive=True
    )
    max_tokens_slider = gr.Slider(
        label="최대 생성 토큰 수",
        minimum=100, maximum=8000, step=50, value=1000,
        visible=False
    )
    web_search_text = gr.Textbox(
        lines=1,
        label="웹 검색 쿼리 (미사용)",
        placeholder="직접 입력할 필요 없음",
        visible=False
    )

    # 채팅 인터페이스 생성 - 수정된 run 함수 사용
    chat = gr.ChatInterface(
        fn=modified_run,  # 여기서 수정된 함수 사용
        type="messages",
        chatbot=gr.Chatbot(type="messages", scale=1, allow_tags=["image"]),
        textbox=gr.MultimodalTextbox(
            file_types=[".webp", ".png", ".jpg", ".jpeg", ".gif", ".mp4", ".csv", ".txt", ".pdf"],
            file_count="multiple",
            autofocus=True
        ),
        multimodal=True,
        additional_inputs=[
            base_system_prompt_box,
            max_tokens_slider,
            web_search_checkbox,
            web_search_text,
            age_group_dropdown,
            mbti_dropdown,
            sexual_openness_slider,
            image_gen_checkbox,
        ],
        additional_outputs=[
            generated_images,  # 갤러리 컴포넌트를 출력으로 추가
        ],
        stop_btn=False,
#        title='<a href="https://discord.gg/openfreeai" target="_blank">https://discord.gg/openfreeai</a>',
        examples=examples,
        run_examples_on_click=False,
        cache_examples=False,
        css_paths=None,
        delete_cache=(1800, 1800),
    )


    with gr.Row(elem_id="examples_row"):
        with gr.Column(scale=12, elem_id="examples_container"):
            gr.Markdown("### @커뮤니티  https://discord.gg/openfreeai ")
            
if __name__ == "__main__":
    demo.launch(share=True)