Update app.py

app.py

@@ -1,9 +1,11 @@
 import discord
 import logging
 import os
+import re
 import asyncio
 from huggingface_hub import InferenceClient
 from googleapiclient.discovery import build
+from youtube_transcript_api import YouTubeTranscriptApi
 from dotenv import load_dotenv
 
 # 환경 변수 로드
@@ -29,9 +31,6 @@ youtube_service = build('youtube', 'v3', developerKey=YOUTUBE_API_KEY)
 # 특정 채널 ID
 SPECIFIC_CHANNEL_ID = int(os.getenv("DISCORD_CHANNEL_ID"))
 
-# 대화 히스토리를 저장할 전역 변수
-conversation_history = []
-
 class MyClient(discord.Client):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -43,7 +42,7 @@ class MyClient(discord.Client):
         # 봇이 시작될 때 안내 메시지를 전송
         channel = self.get_channel(SPECIFIC_CHANNEL_ID)
         if channel:
-            await channel.send("트렌드 보고 싶은 주제를 입력하세요. 예) 최신 테크 뉴스")
+            await channel.send("유튜브 비디오 URL을 입력하면, 자막과 댓글을 기반으로 답글을 작성합니다.")
 
     async def on_message(self, message):
         if message.author == self.user:
@@ -54,18 +53,17 @@ class MyClient(discord.Client):
             return
         self.is_processing = True
         try:
-            # 주제 키워드 추출
-            keywords = await extract_keywords(message)
-            if keywords:
-                # YouTube API로 최신 트렌드 및 인기 비디오 검색
-                video_details = await search_trending_videos(keywords)
-                if video_details:
-                    # 요청자와의 쓰레드 생성 및 결과 전송
-                    await create_thread_and_send_results(message, keywords, video_details)
+            video_id = extract_video_id(message.content)
+            if video_id:
+                transcript = await get_video_transcript(video_id)
+                comments = await get_video_comments(video_id)
+                if comments and transcript:
+                    replies = await generate_replies(comments, transcript)
+                    await create_thread_and_send_replies(message, video_id, comments, replies)
                 else:
-                    await message.channel.send(f"**{keywords}**에 대한 최신 트렌드 비디오를 찾을 수 없습니다.")
+                    await message.channel.send("자막이나 댓글을 가져올 수 없습니다.")
             else:
-                await message.channel.send("키워드를 추출할 수 없습니다.")
+                await message.channel.send("유효한 유튜브 비디오 URL을 제공해 주세요.")
         finally:
             self.is_processing = False
 
@@ -75,56 +73,84 @@ class MyClient(discord.Client):
             isinstance(message.channel, discord.Thread) and message.channel.parent_id == SPECIFIC_CHANNEL_ID
         )
 
-async def extract_keywords(message):
-    user_input = message.content
-    system_prompt = "다음 문장의 의미에 맞는 영문 키워드를 추출하세요: "
-    
-    logging.debug(f'Extracting keywords from user input: {user_input}')
-
-    messages = [{"role": "system", "content": system_prompt + user_input}]
-    logging.debug(f'Messages to be sent to the model: {messages}')
-
-    loop = asyncio.get_event_loop()
-    response = await loop.run_in_executor(None, lambda: hf_client.chat_completion(
-        messages, max_tokens=10, temperature=0.7, top_p=0.85))
-
-    # Hugging Face 응답 파싱
-    if response.choices and response.choices[0].message:
-        keywords = response.choices[0].message['content'].strip()
-    else:
-        keywords = ""
-    logging.debug(f'Extracted keywords: {keywords}')
-    return keywords
-
-async def search_trending_videos(keywords):
-    response = youtube_service.search().list(
-        q=keywords,
+def extract_video_id(url):
+    """
+    YouTube 비디오 URL에서 비디오 ID를 추출합니다.
+    """
+    video_id = None
+    youtube_regex = (
+        r'(https?://)?(www\.)?'
+        '(youtube|youtu|youtube-nocookie)\.(com|be)/'
+        '(watch\?v=|embed/|v/|.+\?v=)?([^&=%\?]{11})')
+
+    match = re.match(youtube_regex, url)
+    if match:
+        video_id = match.group(6)
+    logging.debug(f'Extracted video ID: {video_id}')
+    return video_id
+
+async def get_video_transcript(video_id):
+    """
+    YouTube 비디오의 자막을 가져옵니다.
+    """
+    try:
+        transcript = YouTubeTranscriptApi.get_transcript(video_id)
+        transcript_text = " ".join([entry['text'] for entry in transcript])
+        logging.debug(f'Fetched transcript: {transcript_text}')
+        return transcript_text
+    except Exception as e:
+        logging.error(f'Error fetching transcript: {e}')
+        return None
+
+async def get_video_comments(video_id):
+    """
+    YouTube 비디오의 댓글을 가져옵니다.
+    """
+    comments = []
+    response = youtube_service.commentThreads().list(
         part='snippet',
-        type='video',
-        order='viewCount',  # 인기 있는 비디오 기준으로 정렬
-        maxResults=5  # 원하는 결과 수
+        videoId=video_id,
+        maxResults=100  # 최대 100개의 댓글 가져오기
     ).execute()
-    
-    video_details = []
+
     for item in response.get('items', []):
-        video_details.append({
-            'title': item['snippet']['title'],
-            'url': f"https://www.youtube.com/watch?v={item['id']['videoId']}",
-            'description': item['snippet']['description'],
-            'thumbnail': item['snippet']['thumbnails']['high']['url']
-        })
-    return video_details
-
-async def create_thread_and_send_results(message, keywords, video_details):
-    # 쓰레드 생성
-    thread = await message.channel.create_thread(name=f"{message.author.name}의 트렌드 검색 결과", message=message)
-    if video_details:
-        message_content = f"**{keywords}**에 대한 최신 트렌드 비디오 {len(video_details)}개를 찾았습니다:"
-        await thread.send(message_content)
-        for video in video_details:
-            embed = discord.Embed(title=video['title'], description=video['description'], url=video['url'])
-            embed.set_thumbnail(url=video['thumbnail'])
-            await thread.send(embed=embed)
+        comment = item['snippet']['topLevelComment']['snippet']['textOriginal']
+        comments.append(comment)
+    
+    logging.debug(f'Fetched comments: {comments}')
+    return comments
+
+async def generate_replies(comments, transcript):
+    """
+    댓글과 자막을 기반으로 LLM 답글을 생성합니다.
+    """
+    replies = []
+    for comment in comments:
+        messages = [
+            {"role": "system", "content": f"비디오 자막: {transcript}"},
+            {"role": "user", "content": comment}
+        ]
+        loop = asyncio.get_event_loop()
+        response = await loop.run_in_executor(None, lambda: hf_client.chat_completion(
+            messages, max_tokens=50, temperature=0.7, top_p=0.85))
+        
+        if response.choices and response.choices[0].message:
+            reply = response.choices[0].message['content'].strip()
+        else:
+            reply = "답글을 생성할 수 없습니다."
+        replies.append(reply)
+    
+    logging.debug(f'Generated replies: {replies}')
+    return replies
+
+async def create_thread_and_send_replies(message, video_id, comments, replies):
+    """
+    댓글과 답글을 새로운 쓰레드에 전송합니다.
+    """
+    thread = await message.channel.create_thread(name=f"{message.author.name}의 댓글 답글", message=message)
+    for comment, reply in zip(comments, replies):
+        embed = discord.Embed(description=f"**댓글**: {comment}\n**답글**: {reply}")
+        await thread.send(embed=embed)
 
 if __name__ == "__main__":
     discord_client = MyClient(intents=intents)