0.3.0.4

3 files changed, 89 insertions, 49 deletions

diff --git a/g4f/Provider/needs_auth/Gemini.py b/g4f/Provider/needs_auth/Gemini.py
index ebf5f413..209c2e91 100644
--- a/g4f/Provider/needs_auth/Gemini.py
+++ b/g4f/Provider/needs_auth/Gemini.py
@@ -16,6 +16,7 @@ try:
 except ImportError:
     pass
 
+from ... import debug
 from ...typing import Messages, Cookies, ImageType, AsyncResult
 from ..base_provider import AsyncGeneratorProvider
 from ..helper import format_prompt, get_cookies
@@ -53,6 +54,56 @@ class Gemini(AsyncGeneratorProvider):
     url = "https://gemini.google.com"
     needs_auth = True
     working = True
+    image_models = ["gemini"]
+    default_vision_model = "gemini"
+    _cookies: Cookies = None
+
+    @classmethod
+    async def nodriver_login(cls) -> Cookies:
+        try:
+            import nodriver as uc
+        except ImportError:
+            return
+        try:
+            from platformdirs import user_config_dir
+            user_data_dir = user_config_dir("g4f-nodriver")
+        except:
+            user_data_dir = None
+        if debug.logging:
+            print(f"Open nodriver with user_dir: {user_data_dir}")
+        browser = await uc.start(user_data_dir=user_data_dir)
+        page = await browser.get(f"{cls.url}/app")
+        await page.select("div.ql-editor.textarea", 240)
+        cookies = {}
+        for c in await page.browser.cookies.get_all():
+            if c.domain.endswith(".google.com"):
+                cookies[c.name] = c.value
+        await page.close()
+        return cookies
+
+    @classmethod
+    async def webdriver_login(cls, proxy: str):
+        driver = None
+        try:
+            driver = get_browser(proxy=proxy)
+            try:
+                driver.get(f"{cls.url}/app")
+                WebDriverWait(driver, 5).until(
+                    EC.visibility_of_element_located((By.CSS_SELECTOR, "div.ql-editor.textarea"))
+                )
+            except:
+                login_url = os.environ.get("G4F_LOGIN_URL")
+                if login_url:
+                    yield f"Please login: [Google Gemini]({login_url})\n\n"
+                WebDriverWait(driver, 240).until(
+                    EC.visibility_of_element_located((By.CSS_SELECTOR, "div.ql-editor.textarea"))
+                )
+            cls._cookies = get_driver_cookies(driver)
+        except MissingRequirementsError:
+            pass
+        finally:
+            if driver:
+                driver.close()
 
     @classmethod
     async def create_async_generator(
@@ -72,47 +123,30 @@ class Gemini(AsyncGeneratorProvider):
             if cookies is None:
                 cookies = {}
             cookies["__Secure-1PSID"] = api_key
-        cookies = cookies if cookies else get_cookies(".google.com", False, True)
+        cls._cookies = cookies or cls._cookies or get_cookies(".google.com", False, True)
         base_connector = get_connector(connector, proxy)
         async with ClientSession(
             headers=REQUEST_HEADERS,
             connector=base_connector
         ) as session:
-            snlm0e  = await cls.fetch_snlm0e(session, cookies) if cookies else None
+            snlm0e  = await cls.fetch_snlm0e(session, cls._cookies) if cls._cookies else None
             if not snlm0e:
-                driver = None
-                try:
-                    driver = get_browser(proxy=proxy)
-                    try:
-                        driver.get(f"{cls.url}/app")
-                        WebDriverWait(driver, 5).until(
-                            EC.visibility_of_element_located((By.CSS_SELECTOR, "div.ql-editor.textarea"))
-                        )
-                    except:
-                        login_url = os.environ.get("G4F_LOGIN_URL")
-                        if login_url:
-                            yield f"Please login: [Google Gemini]({login_url})\n\n"
-                        WebDriverWait(driver, 240).until(
-                            EC.visibility_of_element_located((By.CSS_SELECTOR, "div.ql-editor.textarea"))
-                        )
-                    cookies = get_driver_cookies(driver)
-                except MissingRequirementsError:
-                    pass
-                finally:
-                    if driver:
-                        driver.close()
+                cls._cookies = await cls.nodriver_login();
+                if cls._cookies is None:
+                    async for chunk in cls.webdriver_login(proxy):
+                        yield chunk
 
             if not snlm0e:
-                if "__Secure-1PSID" not in cookies:
+                if "__Secure-1PSID" not in cls._cookies:
                     raise MissingAuthError('Missing "__Secure-1PSID" cookie')
-                snlm0e = await cls.fetch_snlm0e(session, cookies)
+                snlm0e = await cls.fetch_snlm0e(session, cls._cookies)
             if not snlm0e:
-                raise RuntimeError("Invalid auth. SNlM0e not found")
+                raise RuntimeError("Invalid cookies. SNlM0e not found")
 
             image_url = await cls.upload_image(base_connector, to_bytes(image), image_name) if image else None
         
             async with ClientSession(
-                cookies=cookies,
+                cookies=cls._cookies,
                 headers=REQUEST_HEADERS,
                 connector=base_connector,
             ) as client:
diff --git a/g4f/Provider/needs_auth/OpenaiAccount.py b/g4f/Provider/needs_auth/OpenaiAccount.py
index 7be60c86..16bfff66 100644
--- a/g4f/Provider/needs_auth/OpenaiAccount.py
+++ b/g4f/Provider/needs_auth/OpenaiAccount.py
@@ -3,4 +3,6 @@ from __future__ import annotations
 from .OpenaiChat import OpenaiChat
 
 class OpenaiAccount(OpenaiChat):
-    needs_auth = True
\ No newline at end of file
+    needs_auth = True
+    parent = "OpenaiChat"
+    image_models = ["dall-e"]
\ No newline at end of file
diff --git a/g4f/Provider/needs_auth/OpenaiChat.py b/g4f/Provider/needs_auth/OpenaiChat.py
index 36b8bd3c..515230f0 100644
--- a/g4f/Provider/needs_auth/OpenaiChat.py
+++ b/g4f/Provider/needs_auth/OpenaiChat.py
@@ -29,6 +29,7 @@ from ...requests.aiohttp import StreamSession
 from ...image import to_image, to_bytes, ImageResponse, ImageRequest
 from ...errors import MissingAuthError, ResponseError
 from ...providers.conversation import BaseConversation
+from ..helper import format_cookies
 from ..openai.har_file import getArkoseAndAccessToken, NoValidHarFileError
 from ... import debug
 
@@ -43,8 +44,14 @@ class OpenaiChat(AsyncGeneratorProvider, ProviderModelMixin):
     supports_message_history = True
     supports_system_message = True
     default_model = None
+    default_vision_model = "gpt-4-vision"
     models = ["gpt-3.5-turbo", "gpt-4", "gpt-4-gizmo"]
-    model_aliases = {"text-davinci-002-render-sha": "gpt-3.5-turbo", "": "gpt-3.5-turbo", "gpt-4-turbo-preview": "gpt-4"}
+    model_aliases = {
+        "text-davinci-002-render-sha": "gpt-3.5-turbo",
+        "": "gpt-3.5-turbo",
+        "gpt-4-turbo-preview": "gpt-4",
+        "dall-e": "gpt-4",
+    }
     _api_key: str = None
     _headers: dict = None
     _cookies: Cookies = None
@@ -334,9 +341,8 @@ class OpenaiChat(AsyncGeneratorProvider, ProviderModelMixin):
         Raises:
             RuntimeError: If an error occurs during processing.
         """
-
         async with StreamSession(
-            proxies={"all": proxy},
+            proxy=proxy,
             impersonate="chrome",
             timeout=timeout
         ) as session:
@@ -358,26 +364,27 @@ class OpenaiChat(AsyncGeneratorProvider, ProviderModelMixin):
                     api_key = cls._api_key = None
                     cls._create_request_args()
                     if debug.logging:
-                        print("OpenaiChat: Load default_model failed")
+                        print("OpenaiChat: Load default model failed")
                         print(f"{e.__class__.__name__}: {e}")
 
             arkose_token = None
             if cls.default_model is None:
+                error = None
                 try:
-                    arkose_token, api_key, cookies = await getArkoseAndAccessToken(proxy)
-                    cls._create_request_args(cookies)
+                    arkose_token, api_key, cookies, headers = await getArkoseAndAccessToken(proxy)
+                    cls._create_request_args(cookies, headers)
                     cls._set_api_key(api_key)
                 except NoValidHarFileError as e:
-                    ...
+                    error = e
                 if cls._api_key is None:
                     await cls.nodriver_access_token()
                 if cls._api_key is None and cls.needs_auth:
-                    raise e
+                    raise error
                 cls.default_model = cls.get_model(await cls.get_default_model(session, cls._headers))
 
             async with session.post(
                 f"{cls.url}/backend-anon/sentinel/chat-requirements"
-                if not cls._api_key else
+                if cls._api_key is None else
                 f"{cls.url}/backend-api/sentinel/chat-requirements",
                 json={"conversation_mode_kind": "primary_assistant"},
                 headers=cls._headers
@@ -393,8 +400,8 @@ class OpenaiChat(AsyncGeneratorProvider, ProviderModelMixin):
                     print(f'Arkose: {need_arkose} Turnstile: {data["turnstile"]["required"]}')
 
             if need_arkose and arkose_token is None:
-                arkose_token, api_key, cookies = await getArkoseAndAccessToken(proxy)
-                cls._create_request_args(cookies)
+                arkose_token, api_key, cookies, headers = await getArkoseAndAccessToken(proxy)
+                cls._create_request_args(cookies, headers)
                 cls._set_api_key(api_key)
                 if arkose_token is None:
                     raise MissingAuthError("No arkose token found in .har file")
@@ -406,7 +413,8 @@ class OpenaiChat(AsyncGeneratorProvider, ProviderModelMixin):
                     print("OpenaiChat: Upload image failed")
                     print(f"{e.__class__.__name__}: {e}")
 
-            model = cls.get_model(model).replace("gpt-3.5-turbo", "text-davinci-002-render-sha")
+            model = cls.get_model(model)
+            model = "text-davinci-002-render-sha" if model == "gpt-3.5-turbo" else model
             if conversation is None:
                 conversation = Conversation(conversation_id, str(uuid.uuid4()) if parent_id is None else parent_id)
             else:
@@ -613,7 +621,7 @@ this.fetch = async (url, options) => {
                 cookies[c.name] = c.value
         user_agent = await page.evaluate("window.navigator.userAgent")
         await page.close()
-        cls._create_request_args(cookies, user_agent)
+        cls._create_request_args(cookies, user_agent=user_agent)
         cls._set_api_key(api_key)
 
     @classmethod
@@ -667,16 +675,12 @@ this.fetch = async (url, options) => {
             "oai-language": "en-US",
         }
 
-    @staticmethod
-    def _format_cookies(cookies: Cookies):
-        return "; ".join(f"{k}={v}" for k, v in cookies.items() if k != "access_token")
-
     @classmethod
-    def _create_request_args(cls, cookies: Cookies = None, user_agent: str = None):
-        cls._headers = cls.get_default_headers()
+    def _create_request_args(cls, cookies: Cookies = None, headers: dict = None, user_agent: str = None):
+        cls._headers = cls.get_default_headers() if headers is None else headers
         if user_agent is not None:
             cls._headers["user-agent"] = user_agent
-        cls._cookies = {} if cookies is None else cookies
+        cls._cookies = {} if cookies is None else {k: v for k, v in cookies.items() if k != "access_token"}
         cls._update_cookie_header()
 
     @classmethod
@@ -693,7 +697,7 @@ this.fetch = async (url, options) => {
 
     @classmethod
     def _update_cookie_header(cls):
-        cls._headers["cookie"] = cls._format_cookies(cls._cookies)
+        cls._headers["cookie"] = format_cookies(cls._cookies)
 
 class Conversation(BaseConversation):
     """