diff options
Diffstat (limited to 'g4f/Provider/You.py')
| -rw-r--r-- | g4f/Provider/You.py | 19 |
1 files changed, 13 insertions, 6 deletions
diff --git a/g4f/Provider/You.py b/g4f/Provider/You.py index e90365ce..61069503 100644 --- a/g4f/Provider/You.py +++ b/g4f/Provider/You.py @@ -8,7 +8,7 @@ import uuid from ..typing import AsyncResult, Messages, ImageType, Cookies from .base_provider import AsyncGeneratorProvider, ProviderModelMixin from .helper import format_prompt -from ..image import ImageResponse, to_bytes, is_accepted_format +from ..image import ImageResponse, ImagePreview, to_bytes, is_accepted_format from ..requests import StreamSession, FormData, raise_for_status from .you.har_file import get_telemetry_ids from .. import debug @@ -34,6 +34,7 @@ class You(AsyncGeneratorProvider, ProviderModelMixin): model_aliases = { "claude-v2": "claude-2" } + image_models = ["dall-e"] _cookies = None _cookies_used = 0 _telemetry_ids = [] @@ -67,7 +68,7 @@ class You(AsyncGeneratorProvider, ProviderModelMixin): timeout=(30, timeout) ) 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 "" headers = { "Accept": "text/event-stream", @@ -102,11 +103,17 @@ class You(AsyncGeneratorProvider, ProviderModelMixin): if event == "youChatToken" and event in data: yield data[event] elif event == "youChatUpdate" and "t" in data and data["t"] is not None: - match = re.search(r"!\[fig\]\((.+?)\)", data["t"]) - if match: - yield ImageResponse(match.group(1), messages[-1]["content"]) + if chat_mode == "create": + match = re.search(r"!\[(.+?)\]\((.+?)\)", data["t"]) + if match: + if match.group(1) == "fig": + yield ImagePreview(match.group(2), messages[-1]["content"]) + else: + yield ImageResponse(match.group(2), match.group(1)) + else: + yield data["t"] else: - yield data["t"] + yield data["t"] @classmethod async def upload_file(cls, client: StreamSession, cookies: Cookies, file: bytes, filename: str = None) -> dict: |