Update main.py

main.py

@@ -13,6 +13,8 @@ import logging
 import os
 import shutil
 from pathlib import Path
+import tempfile
+from tenacity import retry, stop_after_attempt, wait_fixed
 
 app = FastAPI()
 
@@ -39,6 +41,7 @@ class TranscriptResponse(BaseModel):
     error: str | None
     processing_time: float
 
+@retry(stop=stop_after_attempt(3), wait=wait_fixed(2))
 def init_driver():
     options = Options()
     options.add_argument("--headless=new")
@@ -47,6 +50,11 @@ def init_driver():
     options.add_argument("--disable-gpu")
     options.add_argument("--disable-extensions")
 
+    # Create a unique temporary user data directory
+    user_data_dir = tempfile.mkdtemp()
+    options.add_argument(f"--user-data-dir={user_data_dir}")
+    logger.info(f"Using temporary user data directory: {user_data_dir}")
+
     possible_chrome_paths = [
         "/usr/bin/google-chrome",
         "/usr/bin/google-chrome-stable",
@@ -75,22 +83,26 @@ def init_driver():
         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
+        return driver, user_data_dir
     except Exception as e:
         logger.error(f"Driver initialization failed: {str(e)}")
+        # Clean up the temporary directory in case of failure
+        if Path(user_data_dir).exists():
+            shutil.rmtree(user_data_dir, ignore_errors=True)
         raise Exception(f"Driver initialization failed: {str(e)}")
 
 @app.post("/transcript", response_model=TranscriptResponse)
 async def get_transcript(request: VideoRequest):
     start_time = time.time()
     driver = None
+    user_data_dir = None
 
     try:
         video_url = request.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()
+        driver, user_data_dir = init_driver()
         logger.info(f"Processing URL: {video_url}")
         driver.get(video_url)
 
@@ -162,6 +174,9 @@ async def get_transcript(request: VideoRequest):
     finally:
         if driver:
             driver.quit()
+        if user_data_dir and Path(user_data_dir).exists():
+            shutil.rmtree(user_data_dir, ignore_errors=True)
+            logger.info(f"Cleaned up temporary user data directory: {user_data_dir}")
 
 @app.get("/health")
 def health_check():
@@ -180,4 +195,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)), workers=1)