[FIX] update YouTube transcript extraction to handle missing captions and improve error handling

tool.py

@@ -5,6 +5,7 @@ from transformers import pipeline
 import requests
 import io
 from PIL import Image
+from pytubefix import YouTube
 #from dotenv import load_dotenv
 
 #load_dotenv()
@@ -17,7 +18,7 @@ class TranscriptSummarizer(Tool):
 
     def __init__(self, *args, hf_api_key: str = None, **kwargs):
         super().__init__(*args, **kwargs)
-        self.summarizer = pipeline("summarization", model="google/pegasus-xsum")
+        self.summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
         self.api_url = "https://api-inference.huggingface.co/models/ZB-Tech/Text-to-Image"
         self.hf_api_key = hf_api_key
         self.headers = {"Authorization": f"Bearer {self.hf_api_key}"}
@@ -55,7 +56,7 @@ class TranscriptSummarizer(Tool):
             # Concatenate the summaries
             full_summary = "\n".join(summaries)
 
-            key_entities = full_summary.split()[:15]  # Extract first 3 words as key entities
+            key_entities = full_summary.split()[:15]  # Extract first 15 words as key entities
             image_prompt = f"Generate an image related to: {' '.join(key_entities)}, cartoon style"
             image_bytes = self.query({"inputs": image_prompt})
             image = Image.open(io.BytesIO(image_bytes))
@@ -76,17 +77,20 @@ class YouTubeTranscriptExtractor(Tool):
 
     def forward(self, video_url: str) -> str:
         try:
-            from pytubefix import YouTube
             # Create a YouTube object
             yt = YouTube(video_url)
             lang='en'
            # Get the video transcript
             try:
-                transcript = yt.captions[lang].generate_srt_captions()
-            except KeyError:
-                # If the specified language is not available, get the first available transcript
-                transcript = next(iter(yt.captions.values())).generate_srt_captions()
-                lang = next(iter(yt.captions.keys()))
+                if lang in yt.captions:
+                    transcript = yt.captions['en'].generate_srt_captions()
+                else:
+                    transcript = yt.captions.all()[0].generate_srt_captions()
+                    lang = yt.captions.all()[0].code
+            except StopIteration:
+                return "No transcript available for this video."
+            except Exception as e:
+                return f"An unexpected error occurred while accessing captions: {str(e)}"
 
             # Clean up the transcript by removing timestamps and line numbers
             cleaned_transcript = ""