Add colab detection util

modules/utils.py

@@ -12,4 +12,12 @@ def silence_log():
             yield new_target
     finally:
         sys.stdout = old_stdout
-        sys.stderr = old_stderr
+        sys.stderr = old_stderr
+
+
+def is_colab():
+    try:
+        from google import colab
+        return True
+    except:
+        return False

modules/websearch/script.py

@@ -4,20 +4,26 @@ from selenium.webdriver.chrome.options import Options as ChromeOptions
 from selenium.webdriver.firefox.options import Options as FirefoxOptions
 from selenium.webdriver.chrome.service import Service as ChromeService
 from selenium.webdriver.firefox.service import Service as FirefoxService
+from modules.utils import is_colab
 import atexit
 
 
 def get_driver():
     try:
         print("Initializing Chrome driver...")
-        chromeService = ChromeService()
         options = ChromeOptions()
         options.add_argument('--disable-infobars')
         options.add_argument("--headless")
         options.add_argument("--disable-gpu")
         options.add_argument("--no-sandbox")
+        options.add_argument('--disable-dev-shm-usage')
         options.add_argument("--lang=en-GB")
-        return webdriver.Chrome(service=chromeService, options=options)
+
+        if is_colab():
+            return webdriver.Chrome('chromedriver', options=options)
+        else:
+            chromeService = ChromeService()
+            return webdriver.Chrome(service=chromeService, options=options)
     except:
         print("Chrome not found, using Firefox instead.")
         firefoxService = FirefoxService()