Merge pull request #223 from yatikakain/youtube-context-feature

[Feature] Implemented "Generate Quiz from YouTube Video" Context Menu

backend/server.py

@@ -2,6 +2,9 @@
 from flask_cors import CORS
 from pprint import pprint
 import nltk
+import subprocess
+import os
+import glob
 
 from sklearn.metrics.pairwise import cosine_similarity
 from sklearn.feature_extraction.text import TfidfVectorizer
@@ -406,6 +409,56 @@ def upload_file():
 def hello():
     return "The server is working fine"
 
+def clean_transcript(file_path):
+    """Extracts and cleans transcript from a VTT file."""
+    with open(file_path, "r", encoding="utf-8") as file:
+        lines = file.readlines()
+
+    transcript_lines = []
+    skip_metadata = True  # Skip lines until we reach actual captions
+
+    for line in lines:
+        line = line.strip()
+
+        # Skip metadata lines like "Kind: captions" or "Language: en"
+        if line.lower().startswith(("kind:", "language:", "webvtt")):
+            continue
+
+        # Detect timestamps like "00:01:23.456 --> 00:01:25.789"
+        if "-->" in line:
+            skip_metadata = False  # Now real captions start
+            continue
+
+        if not skip_metadata:
+            # Remove formatting tags like <c>...</c> and <00:00:00.000>
+            line = re.sub(r"<[^>]+>", "", line)
+            transcript_lines.append(line)
+
+    return " ".join(transcript_lines).strip()
+
+@app.route('/getTranscript', methods=['GET'])
+def get_transcript():
+    video_id = request.args.get('videoId')
+    if not video_id:
+        return jsonify({"error": "No video ID provided"}), 400
+
+    subprocess.run(["yt-dlp", "--write-auto-sub", "--sub-lang", "en", "--skip-download",
+                "--sub-format", "vtt", "-o", f"subtitles/{video_id}.vtt", f"https://www.youtube.com/watch?v={video_id}"],
+               check=True, capture_output=True, text=True)
+
+    # Find the latest .vtt file in the "subtitles" folder
+    subtitle_files = glob.glob("subtitles/*.vtt")
+    if not subtitle_files:
+        return jsonify({"error": "No subtitles found"}), 404
+
+    latest_subtitle = max(subtitle_files, key=os.path.getctime)
+    transcript_text = clean_transcript(latest_subtitle)
+
+    # Optional: Clean up the file after reading
+    os.remove(latest_subtitle)
+
+    return jsonify({"transcript": transcript_text})
 
 if __name__ == "__main__":
+    os.makedirs("subtitles", exist_ok=True)
     app.run()

extension/public/background.js

@@ -1,14 +1,20 @@
-// Create context menu on installation
-chrome.runtime.onInstalled.addListener(() => {
-    chrome.contextMenus.create({
-      id: "askExtension",
-      title: "Generate Quiz with Selected Text",
-      contexts: ["selection"]
-    });
+// Create context menus
+  chrome.contextMenus.create({
+    id: "askExtension",
+    title: "Generate Quiz with Selected Text",
+    contexts: ["selection"]
   });
-
-  // Handle context menu clicks
-  chrome.contextMenus.onClicked.addListener(async (info, tab) => {
+
+  // YouTube-specific quiz generator
+  chrome.contextMenus.create({
+      id: "generateQuizFromYouTube",
+      title: "Generate Quiz from YouTube Video",
+      contexts: ["page"],
+      documentUrlPatterns: ["*://www.youtube.com/watch?v=*"]
+  });
+
+// Handle context menu clicks
+chrome.contextMenus.onClicked.addListener(async (info, tab) => {
     if (info.menuItemId === "askExtension" && info.selectionText) {
       try {
         // Store the selected text first
@@ -51,10 +57,43 @@ chrome.runtime.onInstalled.addListener(() => {
         console.error("Error in context menu handler:", error);
       }
     }
-  });
-
-  // Listen for messages from content script
-  chrome.runtime.onMessage.addListener((request, sender, sendResponse) => {
+
+    if (info.menuItemId === "generateQuizFromYouTube") {
+        try {
+            // Check if script has already been injected
+            const [{ result }] = await chrome.scripting.executeScript({
+                target: { tabId: tab.id },
+                func: () => window.youtubeScriptLoaded || false
+            });
+
+            if (!result) {
+                await chrome.scripting.executeScript({
+                    target: { tabId: tab.id },
+                    files: ["youtubeContentScript.js"]
+                });
+            } else {
+                console.log("YouTube script already loaded");
+            }
+            // Add badge notification
+            chrome.action.setBadgeText({ 
+                text: "!"
+            });
+            chrome.action.setBadgeBackgroundColor({ 
+                color: "#FF005C"
+            });
+
+            // Clear the badge after 2 seconds
+            setTimeout(() => {
+                chrome.action.setBadgeText({ text: "" });
+            }, 2000);
+        } catch (error) {
+            console.error("Error in YouTube quiz generation handler:", error);
+        }
+    }
+});
+
+// Listen for messages from content script
+chrome.runtime.onMessage.addListener((request, sender, sendResponse) => {
     if (request.type === "TEXT_SELECTED") {
       chrome.storage.local.set({ 
         selectedText: request.text 
@@ -89,12 +128,3 @@ chrome.storage.onChanged.addListener((changes, namespace) => {
     }
 });
 
-
-  // Optional: Handle extension install/update
-  chrome.runtime.onInstalled.addListener((details) => {
-    if (details.reason === "install") {
-      console.log("Extension installed");
-    } else if (details.reason === "update") {
-      console.log("Extension updated");
-    }
-  });

extension/public/youtubeContentScript.js

@@ -0,0 +1,38 @@
+// Function to extract YouTube video ID from the URL
+function getYouTubeVideoId() {
+    const urlParams = new URLSearchParams(window.location.search);
+    return urlParams.get("v");
+}
+
+(() => {
+    const videoId = getYouTubeVideoId();
+
+    if (videoId) {
+        console.log("Extracted YouTube Video ID:", videoId);
+
+        fetch(`http://localhost:5000/getTranscript?videoId=${videoId}`)
+            .then(response => response.json())
+            .then(data => {
+                if (data.transcript) {
+                    console.log("Transcript received:", data.transcript);
+                    chrome.storage.local.set({ videoTranscript: data.transcript }, () => {
+                        console.log("Transcript saved in storage.");
+                    });
+                    generateQuestions(data.transcript);
+                } else {
+                    console.error("Failed to fetch transcript:", data.error);
+                }
+            })
+            .catch(error => {
+                console.error("Error fetching transcript:", error);
+            });
+    }
+})();
+
+// Function to generate quiz questions from transcript
+function generateQuestions(transcript) {
+    console.log("Generating questions from transcript...");
+    chrome.storage.local.set({ "selectedText": transcript }, () => {
+        console.log("Transcript stored for quiz generation.");
+    });
+}