Fix Gemini Proxy Connect call failed (#1768)

Gemini. py file adds get_connector to support proxy

---------

Co-authored-by: HyiKi <dango.little@gmail.com>
This commit is contained in:
HyiKi 2024-04-05 22:16:12 +08:00 committed by GitHub
parent b7ccb9500d
commit b401b6df1c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 98 additions and 93 deletions

View File

@ -5,7 +5,9 @@ import json
import random import random
import re import re
from aiohttp import ClientSession from aiohttp import ClientSession, BaseConnector
from ..helper import get_connector
try: try:
from selenium.webdriver.common.by import By from selenium.webdriver.common.by import By
@ -58,13 +60,19 @@ class Gemini(AsyncGeneratorProvider):
messages: Messages, messages: Messages,
proxy: str = None, proxy: str = None,
cookies: Cookies = None, cookies: Cookies = None,
connector: BaseConnector = None,
image: ImageType = None, image: ImageType = None,
image_name: str = None, image_name: str = None,
**kwargs **kwargs
) -> AsyncResult: ) -> AsyncResult:
prompt = format_prompt(messages) prompt = format_prompt(messages)
cookies = cookies if cookies else get_cookies(".google.com", False, True) cookies = cookies if cookies else get_cookies(".google.com", False, True)
snlm0e = await cls.fetch_snlm0e(cookies, proxy) if cookies else None base_connector = get_connector(connector, proxy)
async with ClientSession(
headers=REQUEST_HEADERS,
connector=base_connector
) as session:
snlm0e = await cls.fetch_snlm0e(session, cookies) if cookies else None
if not snlm0e: if not snlm0e:
driver = None driver = None
try: try:
@ -91,16 +99,17 @@ class Gemini(AsyncGeneratorProvider):
if not snlm0e: if not snlm0e:
if "__Secure-1PSID" not in cookies: if "__Secure-1PSID" not in cookies:
raise MissingAuthError('Missing "__Secure-1PSID" cookie') raise MissingAuthError('Missing "__Secure-1PSID" cookie')
snlm0e = await cls.fetch_snlm0e(cookies, proxy) snlm0e = await cls.fetch_snlm0e(session, cookies)
if not snlm0e: if not snlm0e:
raise RuntimeError("Invalid auth. SNlM0e not found") raise RuntimeError("Invalid auth. SNlM0e not found")
image_url = await cls.upload_image(to_bytes(image), image_name, proxy) if image else None image_url = await cls.upload_image(base_connector, to_bytes(image), image_name) if image else None
async with ClientSession( async with ClientSession(
cookies=cookies, cookies=cookies,
headers=REQUEST_HEADERS headers=REQUEST_HEADERS,
) as session: connector=base_connector,
) as client:
params = { params = {
'bl': REQUEST_BL_PARAM, 'bl': REQUEST_BL_PARAM,
'_reqid': random.randint(1111, 9999), '_reqid': random.randint(1111, 9999),
@ -114,11 +123,10 @@ class Gemini(AsyncGeneratorProvider):
image_name=image_name image_name=image_name
))]) ))])
} }
async with session.post( async with client.post(
REQUEST_URL, REQUEST_URL,
data=data, data=data,
params=params, params=params,
proxy=proxy
) as response: ) as response:
response = await response.text() response = await response.text()
response_part = json.loads(json.loads(response.splitlines()[-5])[0][2]) response_part = json.loads(json.loads(response.splitlines()[-5])[0][2])
@ -138,9 +146,9 @@ class Gemini(AsyncGeneratorProvider):
resolved_images = [] resolved_images = []
preview = [] preview = []
for image in images: for image in images:
async with session.get(image, allow_redirects=False) as fetch: async with client.get(image, allow_redirects=False) as fetch:
image = fetch.headers["location"] image = fetch.headers["location"]
async with session.get(image, allow_redirects=False) as fetch: async with client.get(image, allow_redirects=False) as fetch:
image = fetch.headers["location"] image = fetch.headers["location"]
resolved_images.append(image) resolved_images.append(image)
preview.append(image.replace('=s512', '=s200')) preview.append(image.replace('=s512', '=s200'))
@ -171,11 +179,12 @@ class Gemini(AsyncGeneratorProvider):
0, 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( async with ClientSession(
headers=UPLOAD_IMAGE_HEADERS headers=UPLOAD_IMAGE_HEADERS,
connector=connector
) as session: ) 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() reponse.raise_for_status()
headers = { headers = {
@ -184,7 +193,7 @@ class Gemini(AsyncGeneratorProvider):
} }
data = f"File name: {image_name}" if image_name else None data = f"File name: {image_name}" if image_name else None
async with session.post( async with session.post(
UPLOAD_IMAGE_URL, headers=headers, data=data, proxy=proxy UPLOAD_IMAGE_URL, headers=headers, data=data
) as response: ) as response:
response.raise_for_status() response.raise_for_status()
upload_url = response.headers["X-Goog-Upload-Url"] upload_url = response.headers["X-Goog-Upload-Url"]
@ -195,18 +204,14 @@ class Gemini(AsyncGeneratorProvider):
headers["x-goog-upload-command"] = "upload, finalize" headers["x-goog-upload-command"] = "upload, finalize"
headers["X-Goog-Upload-Offset"] = "0" headers["X-Goog-Upload-Offset"] = "0"
async with session.post( async with session.post(
upload_url, headers=headers, data=image, proxy=proxy upload_url, headers=headers, data=image
) as response: ) as response:
response.raise_for_status() response.raise_for_status()
return await response.text() return await response.text()
@classmethod @classmethod
async def fetch_snlm0e(cls, cookies: Cookies, proxy: str = None): async def fetch_snlm0e(cls, session: ClientSession, cookies: Cookies):
async with ClientSession( async with session.get(cls.url, cookies=cookies) as response:
cookies=cookies,
headers=REQUEST_HEADERS
) as session:
async with session.get(cls.url, proxy=proxy) as response:
text = await response.text() text = await response.text()
match = re.search(r'SNlM0e\":\"(.*?)\"', text) match = re.search(r'SNlM0e\":\"(.*?)\"', text)
if match: if match: