app.py

# -*- coding: utf-8 -*-

import gradio as gr
import os
import requests
import json
import logging

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

# API configuration
FRIENDLI_TOKEN = os.environ.get("FRIENDLI_TOKEN") or "YOUR_FRIENDLI_TOKEN"
FRIENDLI_API_URL = "https://api.friendli.ai/dedicated/v1/chat/completions"

# 장르별 시스템 프롬프트
GENRE_PROMPTS = {
    "판타지": """당신은 판타지 소설 작가입니다. 마법, 드래곤, 엘프, 던전, 마검사 등 판타지 세계관의 요소들을 활용하여 몰입감 있는 이야기를 작성하세요. 
주요 요소: 마법 시스템, 신화적 존재, 판타지 종족, 마법 아이템, 영웅의 여정, 예언과 운명""",
    
    "로맨스": """당신은 로맨스 소설 작가입니다. 감정선, 관계의 발전, 갈등과 화해를 중심으로 독자의 감성을 자극하는 이야기를 작성하세요.
주요 요소: 첫 만남, 감정의 변화, 오해와 화해, 운명적 사랑, 감동적인 고백, 해피엔딩 또는 애틋한 이별""",
    
    "무협": """당신은 무협 소설 작가입니다. 무공, 강호, 문파, 비급 등 동양 무협의 요소들을 활용하여 박진감 넘치는 이야기를 작성하세요.
주요 요소: 무공 수련, 문파 간 대결, 강호의 은원, 무림맹주, 마교와 정파, 절세 고수, 비급 쟁탈""",
    
    "미스터리": """당신은 미스터리 소설 작가입니다. 수수께끼, 단서, 추리 과정을 통해 독자가 함께 사건을 해결해나가는 긴장감 있는 이야기를 작성하세요.
주요 요소: 의문의 사건, 단서 발견, 용의자들, 거짓 단서, 반전, 진실의 폭로, 탐정의 추리""",
    
    "SF": """당신은 SF 소설 작가입니다. 미래 기술, 우주, 외계 생명체, 시간여행 등 과학적 상상력을 바탕으로 한 이야기를 작성하세요.
주요 요소: 미래 기술, 우주 탐험, 외계 문명, 시간 여행, AI와 로봇, 디스토피아/유토피아, 과학적 발견""",
    
    "스릴러": """당신은 스릴러 소설 작가입니다. 긴장감, 서스펜스, 예측불가능한 전개로 독자를 압도하는 이야기를 작성하세요.
주요 요소: 생존의 위기, 추격전, 심리전, 예상치 못한 반전, 시간제한, 극적인 대립, 긴박한 상황""",
    
    "역사": """당신은 역사 소설 작가입니다. 실제 역사적 배경을 바탕으로 당시의 문화와 인물들을 생생하게 그려내는 이야기를 작성하세요.
주요 요소: 시대적 배경, 역사적 사건, 실존 인물, 당시의 문화, 역사적 갈등, 시대정신, 고증""",
    
    "호러": """당신은 호러 소설 작가입니다. 공포, 초자연적 현상, 미지의 존재를 통해 독자에게 오싹한 전율을 선사하는 이야기를 작성하세요.
주요 요소: 초자연적 현상, 저주, 귀신, 미지의 존재, 공포스러운 분위기, 심리적 공포, 생존의 위협"""
}

def format_text(text: str, max_line_length: int = 80) -> str:
    """텍스트 포맷팅 함수"""
    lines = []
    current_line = ""
    
    for paragraph in text.split('\n'):
        words = paragraph.split()
        
        for word in words:
            if len(current_line) + len(word) + 1 <= max_line_length:
                current_line += word + " "
            else:
                lines.append(current_line.strip())
                current_line = word + " "
        
        if current_line:
            lines.append(current_line.strip())
            current_line = ""
        
        lines.append("")  # 문단 구분을 위한 빈 줄
    
    return "\n".join(lines)

def respond(
    message,
    history: list[tuple[str, str]],
    genre="판타지",
    system_message="",  
    max_tokens=8000,  
    temperature=0.8, 
    top_p=0.9,  
):
    # 선택된 장르에 따른 시스템 프롬프트 설정
    genre_prompt = GENRE_PROMPTS.get(genre, GENRE_PROMPTS["판타지"])
    
    system_prefix = f"""당신은 한국어 {genre} 소설 전문 작가입니다. 
{genre_prompt}

작성 지침:
1. 생생한 묘사와 몰입감 있는 서술
2. 입체적인 캐릭터 설정과 대화
3. 긴장감 있는 플롯 전개
4. 장르 특성에 맞는 문체와 분위기
5. 독자를 사로잡는 흥미진진한 전개

각 챕터는 완결성을 가지면서도 전체 이야기의 연속성을 유지해야 합니다.
최대 8000 토큰까지 사용하여 풍부하고 상세한 이야기를 작성하세요."""

    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})

    current_response = ""
    new_history = history.copy()

    try:
        headers = {
            "Authorization": f"Bearer {FRIENDLI_TOKEN}",
            "Content-Type": "application/json"
        }
        
        payload = {
            "model": "dep86pjolcjjnv8",
            "messages": messages,
            "max_tokens": max_tokens,
            "temperature": temperature,
            "top_p": top_p,
            "stream": True,
            "stream_options": {
                "include_usage": True
            }
        }
        
        response = requests.post(FRIENDLI_API_URL, json=payload, headers=headers, stream=True)
        response.raise_for_status()
        
        for line in response.iter_lines():
            if line:
                line = line.decode('utf-8')
                if line.startswith('data: '):
                    line = line[6:]
                    if line == '[DONE]':
                        break
                    
                    try:
                        chunk = json.loads(line)
                        if 'choices' in chunk and len(chunk['choices']) > 0:
                            if 'delta' in chunk['choices'][0] and 'content' in chunk['choices'][0]['delta']:
                                token = chunk['choices'][0]['delta']['content']
                                if token:
                                    current_response += token
                                    formatted_response = format_text(current_response)
                                    new_history = history + [(message, formatted_response)]
                                    yield new_history
                    except json.JSONDecodeError:
                        continue

        final_response = format_text(current_response)
        new_history = history + [(message, final_response)]
        yield new_history

    except Exception as e:
        error_message = f"오류 발생: {str(e)}"
        logging.error(f"응답 생성 실패: {error_message}")
        yield history + [(message, error_message)]

with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css="""
    .message-wrap {
        font-size: 16px !important;
        line-height: 1.6em !important;
        max-width: 90% !important;
        margin: 0 auto !important;
    }
    .message {
        padding: 1em !important;
        margin-bottom: 0.5em !important;
        white-space: pre-wrap !important;
        word-wrap: break-word !important;
        max-width: 100% !important;
    }
    .message p {
        margin: 0 !important;
        padding: 0 !important;
        width: 100% !important;
    }
    .chatbot {
        font-family: 'Noto Sans KR', sans-serif !important;
    }
""") as interface:
    gr.Markdown("# 한국어 소설 생성기")
    gr.Markdown("### 원하는 장르를 선택하고 프롬프트를 입력하여 소설을 생성하세요.")    

    with gr.Row():
        with gr.Column():
            genre_dropdown = gr.Dropdown(
                choices=list(GENRE_PROMPTS.keys()),
                value="판타지",
                label="장르 선택",
                info="원하는 소설 장르를 선택하세요"
            )
            
            chatbot = gr.Chatbot(
                value=[],
                show_label=True,
                label="소설 진행",
                height=600,
                elem_classes="chatbot"
            )
            
            with gr.Row():
                msg = gr.Textbox(
                    label="프롬프트 입력",
                    placeholder="이야기의 시작이나 진행 방향을 입력하세요...",
                    lines=3
                )
                submit_btn = gr.Button("생성", variant="primary")
            
            with gr.Accordion("고급 설정", open=False):
                system_msg = gr.Textbox(
                    label="추가 시스템 메시지",
                    value="독자를 사로잡는 흥미진진한 이야기를 작성하세요.",
                    lines=2
                )
                
                with gr.Row():
                    max_tokens = gr.Slider(
                        minimum=100,
                        maximum=8000,
                        value=8000,
                        step=100,
                        label="최대 토큰 수"
                    )
                    temperature = gr.Slider(
                        minimum=0,
                        maximum=1,
                        value=0.8,
                        step=0.1,
                        label="창의성 수준"
                    )
                    top_p = gr.Slider(
                        minimum=0,
                        maximum=1,
                        value=0.9,
                        step=0.1,
                        label="응답 다양성"
                    )

    examples = gr.Examples(
        examples=[
            ["이야기를 계속 진행해주세요"],
            ["새로운 캐릭터를 등장시켜주세요"],
            ["긴장감 있는 전투 장면을 만들어주세요"],
            ["주인공의 과거를 밝혀주세요"],
            ["예상치 못한 반전을 만들어주세요"],
            ["감동적인 장면을 연출해주세요"],
            ["이야기의 클라이맥스를 만들어주세요"],
            ["새로운 장소로 이동하는 장면을 만들어주세요"],
        ],
        inputs=msg,
        label="예시 프롬프트"
    )

    submit_btn.click(
        fn=respond,
        inputs=[msg, chatbot, genre_dropdown, system_msg, max_tokens, temperature, top_p],
        outputs=[chatbot]
    )
    
    msg.submit(
        fn=respond,
        inputs=[msg, chatbot, genre_dropdown, system_msg, max_tokens, temperature, top_p],
        outputs=[chatbot]
    )

if __name__ == "__main__":
    interface.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True
    )