Add websearch module

README.md

@@ -155,6 +155,7 @@ cd SillyTavern-extras
 | `edge-tts`    | [Microsoft Edge TTS client](https://github.com/rany2/edge-tts)      |
 | `coqui-tts`   | [Coqui TTS server](https://github.com/coqui-ai/TTS)                 |
 | `rvc`         | Real-time voice cloning                                             |
+| `websearch`   | Google search using Selenium headless browser                       |
 
 ## Additional options
 | Flag                     | Description                                                            |
@@ -556,3 +557,14 @@ loadchar (string, required): The URL of the character's image. The URL should po
 `GET /api/talkinghead/result_feed`
 #### **Output**
 Animated transparent image
+
+### Perform web search
+`POST /api/websearch`
+#### **Input**
+```
+{ "query": "what is beauty?" }
+```
+#### **Output**
+```
+{ "results": "that would fall within the purview of your conundrums of philosophy" }
+```

modules/websearch/script.py

@@ -0,0 +1,48 @@
+from selenium import webdriver
+from selenium.webdriver.common.by import By
+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
+
+
+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("--lang=en-GB")
+        return webdriver.Chrome(service=chromeService, options=options)
+    except:
+        print("Chrome not found, using Firefox instead.")
+        firefoxService = FirefoxService()
+        options = FirefoxOptions()
+        options.add_argument("--headless")
+        options.set_preference("intl.accept_languages", "en,en_US")
+        return webdriver.Firefox(service=firefoxService, options=options)
+
+
+def search_google(query: str) -> str:
+    driver = get_driver()
+    print("Searching for " + query + "...")
+    driver.get("https://google.com/search?hl=en&q=" + query)
+    text = ''
+    # Answer box
+    for el in driver.find_elements(By.CSS_SELECTOR, '.hgKElc'):
+        if el and el.text:
+            text += el.text + '\n'
+    # Knowledge panel
+    for el in driver.find_elements(By.CSS_SELECTOR, '.hgKElc'):
+        if el and el.text:
+            text += el.text + '\n'
+    # Page snippets
+    for el in driver.find_elements(By.CSS_SELECTOR, '.yDYNvb.lyLwlc'):
+        if el and el.text:
+            text += el.text + '\n'
+    print("Found: " + text)
+    driver.quit()
+    return text

requirements.txt

@@ -22,6 +22,7 @@ matplotlib
 vosk
 sounddevice
 openai-whisper
+selenium
 
 fastapi
 wxpython; sys_platform == 'win32' or sys_platform == 'darwin'

server.py

@@ -1100,6 +1100,20 @@ def chromadb_import():
     return jsonify({"count": len(ids)})
 
 
+@app.route("/api/websearch", methods=["POST"])
+@require_module("websearch")
+def api_websearch():
+    data = request.get_json()
+    if "query" not in data or not isinstance(data["query"], str):
+        abort(400, '"query" is required')
+
+    query = data["query"]
+    import modules.websearch.script as websearch
+    results = websearch.search_google(query)
+
+    return jsonify({"results": results})
+
+
 if args.share:
     from flask_cloudflared import _run_cloudflared
     import inspect