8 files changed, 68 insertions, 100 deletions

diff --git a/README.md b/README.md
index 7f721d95..0cba63fb 100644
--- a/README.md
+++ b/README.md
@@ -243,6 +243,41 @@ for message in response:
     print(message)
 ```
 
+##### Using Browser
+
+Some providers using a browser to bypass the bot protection.
+They using the selenium webdriver to control the browser.
+The browser settings and the login data are saved in a custom directory.
+If the headless mode is enabled, the browser windows are loaded invisibly.
+For performance reasons, it is recommended to reuse the browser instances
+and close them yourself at the end:
+
+```python
+import g4f
+from undetected_chromedriver import Chrome, ChromeOptions
+from g4f.Provider import (
+    Bard,
+    Poe,
+    AItianhuSpace,
+    MyShell,
+    Phind,
+    PerplexityAi,
+)
+
+options = ChromeOptions()
+options.add_argument("--incognito");
+browser = Chrome(options=options, headless=True)
+for idx in range(10):
+    response = g4f.ChatCompletion.create(
+        model=g4f.models.default,
+        provider=g4f.Provider.Phind,
+        messages=[{"role": "user", "content": "Suggest me a name."}],
+        browser=browser
+    )
+    print(f"{idx}:", response)
+browser.quit()
+```
+
 ##### Cookies Required
 
 Cookies are essential for the proper functioning of some service providers. It is imperative to maintain an active session, typically achieved by logging into your account.
@@ -253,18 +288,16 @@ When running the g4f package locally, the package automatically retrieves cookie
 import g4f
 
 from g4f.Provider import (
-    Bard,
     Bing,
     HuggingChat,
     OpenAssistant,
-    OpenaiChat,
 )
 
 # Usage
 response = g4f.ChatCompletion.create(
     model=g4f.models.default,
     messages=[{"role": "user", "content": "Hello"}],
-    provider=Bard,
+    provider=Bing,
     #cookies=g4f.get_cookies(".google.com"),
     cookies={"cookie_name": "value", "cookie_name2": "value2"},
     auth=True
diff --git a/g4f/Provider/AItianhuSpace.py b/g4f/Provider/AItianhuSpace.py
index a9a824cf..545944c0 100644
--- a/g4f/Provider/AItianhuSpace.py
+++ b/g4f/Provider/AItianhuSpace.py
@@ -24,7 +24,7 @@ class AItianhuSpace(BaseProvider):
         proxy: str = None,
         timeout: int = 120,
         browser: WebDriver = None,
-        hidden_display: bool = True,
+        headless: bool = True,
         **kwargs
     ) -> CreateResult:
         if not model:
@@ -38,13 +38,7 @@ class AItianhuSpace(BaseProvider):
             print(f"AItianhuSpace | using domain: {domain}")
         url = f"https://{domain}"
         prompt = format_prompt(messages)
-        if browser:
-            driver = browser
-        else:
-            if hidden_display:
-                driver, display = get_browser("", True, proxy)
-            else:
-                driver = get_browser("", False, proxy)
+        driver = browser if browser else get_browser("", headless, proxy)
 
         from selenium.webdriver.common.by import By
         from selenium.webdriver.support.ui import WebDriverWait
@@ -66,6 +60,7 @@ document.getElementById('sheet').addEventListener('click', () => {{
         original_window = driver.current_window_handle
         for window_handle in driver.window_handles:
             if window_handle != original_window:
+                driver.close()
                 driver.switch_to.window(window_handle)
                 break
 
@@ -120,9 +115,7 @@ return "";
                 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/MyShell.py b/g4f/Provider/MyShell.py
index 0efeb0e8..5580ff6e 100644
--- a/g4f/Provider/MyShell.py
+++ b/g4f/Provider/MyShell.py
@@ -21,16 +21,10 @@ class MyShell(BaseProvider):
         proxy: str = None,
         timeout: int = 120,
         browser: WebDriver = None,
-        hidden_display: bool = True,
+        headless: bool = True,
         **kwargs
     ) -> CreateResult:
-        if browser:
-            driver = browser
-        else:
-            if hidden_display:
-                driver, display = get_browser("", True, proxy)
-            else:
-                driver = get_browser("", False, proxy)
+        driver = browser if browser else get_browser("", headless, proxy)
 
         from selenium.webdriver.common.by import By
         from selenium.webdriver.support.ui import WebDriverWait
@@ -87,9 +81,7 @@ return content;
                 elif chunk != "":
                     break
         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/PerplexityAi.py b/g4f/Provider/PerplexityAi.py
index b9de429e..1159840a 100644
--- a/g4f/Provider/PerplexityAi.py
+++ b/g4f/Provider/PerplexityAi.py
@@ -22,16 +22,10 @@ class PerplexityAi(BaseProvider):
         timeout: int = 120,
         browser: WebDriver = None,
         copilot: bool = False,
-        hidden_display: bool = True,
+        headless: bool = True,
         **kwargs
     ) -> CreateResult:
-        if browser:
-            driver = browser
-        else:
-            if hidden_display:
-                driver, display = get_browser("", True, proxy)
-            else:
-                driver = get_browser("", False, proxy)
+        driver = browser if browser else get_browser("", headless, proxy)
 
         from selenium.webdriver.common.by import By
         from selenium.webdriver.support.ui import WebDriverWait
@@ -112,9 +106,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
diff --git a/g4f/Provider/Phind.py b/g4f/Provider/Phind.py
index 2612a92d..0cb17ddb 100644
--- a/g4f/Provider/Phind.py
+++ b/g4f/Provider/Phind.py
@@ -23,16 +23,10 @@ class Phind(BaseProvider):
         timeout: int = 120,
         browser: WebDriver = None,
         creative_mode: bool = None,
-        hidden_display: bool = True,
+        headless: bool = True,
         **kwargs
     ) -> CreateResult:
-        if browser:
-            driver = browser
-        else:
-            if hidden_display:
-                driver, display = get_browser("", True, proxy)
-            else:
-                driver = get_browser("", False, proxy)
+        driver = browser if browser else get_browser("", headless, proxy)
 
         from selenium.webdriver.common.by import By
         from selenium.webdriver.support.ui import WebDriverWait
@@ -111,9 +105,7 @@ if(window.reader) {
                 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/helper.py b/g4f/Provider/helper.py
index b10c912a..f879816a 100644
--- a/g4f/Provider/helper.py
+++ b/g4f/Provider/helper.py
@@ -31,14 +31,6 @@ except ImportError:
     class ChromeOptions():
         def add_argument():
             pass
-try:
-    from pyvirtualdisplay import Display
-except ImportError:
-    class Display():
-        def start():
-            pass
-        def stop():
-            pass
 
 from ..typing import Dict, Messages, Union, Tuple
 from .. import debug
@@ -139,24 +131,16 @@ def format_prompt(messages: Messages, add_special_tokens=False) -> str:
 
 def get_browser(
     user_data_dir: str = None,
-    hidden_display: bool = False,
+    headless: bool = False,
     proxy: str = None,
     options: ChromeOptions = None
-) -> Union[Chrome, Tuple[Chrome, Display]] :
+) -> Chrome:
     if user_data_dir == None:
         user_data_dir = user_config_dir("g4f")
 
-    if hidden_display:
-        display = Display(visible=0, size=(1920, 1080))
-        display.start()
-
     if proxy:
         if not options:
             options = ChromeOptions()
         options.add_argument(f'--proxy-server={proxy}')
 
-    browser = Chrome(user_data_dir=user_data_dir, options=options)
-    if hidden_display:
-        return browser, display
-
-    return browser
\ No newline at end of file
+    return Chrome(user_data_dir=user_data_dir, options=options, headless=headless)
\ No newline at end of file
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