1 files changed, 35 insertions, 25 deletions

diff --git a/g4f/Provider/You.py b/g4f/Provider/You.py
index 1fdaf06d..85b60452 100644
--- a/g4f/Provider/You.py
+++ b/g4f/Provider/You.py
@@ -4,13 +4,14 @@ import re
 import json
 import base64
 import uuid
-from aiohttp import ClientSession, FormData, BaseConnector
+from asyncio import get_running_loop
+from aiohttp import ClientSession, FormData, BaseConnector, CookieJar
 
 from ..typing import AsyncResult, Messages, ImageType, Cookies
 from .base_provider import AsyncGeneratorProvider, ProviderModelMixin
-from ..providers.helper import get_connector, format_prompt
+from .helper import format_prompt, get_connector
 from ..image import to_bytes, ImageResponse
-from ..requests.defaults import DEFAULT_HEADERS
+from ..requests import WebDriver, raise_for_status, get_args_from_browser
 
 class You(AsyncGeneratorProvider, ProviderModelMixin):
     url = "https://you.com"
@@ -32,6 +33,8 @@ class You(AsyncGeneratorProvider, ProviderModelMixin):
     model_aliases = {
         "claude-v2": "claude-2"
     }
+    _args: dict = None
+    _cookie_jar: CookieJar = None
     _cookies = None
     _cookies_used = 0
 
@@ -43,25 +46,34 @@ class You(AsyncGeneratorProvider, ProviderModelMixin):
         image: ImageType = None,
         image_name: str = None,
         connector: BaseConnector = None,
+        webdriver: WebDriver = None,
         proxy: str = None,
         chat_mode: str = "default",
         **kwargs,
     ) -> AsyncResult:
+        if cls._args is None:
+            cls._args = get_args_from_browser(cls.url, webdriver, proxy)
+            cls._cookie_jar = CookieJar(loop=get_running_loop())
+        else:
+            if "cookies" in cls._args:
+                del cls._args["cookies"]
+            cls._cookie_jar._loop = get_running_loop()
+        if image is not None:
+            chat_mode = "agent"
+        elif not model or model == cls.default_model:
+            chat_mode = "default"
+        elif model.startswith("dall-e"):
+            chat_mode = "create"
+        else:
+            chat_mode = "custom"
+            model = cls.get_model(model)
         async with ClientSession(
             connector=get_connector(connector, proxy),
-            headers=DEFAULT_HEADERS
-        ) as client:
-            if image is not None:
-                chat_mode = "agent"
-            elif not model or model == cls.default_model:
-                chat_mode = "default"
-            elif model.startswith("dall-e"):
-                chat_mode = "create"
-            else:
-                chat_mode = "custom"
-                model = cls.get_model(model)
-            cookies = await cls.get_cookies(client) if chat_mode != "default" else None
-            upload = json.dumps([await cls.upload_file(client, cookies, to_bytes(image), image_name)]) if image else ""
+            cookie_jar=cls._cookie_jar,
+            **cls._args
+        ) as session:
+            cookies = await cls.get_cookies(session) if chat_mode != "default" else None
+            upload = json.dumps([await cls.upload_file(session, cookies, to_bytes(image), image_name)]) if image else ""
             #questions = [message["content"] for message in messages if message["role"] == "user"]
             # chat = [
             #     {"question": questions[idx-1], "answer": message["content"]}
@@ -70,8 +82,8 @@ class You(AsyncGeneratorProvider, ProviderModelMixin):
             #     and idx < len(questions)
             # ]
             headers = {
-                "Accept": "text/event-stream",
-                "Referer": f"{cls.url}/search?fromSearchBar=true&tbm=youchat",
+                "accept": "text/event-stream",
+                "referer": f"{cls.url}/search?fromSearchBar=true&tbm=youchat",
             }
             data = {
                 "userFiles": upload,
@@ -86,14 +98,14 @@ class You(AsyncGeneratorProvider, ProviderModelMixin):
             }
             if chat_mode == "custom":
                 params["selectedAIModel"] = model.replace("-", "_")
-            async with (client.post if chat_mode == "default" else client.get)(
+            async with (session.post if chat_mode == "default" else session.get)(
                 f"{cls.url}/api/streamingSearch",
                 data=data,
                 params=params,
                 headers=headers,
                 cookies=cookies
             ) as response:
-                response.raise_for_status()
+                await raise_for_status(response)
                 async for line in response.content:
                     if line.startswith(b'event: '):
                         event = line[7:-1].decode()
@@ -115,7 +127,7 @@ class You(AsyncGeneratorProvider, ProviderModelMixin):
             f"{cls.url}/api/get_nonce",
             cookies=cookies,
         ) as response:
-            response.raise_for_status()
+            await raise_for_status(response)
             upload_nonce = await response.text()
         data = FormData()
         data.add_field('file', file, filename=filename)
@@ -127,8 +139,7 @@ class You(AsyncGeneratorProvider, ProviderModelMixin):
             },
             cookies=cookies
         ) as response:
-            if not response.ok:
-                raise RuntimeError(f"Response: {await response.text()}")
+            await raise_for_status(response)
             result = await response.json()
         result["user_filename"] = filename
         result["size"] = len(file)
@@ -177,8 +188,7 @@ class You(AsyncGeneratorProvider, ProviderModelMixin):
                 "session_duration_minutes": 129600
             }
         ) as response:
-            if not response.ok:
-                raise RuntimeError(f"Response: {await response.text()}")
+            await raise_for_status(response)
             session = (await response.json())["data"]
         return {
             "stytch_session": session["session_token"],