From 2ebbcdf0b81670e2b987e4d44f1efee8f68b3b43 Mon Sep 17 00:00:00 2001
From: Heiner Lohaus <hlohaus@users.noreply.github.com>
Date: Fri, 17 Nov 2023 04:50:37 +0100
Subject: Remove virtualdisplay and add headless mode

---
 g4f/Provider/needs_auth/Bard.py | 20 +++++---------------
 g4f/Provider/needs_auth/Poe.py  | 22 +++++++---------------
 2 files changed, 12 insertions(+), 30 deletions(-)

(limited to 'g4f/Provider/needs_auth')

diff --git a/g4f/Provider/needs_auth/Bard.py b/g4f/Provider/needs_auth/Bard.py
index 6cb40c90..7f73f1b3 100644
--- a/g4f/Provider/needs_auth/Bard.py
+++ b/g4f/Provider/needs_auth/Bard.py
@@ -19,17 +19,12 @@ class Bard(BaseProvider):
         stream: bool,
         proxy: str = None,
         browser: WebDriver = None,
-        hidden_display: bool = True,
+        user_data_dir: str = None,
+        headless: bool = True,
         **kwargs
     ) -> CreateResult:
         prompt = format_prompt(messages)
-        if browser:
-            driver = browser
-        else:
-            if hidden_display:
-                driver, display = get_browser(None, True, proxy)
-            else:
-                driver = get_browser(None, False, proxy)
+        driver = browser if browser else get_browser(user_data_dir, headless, proxy)
 
         from selenium.webdriver.common.by import By
         from selenium.webdriver.support.ui import WebDriverWait
@@ -43,9 +38,6 @@ class Bard(BaseProvider):
             # Reopen browser for login
             if not browser:
                 driver.quit()
-                # New browser should be visible
-                if hidden_display:
-                    display.stop()
                 driver = get_browser(None, False, proxy)
                 driver.get(f"{cls.url}/chat")
                 wait = WebDriverWait(driver, 240)
@@ -83,9 +75,7 @@ XMLHttpRequest.prototype.open = function(method, url) {
                 else:
                     time.sleep(0.1)
         finally:
-            driver.close()
             if not browser:
+                driver.close()
                 time.sleep(0.1)
-                driver.quit()
-            if hidden_display:
-                display.stop()
\ No newline at end of file
+                driver.quit()
\ No newline at end of file
diff --git a/g4f/Provider/needs_auth/Poe.py b/g4f/Provider/needs_auth/Poe.py
index 6fbf7fd4..a894bcb1 100644
--- a/g4f/Provider/needs_auth/Poe.py
+++ b/g4f/Provider/needs_auth/Poe.py
@@ -22,6 +22,7 @@ models = {
 class Poe(BaseProvider):
     url = "https://poe.com"
     working = True
+    needs_auth = True
     supports_gpt_35_turbo = True
     supports_stream = True
 
@@ -33,7 +34,8 @@ class Poe(BaseProvider):
         stream: bool,
         proxy: str = None,
         browser: WebDriver = None,
-        hidden_display: bool = True,
+        user_data_dir: str = None,
+        headless: bool = True,
         **kwargs
     ) -> CreateResult:
         if not model:
@@ -41,13 +43,7 @@ class Poe(BaseProvider):
         elif model not in models:
             raise ValueError(f"Model are not supported: {model}")
         prompt = format_prompt(messages)
-        if browser:
-            driver = browser
-        else:
-            if hidden_display:
-                driver, display = get_browser(None, True, proxy)
-            else:
-                driver = get_browser(None, False, proxy)
+        driver = browser if browser else get_browser(user_data_dir, headless, proxy)
 
         script = """
 window._message = window._last_message = "";
@@ -80,14 +76,12 @@ window.WebSocket = ProxiedWebSocket;
 
         try:
             driver.get(f"{cls.url}/{models[model]['name']}")
-            wait = WebDriverWait(driver, 10 if hidden_display else 240)
+            wait = WebDriverWait(driver, 10 if headless else 240)
             wait.until(EC.visibility_of_element_located((By.CSS_SELECTOR, "textarea[class^='GrowingTextArea']")))
         except:
             # Reopen browser for login
             if not browser:
                 driver.quit()
-                if hidden_display:
-                    display.stop()
                 driver = get_browser(None, False, proxy)
                 driver.get(f"{cls.url}/{models[model]['name']}")
                 wait = WebDriverWait(driver, 240)
@@ -121,9 +115,7 @@ if(window._message && window._message != window._last_message) {
                 else:
                     time.sleep(0.1)
         finally:
-            driver.close()
             if not browser:
+                driver.close()
                 time.sleep(0.1)
-                driver.quit()
-            if hidden_display:
-                display.stop()
\ No newline at end of file
+                driver.quit()
\ No newline at end of file
-- 
cgit v1.2.3