From b401b6df1cf020dfe018dee2734b8b081df6cd24 Mon Sep 17 00:00:00 2001 From: HyiKi <55942998+HyiKi@users.noreply.github.com> Date: Fri, 5 Apr 2024 22:16:12 +0800 Subject: [PATCH] Fix Gemini Proxy Connect call failed (#1768) Gemini. py file adds get_connector to support proxy --------- Co-authored-by: HyiKi --- g4f/Provider/needs_auth/Gemini.py | 191 +++++++++++++++--------------- 1 file changed, 98 insertions(+), 93 deletions(-) diff --git a/g4f/Provider/needs_auth/Gemini.py b/g4f/Provider/needs_auth/Gemini.py index bd30d8b2..0621cacd 100644 --- a/g4f/Provider/needs_auth/Gemini.py +++ b/g4f/Provider/needs_auth/Gemini.py @@ -5,7 +5,9 @@ import json import random import re -from aiohttp import ClientSession +from aiohttp import ClientSession, BaseConnector + +from ..helper import get_connector try: from selenium.webdriver.common.by import By @@ -58,93 +60,99 @@ class Gemini(AsyncGeneratorProvider): messages: Messages, proxy: str = None, cookies: Cookies = None, + connector: BaseConnector = None, image: ImageType = None, image_name: str = None, **kwargs ) -> AsyncResult: prompt = format_prompt(messages) cookies = cookies if cookies else get_cookies(".google.com", False, True) - snlm0e = await cls.fetch_snlm0e(cookies, proxy) if 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() - - if not snlm0e: - if "__Secure-1PSID" not in cookies: - raise MissingAuthError('Missing "__Secure-1PSID" cookie') - snlm0e = await cls.fetch_snlm0e(cookies, proxy) - if not snlm0e: - raise RuntimeError("Invalid auth. SNlM0e not found") - - image_url = await cls.upload_image(to_bytes(image), image_name, proxy) if image else None - + base_connector = get_connector(connector, proxy) async with ClientSession( - cookies=cookies, - headers=REQUEST_HEADERS + headers=REQUEST_HEADERS, + connector=base_connector ) as session: - params = { - 'bl': REQUEST_BL_PARAM, - '_reqid': random.randint(1111, 9999), - 'rt': 'c' - } - data = { - 'at': snlm0e, - 'f.req': json.dumps([None, json.dumps(cls.build_request( - prompt, - image_url=image_url, - image_name=image_name - ))]) - } - async with session.post( - REQUEST_URL, - data=data, - params=params, - proxy=proxy - ) as response: - response = await response.text() - response_part = json.loads(json.loads(response.splitlines()[-5])[0][2]) - if response_part[4] is None: - response_part = json.loads(json.loads(response.splitlines()[-7])[0][2]) + snlm0e = await cls.fetch_snlm0e(session, cookies) if 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() - content = response_part[4][0][1][0] - image_prompt = None - match = re.search(r'\[Imagen of (.*?)\]', content) - if match: - image_prompt = match.group(1) - content = content.replace(match.group(0), '') + if not snlm0e: + if "__Secure-1PSID" not in cookies: + raise MissingAuthError('Missing "__Secure-1PSID" cookie') + snlm0e = await cls.fetch_snlm0e(session, cookies) + if not snlm0e: + raise RuntimeError("Invalid auth. SNlM0e not found") - yield content - if image_prompt: - images = [image[0][3][3] for image in response_part[4][0][12][7][0]] - resolved_images = [] - preview = [] - for image in images: - async with session.get(image, allow_redirects=False) as fetch: - image = fetch.headers["location"] - async with session.get(image, allow_redirects=False) as fetch: - image = fetch.headers["location"] - resolved_images.append(image) - preview.append(image.replace('=s512', '=s200')) - yield ImageResponse(resolved_images, image_prompt, {"orginal_links": images, "preview": preview}) + image_url = await cls.upload_image(base_connector, to_bytes(image), image_name) if image else None + + async with ClientSession( + cookies=cookies, + headers=REQUEST_HEADERS, + connector=base_connector, + ) as client: + params = { + 'bl': REQUEST_BL_PARAM, + '_reqid': random.randint(1111, 9999), + 'rt': 'c' + } + data = { + 'at': snlm0e, + 'f.req': json.dumps([None, json.dumps(cls.build_request( + prompt, + image_url=image_url, + image_name=image_name + ))]) + } + async with client.post( + REQUEST_URL, + data=data, + params=params, + ) as response: + response = await response.text() + response_part = json.loads(json.loads(response.splitlines()[-5])[0][2]) + if response_part[4] is None: + response_part = json.loads(json.loads(response.splitlines()[-7])[0][2]) + + content = response_part[4][0][1][0] + image_prompt = None + match = re.search(r'\[Imagen of (.*?)\]', content) + if match: + image_prompt = match.group(1) + content = content.replace(match.group(0), '') + + yield content + if image_prompt: + images = [image[0][3][3] for image in response_part[4][0][12][7][0]] + resolved_images = [] + preview = [] + for image in images: + async with client.get(image, allow_redirects=False) as fetch: + image = fetch.headers["location"] + async with client.get(image, allow_redirects=False) as fetch: + image = fetch.headers["location"] + resolved_images.append(image) + preview.append(image.replace('=s512', '=s200')) + yield ImageResponse(resolved_images, image_prompt, {"orginal_links": images, "preview": preview}) def build_request( prompt: str, @@ -171,11 +179,12 @@ class Gemini(AsyncGeneratorProvider): 0, ] - async def upload_image(image: bytes, image_name: str = None, proxy: str = None): + async def upload_image(connector: BaseConnector, image: bytes, image_name: str = None): async with ClientSession( - headers=UPLOAD_IMAGE_HEADERS + headers=UPLOAD_IMAGE_HEADERS, + connector=connector ) as session: - async with session.options(UPLOAD_IMAGE_URL, proxy=proxy) as reponse: + async with session.options(UPLOAD_IMAGE_URL) as reponse: reponse.raise_for_status() headers = { @@ -184,7 +193,7 @@ class Gemini(AsyncGeneratorProvider): } data = f"File name: {image_name}" if image_name else None async with session.post( - UPLOAD_IMAGE_URL, headers=headers, data=data, proxy=proxy + UPLOAD_IMAGE_URL, headers=headers, data=data ) as response: response.raise_for_status() upload_url = response.headers["X-Goog-Upload-Url"] @@ -195,19 +204,15 @@ class Gemini(AsyncGeneratorProvider): headers["x-goog-upload-command"] = "upload, finalize" headers["X-Goog-Upload-Offset"] = "0" async with session.post( - upload_url, headers=headers, data=image, proxy=proxy + upload_url, headers=headers, data=image ) as response: response.raise_for_status() return await response.text() @classmethod - async def fetch_snlm0e(cls, cookies: Cookies, proxy: str = None): - async with ClientSession( - cookies=cookies, - headers=REQUEST_HEADERS - ) as session: - async with session.get(cls.url, proxy=proxy) as response: - text = await response.text() - match = re.search(r'SNlM0e\":\"(.*?)\"', text) - if match: - return match.group(1) \ No newline at end of file + async def fetch_snlm0e(cls, session: ClientSession, cookies: Cookies): + async with session.get(cls.url, cookies=cookies) as response: + text = await response.text() + match = re.search(r'SNlM0e\":\"(.*?)\"', text) + if match: + return match.group(1) \ No newline at end of file