app.py

import os
import requests
import json
import gradio as gr
import re

# Hugging Face 환경 변수로부터 RapidAPI 키와 호스트 가져오기
AA_KEY = os.getenv("AA_KEY")
AA_HOST = "youtube-transcriptor.p.rapidapi.com"

# 자막 언어 우선순위 리스트 
LANGUAGE_PRIORITY = ['ko', 'en', 'ja', 'zh']

# 유튜브 URL에서 비디오 ID를 추출하는 함수
def get_video_id(youtube_url):
    # 유튜브 URL 또는 youtu.be 단축 URL에서 video_id 추출
    video_id_match = re.search(r"(?<=v=)[^#&?]*", youtube_url) or re.search(r"(?<=youtu.be/)[^#&?]*", youtube_url)
    
    # YouTube Shorts URL 처리
    if not video_id_match:
        video_id_match = re.search(r"(?<=shorts/)[^#&?]*", youtube_url)
    
    return video_id_match.group(0) if video_id_match else None

# 유튜브 자막을 요청하는 함수 (언어 우선순위를 적용하여 시도)
def get_youtube_transcript(youtube_url):
    try:
        # 비디오 ID 추출
        video_id = get_video_id(youtube_url)
        if video_id is None:
            return {"error": "잘못된 유튜브 URL입니다."}
        
        url = "https://youtube-transcriptor.p.rapidapi.com/transcript"
        headers = {
            "x-rapidapi-key": AA_KEY,
            "x-rapidapi-host": AA_HOST
        }
        
        # 1. 우선순위 언어로 시도
        for lang in LANGUAGE_PRIORITY:
            querystring = {"video_id": video_id, "lang": lang}
            response = requests.get(url, headers=headers, params=querystring)
            
            if response.status_code == 200:
                data = response.json()
                if data and not isinstance(data, str) and "error" not in data:
                    return {"language": lang, "data": data}
                if isinstance(data, dict) and "availableLangs" in data:
                    available_langs = data["availableLangs"]
                    # 2. 가능한 언어가 있다면 해당 언어로 시도
                    for available_lang in available_langs:
                        querystring = {"video_id": video_id, "lang": available_lang}
                        response = requests.get(url, headers=headers, params=querystring)
                        if response.status_code == 200:
                            data = response.json()
                            if data and not isinstance(data, str) and "error" not in data:
                                return {"language": available_lang, "data": data}
        
        # 3. 모든 시도 실패시
        return {"error": "자막을 찾을 수 없습니다."}
        
    except Exception as e:
        return {"error": "자막을 불러오는데 실패했습니다."}

# Gradio 인터페이스 함수
def youtube_transcript_interface(youtube_url):
    transcript_data = get_youtube_transcript(youtube_url)
    
    if 'error' in transcript_data:
        return transcript_data['error']
    
    try:
        # 데이터가 리스트인 경우 첫 번째 항목을 사용
        data = transcript_data['data'][0]
        title = data.get('title', '제목 없음')
        transcription = data.get('transcriptionAsText', '자막이 없습니다.')
        
        # 원하는 형식으로 반환
        result = f"제목: {title}\n\n자막:\n{transcription}"
        return result
    except (IndexError, KeyError, TypeError):
        return "자막 데이터를 처리하는 중 오류가 발생했습니다."

# Gradio 인터페이스 생성
interface = gr.Interface(
    fn=youtube_transcript_interface,
    inputs="text",
    outputs="text",
    title="YouTube 자막 추출기",
    description="유튜브 URL을 입력하세요."
)

# Gradio 인터페이스 실행
interface.launch()