Update main.py

main.py

@@ -2,7 +2,7 @@ from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from selenium import webdriver
-from selenium.webdriver.chrome.service import Oceans
+from selenium.webdriver.chrome.service import Service
 from selenium.webdriver.chrome.options import Options
 from selenium.webdriver.common.by import By
 from selenium.webdriver.support.ui import WebDriverWait
@@ -13,9 +13,6 @@ import logging
 import os
 import shutil
 from pathlib import Path
-import re
-from contextlib import contextmanager
-import signal
 
 app = FastAPI()
 
@@ -42,17 +39,6 @@ class TranscriptResponse(BaseModel):
     error: str | None
     processing_time: float
 
-@contextmanager
-def timeout(seconds):
-    def signal_handler(signum, frame):
-        raise TimeoutError("Operation timed out")
-    signal.signal(signal.SIGALRM, signal_handler)
-    signal.alarm(seconds)
-    try:
-        yield
-    finally:
-        signal.alarm(0)
-
 def init_driver():
     options = Options()
     options.add_argument("--headless=new")
@@ -100,71 +86,57 @@ async def get_transcript(request: VideoRequest):
 
     try:
         video_url = request.url
-        if not re.match(r"(?:https?://)?(?:www\.)?(?:youtube\.com/watch\?v=|youtu\.be/)([a-zA-Z0-9_-]{11})", video_url):
+        if not ("youtube.com" in video_url or "youtu.be" in video_url):
             raise HTTPException(status_code=400, detail="Invalid YouTube URL")
 
         driver = init_driver()
         logger.info(f"Processing URL: {video_url}")
         driver.get(video_url)
 
-        with timeout(60):
-            try:
-                cookie_button = WebDriverWait(driver, 3).until(
-                    EC.element_to_be_clickable((By.XPATH, "//*[contains(text(), 'Accept all')]"))
-                )
-                cookie_button.click()
-                logger.info("Accepted cookies")
-            except TimeoutException:
-                logger.info("No cookie consent found")
-
-            logger.info("Clicking 'Show more' button")
-            more_button = WebDriverWait(driver, 10).until(
-                EC.element_to_be_clickable((By.ID, "expand"))
+        try:
+            cookie_button = WebDriverWait(driver, 5).until(
+                EC.element_to_be_clickable((By.XPATH, "//*[contains(text(), 'Accept all')]"))
             )
-            driver.execute_script("arguments[0].click();", more_button)
+            cookie_button.click()
+            logger.info("Accepted cookies")
+        except TimeoutException:
+            logger.info("No cookie consent found")
+            pass
+
+        more_button = WebDriverWait(driver, 10).until(
+            EC.element_to_be_clickable((By.ID, "expand"))
+        )
+        driver.execute_script("arguments[0].click();", more_button)
 
-            logger.info("Clicking transcript button")
-            transcript_button = WebDriverWait(driver, 10).until(
-                EC.element_to_be_clickable((By.CSS_SELECTOR, "button[aria-label='Show transcript']"))
-            )
-            driver.execute_script("arguments[0].click();", transcript_button)
+        transcript_button = WebDriverWait(driver, 10).until(
+            EC.element_to_be_clickable((By.CSS_SELECTOR, "button[aria-label='Show transcript']"))
+        )
+        driver.execute_script("arguments[0].click();", transcript_button)
 
-            logger.info("Waiting for transcript segments")
-            WebDriverWait(driver, 15).until(
-                EC.presence_of_element_located((By.ID, "segments-container"))
-            )
+        WebDriverWait(driver, 15).until(
+            EC.presence_of_element_located((By.ID, "segments-container"))
+        )
 
-            logger.info("Extracting transcript")
-            segments = driver.find_elements(By.CSS_SELECTOR, "div.ytd-transcript-segment-renderer")
-            transcript = []
-            for segment in segments:
-                try:
-                    text = segment.find_element(By.CLASS_NAME, "segment-text").text.strip()
-                    if text:
-                        transcript.append(text)
-                except:
-                    continue
-
-            if not transcript:
-                raise HTTPException(status_code=404, detail="No transcript available")
-
-            logger.info(f"Extracted {len(transcript)} transcript segments")
-            return TranscriptResponse(
-                success=True,
-                transcript=transcript,
-                error=None,
-                processing_time=time.time() - start_time
-            )
+        segments = driver.find_elements(By.CSS_SELECTOR, "div.ytd-transcript-segment-renderer")
+        transcript = []
+        for segment in segments:
+            try:
+                text = segment.find_element(By.CLASS_NAME, "segment-text").text.strip()
+                if text:
+                    transcript.append(text)
+            except:
+                continue
+
+        if not transcript:
+            raise HTTPException(status_code=404, detail="No transcript available")
 
-    except TimeoutError:
-        error_msg = "Request timed out"
-        logger.error(error_msg)
         return TranscriptResponse(
-            success=False,
-            transcript=None,
-            error=error_msg,
+            success=True,
+            transcript=transcript,
+            error=None,
             processing_time=time.time() - start_time
         )
+
     except TimeoutException as e:
         error_msg = "Timed out waiting for page elements - the video might not have transcripts"
         logger.error(error_msg)