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 Service
+from selenium.webdriver.chrome.service import Oceans
 from selenium.webdriver.chrome.options import Options
 from selenium.webdriver.common.by import By
 from selenium.webdriver.support.ui import WebDriverWait
@@ -13,6 +13,9 @@ import logging
 import os
 import shutil
 from pathlib import Path
+import re
+from contextlib import contextmanager
+import signal
 
 app = FastAPI()
 
@@ -39,7 +42,16 @@ class TranscriptResponse(BaseModel):
     error: str | None
     processing_time: float
 
-# Driver init and route handlers here...
+@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()
@@ -47,36 +59,35 @@ def init_driver():
     options.add_argument("--no-sandbox")
     options.add_argument("--disable-dev-shm-usage")
     options.add_argument("--disable-gpu")
-    
-    # Try multiple possible Chrome binary locations
+
     possible_chrome_paths = [
         "/usr/bin/google-chrome",
         "/usr/bin/google-chrome-stable",
-        "/usr/lib/chromium-browser/chrome",
-        "/usr/bin/chromium"
     ]
     chrome_path = None
     for path in possible_chrome_paths:
         if os.path.exists(path):
             chrome_path = path
             break
-    
+
     if not chrome_path:
         logger.error(f"No Chrome binary found in paths: {possible_chrome_paths}")
         raise Exception(f"No Chrome binary found in paths: {possible_chrome_paths}")
-    
+
     options.binary_location = chrome_path
     logger.info(f"Using Chrome binary: {chrome_path}")
-    
+
     try:
-        chromedriver_path = "/usr/bin/chromedriver"
-        if not os.path.exists(chromedriver_path):
+        chromedriver_path = shutil.which("chromedriver")
+        if not chromedriver_path or not os.path.exists(chromedriver_path):
             logger.error(f"ChromeDriver not found at {chromedriver_path}")
             raise Exception(f"ChromeDriver not found at {chromedriver_path}")
-        
+
         service = Service(executable_path=chromedriver_path)
         driver = webdriver.Chrome(service=service, options=options)
-        logger.info("ChromeDriver initialized successfully")
+        chrome_version = driver.capabilities["browserVersion"]
+        chromedriver_version = driver.capabilities["chrome"]["chromedriverVersion"].split()[0]
+        logger.info(f"Chrome version: {chrome_version}, ChromeDriver version: {chromedriver_version}")
         return driver
     except Exception as e:
         logger.error(f"Driver initialization failed: {str(e)}")
@@ -86,65 +97,74 @@ def init_driver():
 async def get_transcript(request: VideoRequest):
     start_time = time.time()
     driver = None
-    
+
     try:
         video_url = request.url
-        if not ("youtube.com" in video_url or "youtu.be" in video_url):
+        if not re.match(r"(?:https?://)?(?:www\.)?(?:youtube\.com/watch\?v=|youtu\.be/)([a-zA-Z0-9_-]{11})", 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)
 
-        # Handle cookie consent if it appears
-        try:
-            cookie_button = WebDriverWait(driver, 5).until(
-                EC.element_to_be_clickable((By.XPATH, "//*[contains(text(), 'Accept all')]"))
+        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"))
             )
-            cookie_button.click()
-            logger.info("Accepted cookies")
-        except TimeoutException:
-            logger.info("No cookie consent found")
-            pass
-
-        # Click more button
-        more_button = WebDriverWait(driver, 10).until(
-            EC.element_to_be_clickable((By.ID, "expand"))
-        )
-        driver.execute_script("arguments[0].click();", more_button)
-
-        # 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)
+            driver.execute_script("arguments[0].click();", more_button)
 
-        # Wait for transcript
-        WebDriverWait(driver, 15).until(
-            EC.presence_of_element_located((By.ID, "segments-container"))
-        )
+            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)
 
-        # Extract 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
+            logger.info("Waiting for transcript segments")
+            WebDriverWait(driver, 15).until(
+                EC.presence_of_element_located((By.ID, "segments-container"))
+            )
 
-        if not transcript:
-            raise HTTPException(status_code=404, detail="No transcript available")
+            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
+            )
 
+    except TimeoutError:
+        error_msg = "Request timed out"
+        logger.error(error_msg)
         return TranscriptResponse(
-            success=True,
-            transcript=transcript,
-            error=None,
+            success=False,
+            transcript=None,
+            error=error_msg,
             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)
@@ -165,7 +185,7 @@ async def get_transcript(request: VideoRequest):
     finally:
         if driver:
             driver.quit()
- 
+
 @app.get("/health")
 def health_check():
     chrome_path = shutil.which("google-chrome")
@@ -183,4 +203,4 @@ async def root():
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", 7860)))
+    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", 7860)))