Update app.py

app.py

@@ -3,242 +3,294 @@ import sys
 import subprocess
 import threading
 import time
+import requests
+import socket
 import io
 import wave
-import audioop
-import uuid
-from flask import Flask, request
+import json
+import yt_dlp
+from flask import Flask, render_template, request
 from flask_socketio import SocketIO, emit
 from dotenv import load_dotenv
 import google.generativeai as genai
-from groq import Groq
+
+# DNS Patch
+_original_getaddrinfo = socket.getaddrinfo
+def patched_getaddrinfo(host, port, family=0, type=0, proto=0, flags=0):
+    if host and ("youtube.com" in host or "youtu.be" in host or "googlevideo.com" in host):
+        try:
+            url = f"https://8.8.8.8/resolve?name={host}&type=A"
+            r = requests.get(url, timeout=2)
+            data = r.json()
+            if "Answer" in data:
+                ip = data["Answer"][0]["data"]
+                return _original_getaddrinfo(ip, port, family, type, proto, flags)
+        except Exception:
+            pass 
+    return _original_getaddrinfo(host, port, family, type, proto, flags)
+socket.getaddrinfo = patched_getaddrinfo
 
 load_dotenv()
 
-# === API設定 ===
-GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
-GROQ_API_KEY = os.getenv("GROQ_API_KEY")
+try:
+    cookie_content = os.getenv("YOUTUBE_COOKIES")
+    if cookie_content:
+        cookie_content = cookie_content.replace("\\n", "\n")
+        with open("cookies.txt", "w") as f:
+            f.write(cookie_content)
+except Exception:
+    pass
 
-if not GEMINI_API_KEY: print("⚠️ GEMINI_API_KEYが設定されていません")
-if not GROQ_API_KEY: print("⚠️ GROQ_API_KEYが設定されていません")
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
+model_name = "gemini-2.0-flash" 
 
-genai.configure(api_key=GEMINI_API_KEY)
-groq_client = Groq(api_key=GROQ_API_KEY) if GROQ_API_KEY else None
+if not GEMINI_API_KEY:
+    print("⚠️ GEMINI_API_KEYが設定されていません")
+else:
+    genai.configure(api_key=GEMINI_API_KEY)
 
 sys.stdout.reconfigure(line_buffering=True)
 
 app = Flask(__name__)
 app.config['SECRET_KEY'] = os.getenv("SECRET_KEY", "secret!")
-socketio = SocketIO(app, cors_allowed_origins="*", max_http_buffer_size=10 * 1024 * 1024)
+socketio = SocketIO(app, cors_allowed_origins="*")
 
-# Geminiモデル (確定用)
-model = genai.GenerativeModel("models/gemini-2.0-flash")
+MAX_CONCURRENT_STREAMS = 5
+SPEAKER_COLORS = ["#FFFFFF", "#03DAC6", "#FFB400", "#CF6679", "#BB86FC"]
 
-active_processors = {}
-SILENCE_THRESHOLD = 500 
+session_speakers = {}
+active_sessions = {}
 
-class AudioProcessor:
-    def __init__(self, sid):
+class StreamWorker:
+    def __init__(self, sid, url, settings, socket_instance):
         self.sid = sid
-        self.running = True
-        self.settings = {'target_lang': 'ja', 'max_speakers': 2}
-        self.reset_flag = False # バッファリセット用フラグ
-        
-        # FFmpegプロセス
-        self.ffmpeg_proc = subprocess.Popen(
-            [
-                'ffmpeg', '-i', 'pipe:0', '-f', 's16le', '-ac', '1',
-                '-ar', '16000', '-acodec', 'pcm_s16le', 'pipe:1', '-v', 'error'
-            ],
-            stdin=subprocess.PIPE,
-            stdout=subprocess.PIPE,
-            stderr=sys.stderr
-        )
-
-        self.thread = threading.Thread(target=self._process_loop, daemon=True)
-        self.thread.start()
-
-    def write_data(self, chunk):
-        if self.running and self.ffmpeg_proc.stdin:
-            try:
-                self.ffmpeg_proc.stdin.write(chunk)
-                self.ffmpeg_proc.stdin.flush()
-            except Exception as e:
-                print(f"❌ FFmpeg書き込みエラー: {e}")
-
-    def update_settings(self, data):
-        self.settings.update(data)
-
-    def trigger_reset(self):
-        """シーク時などにバッファを強制クリアする"""
-        self.reset_flag = True
-        print(f"🔄 バッファリセット予約: {self.sid}")
+        self.url = url
+        self.settings = settings
+        self.socketio = socket_instance
+        self.is_running = False
+        self.ffmpeg_proc = None
+        self.dl_proc = None
+        self.fifo_path = f"/tmp/stream_{sid}.fifo"
+        self.model = genai.GenerativeModel(model_name)
+
+    def start(self):
+        if self.is_running: return
+        self.is_running = True
+        self.socketio.start_background_task(target=self._main_loop)
 
     def stop(self):
-        self.running = False
-        if self.ffmpeg_proc:
-            try:
-                self.ffmpeg_proc.stdin.close()
-                self.ffmpeg_proc.kill()
+        self.is_running = False
+        if self.ffmpeg_proc and self.ffmpeg_proc.poll() is None:
+            self.ffmpeg_proc.kill()
+        if os.path.exists(self.fifo_path):
+            try: os.remove(self.fifo_path)
             except: pass
-        print(f"🛑 プロセッサ停止: {self.sid}")
-
-    # --- 速報: Groq (Whisper) ---
-    def _call_groq_fast(self, wav_bytes, chunk_id, tgt_name):
-        if not groq_client: return
+        print(f"🛑 Worker停止: {self.sid}")
+
+    def _run_ytdlp(self):
+        if os.path.exists(self.fifo_path): os.remove(self.fifo_path)
+        os.mkfifo(self.fifo_path)
+        ydl_opts = {
+            'format': 'bestaudio/best',
+            'outtmpl': self.fifo_path,
+            'quiet': True, 'no_warnings': True, 'nocheckcertificate': True,
+            'cookiefile': 'cookies.txt' if os.path.exists("cookies.txt") else None,
+            'socket_timeout': 10, 'retries': 10,
+        }
         try:
-            # Whisperへの指示
-            prompt_text = f"Translate everything into {tgt_name}."
-            
-            completion = groq_client.audio.transcriptions.create(
-                file=("audio.wav", wav_bytes),
-                model="whisper-large-v3",
-                prompt=prompt_text,
-                response_format="json",
-                temperature=0.0
-            )
-            
-            text = completion.text.strip()
-            if text:
-                # 速報を送信
-                socketio.emit('subtitle_preview', {
-                    'chunk_id': chunk_id,
-                    'text': text
-                }, room=self.sid)
-                print(f"🚀 Groq速報: {text[:20]}...")
-
+            with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+                ydl.download([self.url])
         except Exception as e:
-            print(f"⚠️ Groqエラー: {e}")
+            print(f"❌ yt-dlpエラー ({self.sid}): {e}")
+            self.is_running = False
 
-    # --- 確定: Gemini ---
-    def _call_gemini_accurate(self, wav_bytes, chunk_id, tgt_name, max_spk):
+    def _process_audio_with_gemini(self, audio_data):
         try:
+            tgt_lang = self.settings.get('target_lang', 'ja')
+            tgt_name = "Japanese" if tgt_lang == 'ja' else "English"
+            max_spk = self.settings.get('max_speakers', 2)
+
             prompt = f"""
-            Task: Listen to the audio.
-            1. Identify distinct speakers (Max {max_spk}).
+            Task: Listen to this audio.
+            1. Identify distinct speakers. (Context implies MAX {max_spk} speakers).
             2. Translate speech to {tgt_name}.
             
-            Output format:
-            Speaker 1: [Translation]
+            Output format (Strictly plain text):
+            Speaker 1: [Translated Text]
+            Speaker 2: [Translated Text]
             
             Rules:
-            - Output ONLY translation.
-            - If one person, use "Speaker 1".
-            - Adapt tone/context.
+            - If it sounds like one person but tone changes, keep it as Speaker 1 if limit is 1.
+            - Do NOT output original text.
+            - If no speech, output nothing.
             """
 
-            response = model.generate_content([
+            response = self.model.generate_content([
                 prompt,
-                {"mime_type": "audio/wav", "data": wav_bytes}
+                {"mime_type": "audio/wav", "data": audio_data}
             ])
-            text = response.text.strip()
-            
-            if text:
-                parsed_lines = []
-                for line in text.split('\n'):
+            return response.text.strip()
+        except Exception:
+            return ""
+
+    def _main_loop(self):
+        print(f"🔁 処理開始 [User: {self.sid}] URL: {self.url}")
+        self.download_thread = threading.Thread(target=self._run_ytdlp)
+        self.download_thread.daemon = True
+        self.download_thread.start()
+
+        time.sleep(1)
+        if not os.path.exists(self.fifo_path):
+            self.socketio.emit("error_msg", {"msg": "開始失敗"}, room=self.sid)
+            self.stop()
+            return
+
+        ffmpeg_cmd = ["ffmpeg", "-y", "-i", self.fifo_path, "-f", "s16le", "-acodec", "pcm_s16le", "-ac", "1", "-ar", "16000", "pipe:1"]
+        try:
+            self.ffmpeg_proc = subprocess.Popen(ffmpeg_cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
+        except Exception:
+            self.stop(); return
+
+        chunk_seconds = 8 
+        chunk_size = 16000 * 2 * chunk_seconds
+
+        while self.is_running:
+            try:
+                audio_chunk_raw = self.ffmpeg_proc.stdout.read(chunk_size)
+                if not audio_chunk_raw:
+                    if self.is_running and self.ffmpeg_proc.poll() is None:
+                        time.sleep(0.5); continue
+                    else: break
+                
+                wav_io = io.BytesIO()
+                with wave.open(wav_io, 'wb') as wf:
+                    wf.setnchannels(1); wf.setsampwidth(2); wf.setframerate(16000)
+                    wf.writeframes(audio_chunk_raw)
+                audio_bytes = wav_io.getvalue()
+
+                result_text = self._process_audio_with_gemini(audio_bytes)
+                if not result_text: continue
+
+                for line in result_text.split('\n'):
                     if ':' in line:
                         parts = line.split(':', 1)
-                        parsed_lines.append({
-                            'speaker': parts[0].strip(),
-                            'text': parts[1].strip()
-                        })
+                        if len(parts) < 2: continue
                         
-                if parsed_lines:
-                    # 確定を送信
-                    socketio.emit('subtitle_final', {
-                        'chunk_id': chunk_id,
-                        'results': parsed_lines
-                    }, room=self.sid)
-                    print(f"✅ Gemini確定: {parsed_lines[0]['text'][:20]}...")
+                        original_speaker_id = parts[0].strip()
+                        text_content = parts[1].strip()
+                        if not text_content: continue
 
-        except Exception as e:
-            print(f"⚠️ Geminiエラー: {e}")
+                        # 【強制補正ロジック】
+                        max_spk = self.settings.get('max_speakers', 2)
+                        
+                        # 1人の設定なら、何を言われても "Speaker 1" にする
+                        if max_spk == 1:
+                            original_speaker_id = "Speaker 1"
+                        
+                        # 数字部分を抽出してチェック (例: "Speaker 5" -> 5)
+                        spk_num_match = re.search(r'\d+', original_speaker_id)
+                        if spk_num_match:
+                            spk_num = int(spk_num_match.group())
+                            if spk_num > max_spk:
+                                # 制限を超えた話者は Speaker 1 にまとめる（または無視）
+                                original_speaker_id = "Speaker 1"
+
+                        if self.sid not in session_speakers: session_speakers[self.sid] = {}
+                        
+                        # 新規登録
+                        if original_speaker_id not in session_speakers[self.sid]:
+                            # もし登録済み人数が制限を超えていたら、既存の最後に割り当てる等の処理も可能だが
+                            # 上記のID書き換えで対応済み
+                            new_idx = len(session_speakers[self.sid])
+                            color = SPEAKER_COLORS[new_idx % len(SPEAKER_COLORS)]
+                            session_speakers[self.sid][original_speaker_id] = {"name": original_speaker_id, "color": color}
+                            self.socketio.emit("initial_speaker_profiles", session_speakers[self.sid], room=self.sid)
+                        
+                        spk_data = session_speakers[self.sid][original_speaker_id]
+                        display_name = spk_data["name"]
+                        display_color = spk_data["color"]
+
+                        self.socketio.emit("new_subtitle", {
+                            "original": "...", 
+                            "translated": text_content,
+                            "speaker": display_name,
+                            "speaker_id": original_speaker_id,
+                            "color": display_color
+                        }, room=self.sid)
+                        
+                        print(f"✅ 送信({self.sid}): {display_name}: {text_content}")
 
-    def _process_loop(self):
-        CHUNK_DURATION = 4 # バッファサイズ (秒)
-        REQUIRED_BYTES = 16000 * 2 * CHUNK_DURATION
-        pcm_buffer = bytearray()
+            except Exception: continue
+        
+        self.stop()
+
+@app.route("/")
+def index():
+    return render_template("index.html")
 
-        while self.running:
-            try:
-                # バッファリセット判定
-                if self.reset_flag:
-                    pcm_buffer.clear()
-                    self.reset_flag = False
-                    print("✨ バッファクリア完了")
-
-                raw_chunk = self.ffmpeg_proc.stdout.read(4096)
-                if not raw_chunk:
-                    time.sleep(0.1)
-                    if self.ffmpeg_proc.poll() is not None: break
-                    continue
-
-                pcm_buffer.extend(raw_chunk)
-
-                if len(pcm_buffer) >= REQUIRED_BYTES:
-                    to_process = pcm_buffer[:REQUIRED_BYTES]
-                    pcm_buffer = pcm_buffer[REQUIRED_BYTES:] # スライドさせるか、オーバーラップさせるか
-
-                    # 無音検知 (VAD)
-                    try:
-                        rms = audioop.rms(to_process, 2)
-                        if rms < SILENCE_THRESHOLD:
-                            continue # 無音ならスキップ
-                    except: pass
-
-                    # WAV作成
-                    wav_io = io.BytesIO()
-                    with wave.open(wav_io, 'wb') as wf:
-                        wf.setnchannels(1); wf.setsampwidth(2); wf.setframerate(16000)
-                        wf.writeframes(to_process)
-                    wav_bytes = wav_io.getvalue()
-
-                    # 設定取得
-                    tgt_lang = self.settings.get('target_lang', 'ja')
-                    tgt_name = "Japanese" if tgt_lang == 'ja' else "English"
-                    max_spk = self.settings.get('max_speakers', 2)
-                    chunk_id = str(uuid.uuid4())
-
-                    # ★並列処理開始
-                    t1 = threading.Thread(target=self._call_groq_fast, args=(wav_bytes, chunk_id, tgt_name))
-                    t2 = threading.Thread(target=self._call_gemini_accurate, args=(wav_bytes, chunk_id, tgt_name, max_spk))
-                    t1.start()
-                    t2.start()
-
-            except Exception as e:
-                print(f"❌ ループエラー: {e}")
-                time.sleep(1)
-
-# === Socket.IO イベント ===
 @socketio.on('connect')
-def handle_connect():
-    active_processors[request.sid] = AudioProcessor(request.sid)
+def handle_connect(): pass
 
 @socketio.on('disconnect')
 def handle_disconnect():
-    if request.sid in active_processors:
-        active_processors[request.sid].stop()
-        del active_processors[request.sid]
+    sid = request.sid
+    if sid in active_sessions:
+        active_sessions[sid].stop()
+        del active_sessions[sid]
+    if sid in session_speakers: del session_speakers[sid]
+
+@socketio.on('start_stream')
+def handle_start_stream(data):
+    sid = request.sid
+    if sid in active_sessions: active_sessions[sid].stop()
+    if len(active_sessions) >= MAX_CONCURRENT_STREAMS:
+        emit("error_msg", {"msg": "混雑中"}, room=sid); return
+
+    settings = {
+        'target_lang': data.get('target_lang', 'ja'),
+        'max_speakers': int(data.get('max_speakers', 2))
+    }
+    worker = StreamWorker(sid, data.get('url'), settings, socketio)
+    active_sessions[sid] = worker
+    worker.start()
+    emit("status_msg", {"msg": "開始"}, room=sid)
+
+@socketio.on('stop_stream')
+def handle_stop_stream():
+    if request.sid in active_sessions: active_sessions[request.sid].stop()
+    emit("status_msg", {"msg": "停止"}, room=request.sid)
 
 @socketio.on('update_settings')
-def handle_settings(data):
-    if request.sid in active_processors:
-        active_processors[request.sid].update_settings(data)
-
-@socketio.on('audio_stream')
-def handle_audio_stream(data):
-    if request.sid in active_processors:
-        active_processors[request.sid].write_data(data)
-
-# ★追加: クライアントからのリセット要求
-@socketio.on('reset_stream')
-def handle_stream_reset():
+def handle_settings_update(data):
+    if request.sid in active_sessions:
+        active_sessions[request.sid].settings.update(data)
+
+
+# === リセットボタン ===
+@socketio.on('reset_speakers')
+def handle_reset_speakers():
     sid = request.sid
-    if sid in active_processors:
-        active_processors[sid].trigger_reset()
+    if sid in session_speakers:
+        session_speakers[sid] = {}
+        # 空の状態を通知
+        emit("initial_speaker_profiles", {}, room=sid)
+    print(f"🔄 話者リセット: {sid}")
+
+@socketio.on('update_speaker_name')
+def handle_speaker_name_update(data):
+    sid = request.sid
+    id = data.get("id"); new_name = data.get("new_name")
+    if sid in session_speakers and id in session_speakers[sid]:
+        session_speakers[sid][id]["name"] = new_name
+        emit("speaker_name_changed", {"id": id, "new_name": new_name}, room=sid)
 
-@socketio.on('request_speaker_profiles')
-def req_prof(): emit('initial_speaker_profiles', {}, room=request.sid)
+@socketio.on('update_speaker_color')
+def handle_speaker_color_update(data):
+    sid = request.sid
+    id = data.get('id'); color = data.get('color')
+    if sid in session_speakers and id in session_speakers[sid]:
+        session_speakers[sid][id]["color"] = color
+        emit("speaker_color_changed", {"id": id, "color": color}, room=sid)
 
 if __name__ == "__main__":
     socketio.run(app, host='0.0.0.0', port=7860, allow_unsafe_werkzeug=True)