From 9239cadd8b3e539a7d5da0eb22b2b047417fb426 Mon Sep 17 00:00:00 2001 From: Heiner Lohaus Date: Sun, 19 Nov 2023 05:36:04 +0100 Subject: Add Response Handler to OpenaiChat Update Providers with WebDriver Add WebDriverSession helper Use native streaming in curl_cffi --- g4f/Provider/needs_auth/Theb.py | 92 ++++++++++++++++++++++------------------- 1 file changed, 49 insertions(+), 43 deletions(-) (limited to 'g4f/Provider/needs_auth/Theb.py') diff --git a/g4f/Provider/needs_auth/Theb.py b/g4f/Provider/needs_auth/Theb.py index 89c69727..cf33f0c6 100644 --- a/g4f/Provider/needs_auth/Theb.py +++ b/g4f/Provider/needs_auth/Theb.py @@ -4,7 +4,7 @@ import time from ...typing import CreateResult, Messages from ..base_provider import BaseProvider -from ..helper import WebDriver, format_prompt, get_browser +from ..helper import WebDriver, WebDriverSession, format_prompt models = { "theb-ai": "TheB.AI", @@ -44,26 +44,60 @@ class Theb(BaseProvider): messages: Messages, stream: bool, proxy: str = None, - browser: WebDriver = None, - headless: bool = True, + web_driver: WebDriver = None, + virtual_display: bool = True, **kwargs ) -> CreateResult: if model in models: model = models[model] prompt = format_prompt(messages) - driver = browser if browser else get_browser(None, headless, proxy) + web_session = WebDriverSession(web_driver, virtual_display=virtual_display, proxy=proxy) + with web_session as driver: + from selenium.webdriver.common.by import By + from selenium.webdriver.support.ui import WebDriverWait + from selenium.webdriver.support import expected_conditions as EC + from selenium.webdriver.common.keys import Keys - from selenium.webdriver.common.by import By - from selenium.webdriver.support.ui import WebDriverWait - from selenium.webdriver.support import expected_conditions as EC - from selenium.webdriver.common.keys import Keys + # Register fetch hook + script = """ +window._fetch = window.fetch; +window.fetch = (url, options) => { + // Call parent fetch method + const result = window._fetch(url, options); + if (!url.startsWith("/api/conversation")) { + return result; + } + // Load response reader + result.then((response) => { + if (!response.body.locked) { + window._reader = response.body.getReader(); + } + }); + // Return dummy response + return new Promise((resolve, reject) => { + resolve(new Response(new ReadableStream())) + }); +} +window._last_message = ""; +""" + driver.execute_cdp_cmd("Page.addScriptToEvaluateOnNewDocument", { + "source": script + }) + + try: + driver.get(f"{cls.url}/home") + wait = WebDriverWait(driver, 5) + wait.until(EC.visibility_of_element_located((By.ID, "textareaAutosize"))) + except: + driver = web_session.reopen() + driver.execute_cdp_cmd("Page.addScriptToEvaluateOnNewDocument", { + "source": script + }) + driver.get(f"{cls.url}/home") + wait = WebDriverWait(driver, 240) + wait.until(EC.visibility_of_element_located((By.ID, "textareaAutosize"))) - - try: - driver.get(f"{cls.url}/home") - wait = WebDriverWait(driver, 10 if headless else 240) - wait.until(EC.visibility_of_element_located((By.TAG_NAME, "body"))) - time.sleep(0.1) + time.sleep(200) try: driver.find_element(By.CSS_SELECTOR, ".driver-overlay").click() driver.find_element(By.CSS_SELECTOR, ".driver-overlay").click() @@ -87,29 +121,6 @@ class Theb(BaseProvider): button = container.find_element(By.CSS_SELECTOR, "button.btn-blue.btn-small.border") button.click() - # Register fetch hook - script = """ -window._fetch = window.fetch; -window.fetch = (url, options) => { - // Call parent fetch method - const result = window._fetch(url, options); - if (!url.startsWith("/api/conversation")) { - return result; - } - // Load response reader - result.then((response) => { - if (!response.body.locked) { - window._reader = response.body.getReader(); - } - }); - // Return dummy response - return new Promise((resolve, reject) => { - resolve(new Response(new ReadableStream())) - }); -} -window._last_message = ""; -""" - driver.execute_script(script) # Submit prompt wait.until(EC.visibility_of_element_located((By.ID, "textareaAutosize"))) @@ -150,9 +161,4 @@ return ''; elif chunk != "": break else: - time.sleep(0.1) - finally: - if not browser: - driver.close() - time.sleep(0.1) - driver.quit() \ No newline at end of file + time.sleep(0.1) \ No newline at end of file -- cgit v1.2.3