thecookingsenpai 96d48b3470 Initial commit
2023-12-25 13:28:54 +01:00

72 lines
2.5 KiB
Python

import gradio as gr
import modules.shared as shared
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.chrome.service import Service
import urllib
import re
search_access = True
service = Service()
options = webdriver.ChromeOptions()
options.add_argument('headless')
options.add_argument('--disable-infobars')
options.add_argument('--disable-dev-shm-usage')
options.add_argument('--no-sandbox')
options.add_argument('--disable-gpu')
options.add_argument('--remote-debugging-port=9222')
def google_results(query):
driver = webdriver.Chrome(service=service,options=options)
query = urllib.parse.quote_plus(query)
url="https://www.google.com/search?hl=en&q="+query
driver.get(url)
html = driver.find_element(By.CLASS_NAME, 'ULSxyf').text
driver.quit()
return html
def ui():
global search_access
checkbox = gr.Checkbox(value=search_access, label="Enable Google Search")
checkbox.change(fn=update_search_access, inputs=checkbox)
return checkbox, search_access
def update_search_access(checkbox_value):
global search_access
search_access = checkbox_value # assign the value of the checkbox to the variable
return search_access, checkbox_value
def input_modifier(user_input, state):
global search_access
if search_access:
search_query = re.search(r'search\s+"([^"]+)"', user_input, re.IGNORECASE)
if search_query:
query = search_query.group(1)
elif user_input.lower().startswith("search"):
query = user_input.replace("search", "").strip()
else:
query = ""
if not query:
shared.processing_message = "*Typing...*"
return user_input
else:
shared.processing_message = f"*Searching online for {query}*"
state["context"] = state["context"] + "Relevant search results are in the Google search results. Use this info in the response. If no results are found, say so and offer information from the memory."
search_data = google_results(query)
if not search_data:
user_prompt = f"User question: {user_input}\n Google search results: NO RESULTS FOUND"
else:
user_prompt = f"User question: {user_input}\n Google search results: {search_data}"
return str(user_prompt)
shared.processing_message = "*Typing...*"
return user_input
def output_modifier(output):
return output
def bot_prefix_modifier(prefix):
return prefix