modules/gemini_handler.py

import os
import json
import requests
import time
from dotenv import load_dotenv

# 환경 변수 로드 시도
# HuggingFace Spaces에서는 환경 변수가 자동으로 로드됨
# 로컬 개발 환경에서는 .env 파일 사용
load_dotenv()

# Gemini API 키 및 기본 URL
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY", "")
GEMINI_API_URL = "https://generativelanguage.googleapis.com/v1/models/gemini-pro:generateContent"

def gemini_query(prompt, retry_attempts=3, retry_delay=2):
    """
    Gemini API에 텍스트 생성 요청을 보냅니다.
    
    Args:
        prompt: 프롬프트 텍스트
        retry_attempts: 재시도 횟수
        retry_delay: 재시도 대기 시간 (초)
    
    Returns:
        생성된 텍스트 응답
    
    Raises:
        Exception: API 요청 실패 시
    """
    if not GEMINI_API_KEY:
        # 허깅페이스 환경에서 API 키가 없을 경우 대체 메시지
        return "API 키가 설정되지 않아 응답을 생성할 수 없습니다. HuggingFace Spaces의 Settings에서 GEMINI_API_KEY를 시크릿으로 추가해주세요."
    
    headers = {
        "Content-Type": "application/json",
        "x-goog-api-key": GEMINI_API_KEY
    }
    
    data = {
        "contents": [{
            "parts": [{
                "text": prompt
            }]
        }],
        "generationConfig": {
            "temperature": 0.7,
            "topP": 0.95,
            "topK": 40,
            "maxOutputTokens": 2048
        }
    }
    
    # 재시도 로직
    attempt = 0
    while attempt < retry_attempts:
        try:
            response = requests.post(
                GEMINI_API_URL,
                headers=headers,
                json=data,
                timeout=60  # 타임아웃 설정
            )
            
            response.raise_for_status()  # HTTP 오류 체크
            
            result = response.json()
            
            # 응답 파싱
            if "candidates" in result and result["candidates"]:
                return result["candidates"][0]["content"]["parts"][0]["text"]
            else:
                raise Exception("유효한 응답을 받지 못했습니다.")
                
        except requests.exceptions.RequestException as e:
            attempt += 1
            if attempt < retry_attempts:
                print(f"API 요청 실패, {retry_delay}초 후 재시도 ({attempt}/{retry_attempts}): {str(e)}")
                time.sleep(retry_delay)
            else:
                raise Exception(f"Gemini API 요청이 {retry_attempts}회 실패했습니다: {str(e)}")
    
    raise Exception("알 수 없는 오류가 발생했습니다.")

def get_persona_enhancement(persona_data):
    """
    LLM을 통해 페르소나를 강화합니다.
    
    Args:
        persona_data: 페르소나 기본 정보
    
    Returns:
        강화된 페르소나 데이터
    """
    # 기본 정보 추출
    name = persona_data.get("기본정보", {}).get("이름", "")
    object_type = persona_data.get("기본정보", {}).get("유형", "")
    description = persona_data.get("기본정보", {}).get("설명", "")
    
    # 성격 특성 추출
    traits = []
    for trait, value in persona_data.get("성격특성", {}).items():
        level = "높은" if value >= 70 else "중간" if value >= 40 else "낮은"
        traits.append(f"{trait}: {level} ({value}/100)")
    
    # 배경 이야기 및 경험
    backstory = persona_data.get("배경이야기", "")
    experiences = persona_data.get("경험", [])
    
    # 프롬프트 구성
    prompt = f"""
당신은 물체에 인격을 부여하는 전문가입니다. 다음 정보를 기반으로 매력적이고 개성 있는 페르소나를 강화해주세요.

## 기본 정보
- 이름: {name}
- 유형: {object_type}
- 설명: {description}

## 성격 특성
{', '.join(traits)}

## 매력적 결함
{', '.join(persona_data.get("매력적결함", []))}

## 소통 방식
- 대화 스타일: {persona_data.get("소통방식", "")}
- 유머 스타일: {persona_data.get("유머스타일", "")}
- 말투 패턴: {persona_data.get("말투패턴", "")}

## 관계 성향
- 애착 스타일: {persona_data.get("관계성향", {}).get("애착스타일", "")}
- 관계 깊이 선호도: {persona_data.get("관계성향", {}).get("관계깊이선호도", "")}
- 초기 태도: {persona_data.get("관계성향", {}).get("초기태도", "")}

## 배경 이야기
{backstory}

## 주요 경험
{', '.join(experiences) if experiences else "정보 없음"}

-----

위 정보를 기반으로 다음 작업을 수행해주세요:

1. 상세한 배경 이야기 확장 (최소 2문장, 최대 4문장)
2. 최소 3개 이상의 구체적인 관심사 추가
3. 말투와 표현 패턴 구체화 (실제 대화에서 쓸만한 특징적 표현 3개 이상)
4. 독특한 성격 특성 추가 (기존 특성 유지하되 개성을 살릴 수 있는 디테일 추가)

강화된 페르소나 정보를 원본 JSON 구조를 유지하면서 제공해주세요. 단, 일부 필드는 세부적으로 확장하여 더 풍부하게 만들어주세요.
JSON 형식만 제공하고, 다른 설명은 하지 마세요.
"""
    
    try:
        response = gemini_query(prompt)
        
        # JSON 형식 추출
        json_str = extract_json(response)
        
        if json_str:
            enhanced_persona = json.loads(json_str)
            
            # 기존 키를 보존하면서 새 내용 병합
            for key in persona_data:
                if key not in enhanced_persona:
                    enhanced_persona[key] = persona_data[key]
            
            return enhanced_persona
        else:
            print("유효한 JSON 응답을 추출할 수 없습니다.")
            return persona_data
        
    except Exception as e:
        print(f"페르소나 강화 중 오류 발생: {str(e)}")
        return persona_data

def generate_response(persona, conversation_history):
    """
    페르소나 특성에 맞는 대화 응답을 생성합니다.
    
    Args:
        persona: 페르소나 정보
        conversation_history: 이전 대화 내역
    
    Returns:
        생성된 응답 텍스트
    """
    # 페르소나 정보 요약
    name = persona.get("기본정보", {}).get("이름", "무명")
    object_type = persona.get("기본정보", {}).get("유형", "물건")
    description = persona.get("기본정보", {}).get("설명", "")
    
    # 성격 특성 요약
    traits = []
    for trait, value in persona.get("성격특성", {}).items():
        level = "높은" if value >= 70 else "중간" if value >= 40 else "낮은"
        traits.append(f"{trait}: {level} ({value}/100)")
    
    # 최근 대화 내역 추출 (최대 10개)
    recent_conversation = []
    for msg in conversation_history[-10:]:
        role = "User" if msg["role"] == "user" else "Assistant" if msg["role"] == "assistant" else "System"
        recent_conversation.append(f"{role}: {msg['content']}")
    
    # 프롬프트 구성
    prompt = f"""
당신은 이제 다음 페르소나를 구현해야 합니다:

## 페르소나 정보
- 이름: {name}
- 유형: {object_type}
- 설명: {description}

## 성격 특성
{', '.join(traits)}

## 배경
{persona.get("배경이야기", "")}

## 성격 요약
{persona.get("성격요약", "")}

## 소통 방식
- 대화 스타일: {persona.get("소통방식", "")}
- 유머 스타일: {persona.get("유머스타일", "")}
- 매력적 결함: {', '.join(persona.get("매력적결함", []))}

## 말투 패턴 예시
{persona.get("말투패턴", "")}

## 관심사
{', '.join(persona.get("관심사", []))}

당신은 위 페르소나의 역할을 완벽하게 구현하여 사용자와 대화해야 합니다. 
온기, 능력, 신뢰성 등의 점수에 따라 성격 특성을 정확히 반영하세요.
관심사와 배경을 자연스럽게 대화에 활용하세요.
말투 패턴과 매력적 결함을 일관되게 표현하세요.

## 최근 대화 내역
{' '.join(recent_conversation)}

위 대화를 이어서, {name}으로서 답변하세요. 페르소나에 충실하되 사용자의 질문에 직접적으로 답변하세요.
답변은 한국어로만 작성하고, 절대 다른 언어를 사용하지 마세요.
"""
    
    try:
        response = gemini_query(prompt)
        
        # 응답의 첫 줄이 "Assistant:" 또는 유사한 형태로 시작하면 제거
        if response.startswith("Assistant:") or response.startswith(f"{name}:"):
            response = response.split(":", 1)[1].strip()
        
        return response
        
    except Exception as e:
        print(f"응답 생성 중 오류 발생: {str(e)}")
        return f"죄송합니다, 응답을 생성하는 중에 문제가 발생했습니다. 잠시 후 다시 시도해주세요."

def extract_json(text):
    """
    텍스트에서 JSON 형식의 데이터를 추출합니다.
    
    Args:
        text: 텍스트 데이터
    
    Returns:
        추출된 JSON 문자열 또는 None
    """
    # JSON 블록 추출 시도
    if "```json" in text:
        # 마크다운 코드 블록에서 JSON 추출
        start = text.find("```json") + 7
        end = text.find("```", start)
        if end != -1:
            return text[start:end].strip()
    elif "```" in text:
        # 일반 코드 블록에서 JSON 추출
        start = text.find("```") + 3
        end = text.find("```", start)
        if end != -1:
            return text[start:end].strip()
    
    # 중괄호를 기준으로 추출 시도
    if "{" in text and "}" in text:
        start = text.find("{")
        # 중첩된 중괄호 처리를 위한 간단한 로직
        nested = 0
        for i in range(start, len(text)):
            if text[i] == "{":
                nested += 1
            elif text[i] == "}":
                nested -= 1
                if nested == 0:
                    return text[start:i+1]
    
    # JSON 형식이 감지되지 않음
    return None