app.py

import gradio as gr
import websocket
import json
import base64
import numpy as np
import threading
import queue
import os
from datetime import datetime
import pyaudio
import wave
import io

class RealtimeTranslator:
    def __init__(self):
        self.ws = None
        self.api_key = os.getenv("OPENAI_API_KEY")
        self.audio_queue = queue.Queue()
        self.transcript_queue = queue.Queue()
        self.translation_queue = queue.Queue()
        self.is_connected = False
        self.is_recording = False
        self.source_lang = "ko"
        self.target_lang = "en"
        
        # PyAudio 설정
        self.p = pyaudio.PyAudio()
        self.sample_rate = 24000
        self.chunk_size = 1024
        self.audio_format = pyaudio.paInt16
        
    def connect_websocket(self):
        """WebSocket 연결 설정"""
        try:
            url = "wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview-2024-12-17"
            headers = {
                "Authorization": f"Bearer {self.api_key}",
                "OpenAI-Beta": "realtime=v1"
            }
            
            self.ws = websocket.WebSocketApp(
                url,
                header=headers,
                on_open=self.on_open,
                on_message=self.on_message,
                on_error=self.on_error,
                on_close=self.on_close
            )
            
            # WebSocket을 별도 스레드에서 실행
            wst = threading.Thread(target=self.ws.run_forever)
            wst.daemon = True
            wst.start()
            
            return "연결 성공"
        except Exception as e:
            return f"연결 실패: {str(e)}"
    
    def on_open(self, ws):
        """WebSocket 연결 시 호출"""
        self.is_connected = True
        print("WebSocket 연결됨")
        
        # 세션 설정
        session_update = {
            "type": "session.update",
            "session": {
                "modalities": ["text", "audio"],
                "instructions": f"You are a helpful translator. Translate between {self.get_language_name(self.source_lang)} and {self.get_language_name(self.target_lang)}. Respond with both the transcription and translation.",
                "voice": "alloy",
                "input_audio_format": "pcm16",
                "output_audio_format": "pcm16",
                "input_audio_transcription": {
                    "model": "whisper-1"
                },
                "turn_detection": {
                    "type": "server_vad",
                    "threshold": 0.5,
                    "prefix_padding_ms": 300,
                    "silence_duration_ms": 500
                }
            }
        }
        ws.send(json.dumps(session_update))
    
    def on_message(self, ws, message):
        """메시지 수신 시 호출"""
        try:
            event = json.loads(message)
            event_type = event.get("type")
            
            if event_type == "conversation.item.input_audio_transcription.completed":
                # 음성 전사 완료
                transcript = event.get("transcript", "")
                self.transcript_queue.put(transcript)
                
                # 번역 요청
                self.request_translation(transcript)
                
            elif event_type == "response.text.delta":
                # 번역 결과 수신
                delta = event.get("delta", "")
                self.translation_queue.put(delta)
                
            elif event_type == "response.audio.delta":
                # 오디오 데이터 수신
                audio_data = base64.b64decode(event.get("delta", ""))
                self.audio_queue.put(audio_data)
                
            elif event_type == "error":
                error_msg = event.get("error", {}).get("message", "Unknown error")
                print(f"Error: {error_msg}")
                
        except Exception as e:
            print(f"메시지 처리 오류: {str(e)}")
    
    def on_error(self, ws, error):
        """오류 발생 시 호출"""
        print(f"WebSocket 오류: {error}")
        self.is_connected = False
    
    def on_close(self, ws, close_status_code, close_msg):
        """연결 종료 시 호출"""
        print("WebSocket 연결 종료")
        self.is_connected = False
    
    def get_language_name(self, lang_code):
        """언어 코드를 언어 이름으로 변환"""
        languages = {
            "ko": "Korean",
            "en": "English",
            "ja": "Japanese",
            "zh": "Chinese",
            "es": "Spanish",
            "fr": "French"
        }
        return languages.get(lang_code, lang_code)
    
    def request_translation(self, text):
        """번역 요청"""
        if not self.ws or not self.is_connected:
            return
        
        message = {
            "type": "conversation.item.create",
            "item": {
                "type": "message",
                "role": "user",
                "content": [{
                    "type": "input_text",
                    "text": f"Translate this {self.get_language_name(self.source_lang)} text to {self.get_language_name(self.target_lang)}: '{text}'"
                }]
            }
        }
        
        self.ws.send(json.dumps(message))
        
        # 응답 생성 요청
        response_create = {"type": "response.create"}
        self.ws.send(json.dumps(response_create))
    
    def send_audio_chunk(self, audio_data):
        """오디오 청크 전송"""
        if not self.ws or not self.is_connected:
            return
        
        # PCM16 형식으로 인코딩
        audio_base64 = base64.b64encode(audio_data).decode('utf-8')
        
        message = {
            "type": "input_audio_buffer.append",
            "audio": audio_base64
        }
        
        self.ws.send(json.dumps(message))
    
    def process_audio(self, audio_file):
        """오디오 파일 처리 및 전송"""
        if not self.is_connected:
            return "WebSocket이 연결되지 않았습니다.", ""
        
        try:
            # 오디오 파일 읽기
            with wave.open(audio_file, 'rb') as wf:
                # 오디오를 24kHz PCM16으로 변환 필요
                audio_data = wf.readframes(wf.getnframes())
                
            # 오디오 데이터를 청크로 나누어 전송
            chunk_size = 4096
            for i in range(0, len(audio_data), chunk_size):
                chunk = audio_data[i:i+chunk_size]
                self.send_audio_chunk(chunk)
            
            # 오디오 버퍼 커밋
            commit_message = {"type": "input_audio_buffer.commit"}
            self.ws.send(json.dumps(commit_message))
            
            # 전사 및 번역 결과 대기
            transcript = ""
            translation = ""
            
            # 타임아웃 설정 (10초)
            import time
            timeout = 10
            start_time = time.time()
            
            while time.time() - start_time < timeout:
                # 전사 결과 확인
                try:
                    transcript = self.transcript_queue.get(timeout=0.1)
                except queue.Empty:
                    pass
                
                # 번역 결과 확인
                try:
                    while not self.translation_queue.empty():
                        translation += self.translation_queue.get()
                except queue.Empty:
                    pass
                
                if transcript and translation:
                    break
            
            return transcript, translation
            
        except Exception as e:
            return f"오류: {str(e)}", ""
    
    def disconnect(self):
        """WebSocket 연결 종료"""
        if self.ws:
            self.ws.close()
        self.is_connected = False
        return "연결 종료됨"

# Gradio 인터페이스 생성
def create_interface():
    translator = RealtimeTranslator()
    
    def connect():
        if not translator.api_key:
            return "API 키가 설정되지 않았습니다. 환경 변수 OPENAI_API_KEY를 설정하세요.", gr.update(value=False)
        result = translator.connect_websocket()
        return result, gr.update(value=translator.is_connected)
    
    def disconnect():
        result = translator.disconnect()
        return result, gr.update(value=False)
    
    def translate_audio(audio_file, source_lang, target_lang):
        if not audio_file:
            return "오디오 파일을 선택하세요.", "", None
        
        translator.source_lang = source_lang
        translator.target_lang = target_lang
        
        transcript, translation = translator.process_audio(audio_file)
        
        # 오디오 응답 처리 (현재는 텍스트만 반환)
        return transcript, translation, None
    
    def swap_languages(source, target):
        return target, source
    
    with gr.Blocks(title="실시간 음성 번역기") as demo:
        gr.Markdown("# 🎙️ OpenAI Realtime API 음성 번역기")
        gr.Markdown("실시간으로 음성을 전사하고 번역합니다.")
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 연결 상태")
                connection_status = gr.Checkbox(label="연결됨", value=False, interactive=False)
                connect_btn = gr.Button("연결", variant="primary")
                disconnect_btn = gr.Button("연결 종료", variant="secondary")
                status_text = gr.Textbox(label="상태 메시지", value="연결되지 않음")
        
        with gr.Row():
            with gr.Column(scale=2):
                gr.Markdown("### 언어 설정")
                with gr.Row():
                    source_lang = gr.Dropdown(
                        choices=[("한국어", "ko"), ("영어", "en"), ("일본어", "ja"), 
                                ("중국어", "zh"), ("스페인어", "es"), ("프랑스어", "fr")],
                        value="ko",
                        label="입력 언어"
                    )
                    swap_btn = gr.Button("↔️", scale=0)
                    target_lang = gr.Dropdown(
                        choices=[("한국어", "ko"), ("영어", "en"), ("일본어", "ja"), 
                                ("중국어", "zh"), ("스페인어", "es"), ("프랑스어", "fr")],
                        value="en",
                        label="출력 언어"
                    )
        
        with gr.Row():
            with gr.Column():
                gr.Markdown("### 음성 입력")
                audio_input = gr.Audio(
                    source="microphone",
                    type="filepath",
                    label="녹음하기"
                )
                translate_btn = gr.Button("번역하기", variant="primary")
        
        with gr.Row():
            with gr.Column():
                gr.Markdown("### 결과")
                transcript_output = gr.Textbox(
                    label="전사된 텍스트",
                    placeholder="음성 전사 결과가 여기에 표시됩니다...",
                    lines=3
                )
                translation_output = gr.Textbox(
                    label="번역된 텍스트",
                    placeholder="번역 결과가 여기에 표시됩니다...",
                    lines=3
                )
                audio_output = gr.Audio(
                    label="번역된 음성",
                    type="filepath"
                )
        
        # 이벤트 핸들러
        connect_btn.click(
            fn=connect,
            outputs=[status_text, connection_status]
        )
        
        disconnect_btn.click(
            fn=disconnect,
            outputs=[status_text, connection_status]
        )
        
        swap_btn.click(
            fn=swap_languages,
            inputs=[source_lang, target_lang],
            outputs=[source_lang, target_lang]
        )
        
        translate_btn.click(
            fn=translate_audio,
            inputs=[audio_input, source_lang, target_lang],
            outputs=[transcript_output, translation_output, audio_output]
        )
        
        gr.Markdown("""
        ### 📝 사용 방법
        1. **연결** 버튼을 클릭하여 OpenAI Realtime API에 연결합니다.
        2. 입력 언어와 출력 언어를 선택합니다.
        3. 마이크 버튼을 클릭하여 음성을 녹음합니다.
        4. **번역하기** 버튼을 클릭하면 전사 및 번역이 진행됩니다.
        
        ### ⚠️ 주의사항
        - 환경 변수 `OPENAI_API_KEY`가 설정되어 있어야 합니다.
        - 긴 오디오의 경우 처리 시간이 오래 걸릴 수 있습니다.
        """)
    
    return demo

# 실행
if __name__ == "__main__":
    # 필요한 패키지 설치 안내
    print("""
    필요한 패키지:
    pip install gradio websocket-client pyaudio wave numpy
    
    환경 변수 설정:
    export OPENAI_API_KEY="your-api-key-here"
    """)
    
    demo = create_interface()
    demo.launch(share=True)