Add upload image support in You

g4f/Provider/You.py

@@ -3,67 +3,114 @@ from __future__ import annotations
 import json
 import base64
 import uuid
+from aiohttp import ClientSession, FormData
 
-from ..requests import StreamSession
-from ..typing import AsyncGenerator, Messages
-from .base_provider import AsyncGeneratorProvider, format_prompt
-
+from ..typing import AsyncGenerator, Messages, ImageType, Cookies
+from .base_provider import AsyncGeneratorProvider
+from .helper import get_connector, format_prompt
+from ..image import to_bytes
+from ..defaults import DEFAULT_HEADERS
 
 class You(AsyncGeneratorProvider):
     url = "https://you.com"
     working = True
     supports_gpt_35_turbo = True
     supports_gpt_4 = True
-    _session_used = 0
-    _session_token = None
+    _cookies = None
+    _cookies_used = 0
 
     @classmethod
     async def create_async_generator(
         cls,
         model: str,
         messages: Messages,
+        image: ImageType = None,
+        image_name: str = None,
         proxy: str = None,
-        timeout: int = 120,
+        chat_mode: str = "default",
         **kwargs,
     ) -> AsyncGenerator:
-        async with StreamSession(proxies={"https": proxy}, impersonate="chrome107", timeout=timeout) as session:
+        async with ClientSession(
+            connector=get_connector(kwargs.get("connector"), proxy),
+            headers=DEFAULT_HEADERS
+        ) as client:
+            if image:
+                chat_mode = "agent"
+            elif model == "gpt-4":
+                chat_mode = 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 ""
+            #questions = [message["content"] for message in messages if message["role"] == "user"]
+            # chat = [
+            #     {"question": questions[idx-1], "answer": message["content"]}
+            #     for idx, message in enumerate(messages)
+            #     if message["role"] == "assistant"
+            #     and idx < len(questions)
+            # ]
             headers = {
                 "Accept": "text/event-stream",
                 "Referer": f"{cls.url}/search?fromSearchBar=true&tbm=youchat",
             }
             data = {
+                "userFiles": upload,
+                "selectedChatMode": chat_mode,
                 "q": format_prompt(messages),
                 "domain": "youchat",
-                "chat": "", "selectedChatMode": "gpt-4" if model == "gpt-4" else "default"
+                #"chat": json.dumps(chat),
             }
-            async with session.get(
+            async with client.post(
                 f"{cls.url}/api/streamingSearch",
-                params=data,
+                data=data,
                 headers=headers,
-                cookies=cls.get_cookies(await cls.get_session_token(proxy, timeout)) if model == "gpt-4" else None
+                cookies=cookies
             ) as response:
                 response.raise_for_status()
-                start = b'data: {"youChatToken": '
-                async for line in response.iter_lines():
-                    if line.startswith(start):
-                        yield json.loads(line[len(start):-1])
+                async for line in response.content:
+                    if line.startswith(b'event: '):
+                        event = line[7:-1]
+                    elif line.startswith(b'data: '):
+                        if event == b"youChatUpdate" or event == b"youChatToken":
+                            data = json.loads(line[6:-1])
+                        if event == b"youChatToken" and "youChatToken" in data:
+                            yield data["youChatToken"]
+                        elif event == b"youChatUpdate" and "t" in data:
+                            yield data["t"]                         
 
     @classmethod
-    async def get_session_token(cls, proxy: str, timeout: int):
-        if not cls._session_token or cls._session_used >= 5:
-            cls._session_token = await cls.create_session_token(proxy, timeout)
-        cls._session_used += 1
-        return cls._session_token
-
-    def get_cookies(access_token: str, session_jwt: str = "0"):
-        return {
-            'stytch_session_jwt': session_jwt,
-            'ydc_stytch_session': access_token,
-            'ydc_stytch_session_jwt': session_jwt
-        }
+    async def upload_file(cls, client: ClientSession, cookies: Cookies, file: bytes, filename: str = None) -> dict:
+        async with client.get(
+            f"{cls.url}/api/get_nonce",
+            cookies=cookies,
+        ) as response:
+            response.raise_for_status()
+            upload_nonce = await response.text()
+        data = FormData()
+        data.add_field('file', file, filename=filename)
+        async with client.post(
+            f"{cls.url}/api/upload",
+            data=data,
+            headers={
+                "X-Upload-Nonce": upload_nonce,
+            },
+            cookies=cookies
+        ) as response:
+            if not response.ok:
+                raise RuntimeError(f"Response: {await response.text()}")
+            result = await response.json()
+        result["user_filename"] = filename
+        result["size"] = len(file)
+        return result
 
     @classmethod
-    def get_jwt(cls):
+    async def get_cookies(cls, client: ClientSession) -> Cookies:
+        if not cls._cookies or cls._cookies_used >= 5:
+            cls._cookies = await cls.create_cookies(client)
+            cls._cookies_used = 0
+        cls._cookies_used += 1
+        return cls._cookies        
+
+    @classmethod
+    def get_sdk(cls) -> str:
         return base64.standard_b64encode(json.dumps({
             "event_id":f"event-id-{str(uuid.uuid4())}",
             "app_session_id":f"app-session-id-{str(uuid.uuid4())}",
@@ -75,19 +122,21 @@ class You(AsyncGeneratorProvider):
             "sdk":{"identifier":"Stytch.js Javascript SDK","version":"3.3.0"
         }}).encode()).decode()
 
-    @classmethod
-    async def create_session_token(cls, proxy: str, timeout: int):
-        async with StreamSession(proxies={"https": proxy}, impersonate="chrome110", timeout=timeout) as session:
-            user_uuid = str(uuid.uuid4())
+    def get_auth() -> str:
         auth_uuid = "507a52ad-7e69-496b-aee0-1c9863c7c8"
         auth_token = f"public-token-live-{auth_uuid}bb:public-token-live-{auth_uuid}19"
         auth = base64.standard_b64encode(auth_token.encode()).decode()
-            async with session.post(
+        return f"Basic {auth}",
+
+    @classmethod
+    async def create_cookies(cls, client: ClientSession) -> Cookies:
+        user_uuid = str(uuid.uuid4())
+        async with client.post(
             "https://web.stytch.com/sdk/v1/passwords",
             headers={
-                    "Authorization": f"Basic {auth}",
-                    "X-SDK-Client": cls.get_jwt(),
-                    "X-SDK-Parent-Host": "https://you.com"
+                "Authorization": cls.get_auth(),
+                "X-SDK-Client": cls.get_sdk(),
+                "X-SDK-Parent-Host": cls.url
             },
             json={
                 "email": f"{user_uuid}@gmail.com",
@@ -97,4 +146,10 @@ class You(AsyncGeneratorProvider):
         ) as response:
             if not response.ok:
                 raise RuntimeError(f"Response: {await response.text()}")
-                return (await response.json())["data"]["session_token"]
+            session = (await response.json())["data"]
+        return {
+            "stytch_session": session["session_token"],
+            'stytch_session_jwt': session["session_jwt"],
+            'ydc_stytch_session': session["session_token"],
+            'ydc_stytch_session_jwt': session["session_jwt"],
+        }