From a26421bcd8b8580650fa05b3fb4f8fdfa0ef9921 Mon Sep 17 00:00:00 2001 From: Heiner Lohaus Date: Sun, 21 Apr 2024 15:15:55 +0200 Subject: Add image model list --- g4f/Provider/needs_auth/Gemini.py | 1 + g4f/Provider/needs_auth/OpenaiAccount.py | 3 ++- g4f/Provider/needs_auth/OpenaiChat.py | 30 ++++++++++++++++-------------- 3 files changed, 19 insertions(+), 15 deletions(-) (limited to 'g4f/Provider/needs_auth') diff --git a/g4f/Provider/needs_auth/Gemini.py b/g4f/Provider/needs_auth/Gemini.py index ebf5f413..3917df80 100644 --- a/g4f/Provider/needs_auth/Gemini.py +++ b/g4f/Provider/needs_auth/Gemini.py @@ -53,6 +53,7 @@ class Gemini(AsyncGeneratorProvider): url = "https://gemini.google.com" needs_auth = True working = True + image_models = ["gemini"] @classmethod async def create_async_generator( diff --git a/g4f/Provider/needs_auth/OpenaiAccount.py b/g4f/Provider/needs_auth/OpenaiAccount.py index 7be60c86..6260d343 100644 --- a/g4f/Provider/needs_auth/OpenaiAccount.py +++ b/g4f/Provider/needs_auth/OpenaiAccount.py @@ -3,4 +3,5 @@ from __future__ import annotations from .OpenaiChat import OpenaiChat class OpenaiAccount(OpenaiChat): - needs_auth = True \ No newline at end of file + needs_auth = True + 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..7952d606 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 @@ -44,7 +45,12 @@ class OpenaiChat(AsyncGeneratorProvider, ProviderModelMixin): supports_system_message = True default_model = None 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 @@ -364,8 +370,8 @@ class OpenaiChat(AsyncGeneratorProvider, ProviderModelMixin): arkose_token = None if cls.default_model is 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: ... @@ -393,8 +399,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") @@ -613,7 +619,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 +673,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 +695,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): """ -- cgit v1.2.3