Update app.py

app.py

@@ -6,6 +6,10 @@ import whisper
 import torch
 from gtts import gTTS
 import os
+import yt_dlp
+import re
+
+
 hf_token = os.getenv("HF_TOKEN")
 app = FastAPI()
 
@@ -25,6 +29,33 @@ whisper_model = whisper.load_model("base")
 # Lưu hội thoại
 conversation = [{"role": "system", "content": "Bạn là một trợ lý AI. Hãy trả lời ngắn gọn, súc tích, tối đa 2 câu."}]
 
+# Hàm trích xuất tên bài hát từ văn bản
+def extract_song_name(text):
+    import re
+    match = re.search(r"(bài|bài hát|nghe nhạc|mở nhạc)\s+(.*)", text.lower())
+    if match:
+        return match.group(2).strip()
+    return None
+
+def download_youtube_as_wav(song_name, output_path="song.wav"):
+    search_query = f"ytsearch1:{song_name}"
+    ydl_opts = {
+        'format': 'bestaudio/best',
+        'outtmpl': 'temp_audio.%(ext)s',
+        'postprocessors': [{
+            'key': 'FFmpegExtractAudio',
+            'preferredcodec': 'wav',
+            'preferredquality': '192',
+        }],
+        'quiet': True,
+    }
+    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+        ydl.download([search_query])
+    if os.path.exists("temp_audio.wav"):
+        os.rename("temp_audio.wav", output_path)
+        return output_path
+    return None
+
 class ChatRequest(BaseModel):
     message: str
 
@@ -45,23 +76,31 @@ async def chat(request: ChatRequest):
 # Endpoint voice chat + TTS
 @app.post("/voice_chat")
 async def voice_chat(file: UploadFile = File(...)):
-    # Lưu file tạm
     file_location = f"temp_{file.filename}"
     with open(file_location, "wb") as f:
         f.write(await file.read())
 
-    # Chuyển âm thanh thành text
     result = whisper_model.transcribe(file_location, language="vi")
     user_text = result["text"]
+    os.remove(file_location)
+
+    # Kiểm tra yêu cầu mở nhạc
+    if any(kw in user_text.lower() for kw in ["nghe nhạc", "mở bài hát", "bài hát", "bài"]):
+        song_name = extract_song_name(user_text)
+        if song_name:
+            wav_path = download_youtube_as_wav(song_name)
+            if wav_path:
+                return FileResponse(wav_path, media_type="audio/wav", filename="song.wav")
+            else:
+                return {"error": "Không tìm thấy hoặc tải được bài hát."}
 
-    # Gọi mô hình Qwen để trả lời
+    # Nếu không phải yêu cầu mở nhạc → xử lý như cũ
     conversation.append({"role": "user", "content": user_text})
     text = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
     model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
     response_text = generate_full_response(model_inputs)
     conversation.append({"role": "assistant", "content": response_text})
 
-    # Tạo file âm thanh từ phản hồi
     tts = gTTS(response_text, lang="vi")
     audio_file = "response.mp3"
     tts.save(audio_file)