From 69ef224f92fcec3335a857305191678b886651e1 Mon Sep 17 00:00:00 2001 From: Heiner Lohaus Date: Wed, 24 Jan 2024 00:46:35 +0100 Subject: Add get_connector helper --- g4f/Provider/Bing.py | 10 ++-------- g4f/Provider/HuggingChat.py | 10 ++-------- g4f/Provider/Liaobots.py | 9 ++------- g4f/Provider/PerplexityLabs.py | 9 ++------- g4f/Provider/bing/create_images.py | 10 ++-------- g4f/Provider/helper.py | 15 +++++++++++++-- g4f/errors.py | 3 +++ g4f/typing.py | 4 +++- 8 files changed, 29 insertions(+), 41 deletions(-) diff --git a/g4f/Provider/Bing.py b/g4f/Provider/Bing.py index 751b0d3a..11bb1414 100644 --- a/g4f/Provider/Bing.py +++ b/g4f/Provider/Bing.py @@ -11,6 +11,7 @@ from aiohttp import ClientSession, ClientTimeout, BaseConnector from ..typing import AsyncResult, Messages, ImageType from ..image import ImageResponse from .base_provider import AsyncGeneratorProvider +from .helper import get_connector from .bing.upload_image import upload_image from .bing.create_images import create_images from .bing.conversation import Conversation, create_conversation, delete_conversation @@ -68,15 +69,8 @@ class Bing(AsyncGeneratorProvider): cookies = {**Defaults.cookies, **cookies} if cookies else Defaults.cookies gpt4_turbo = True if model.startswith("gpt-4-turbo") else False - - if proxy and not connector: - try: - from aiohttp_socks import ProxyConnector - connector = ProxyConnector.from_url(proxy) - except ImportError: - raise RuntimeError('Install "aiohttp_socks" package for proxy support') - return stream_generate(prompt, tone, image, context, cookies, connector, web_search, gpt4_turbo, timeout) + return stream_generate(prompt, tone, image, context, cookies, get_connector(connector, proxy), web_search, gpt4_turbo, timeout) def create_context(messages: Messages) -> str: """ diff --git a/g4f/Provider/HuggingChat.py b/g4f/Provider/HuggingChat.py index a481ce7c..79e4ae38 100644 --- a/g4f/Provider/HuggingChat.py +++ b/g4f/Provider/HuggingChat.py @@ -6,7 +6,7 @@ from aiohttp import ClientSession, BaseConnector from ..typing import AsyncResult, Messages from .base_provider import AsyncGeneratorProvider, ProviderModelMixin -from .helper import format_prompt, get_cookies +from .helper import format_prompt, get_cookies, get_connector class HuggingChat(AsyncGeneratorProvider, ProviderModelMixin): @@ -44,16 +44,10 @@ class HuggingChat(AsyncGeneratorProvider, ProviderModelMixin): headers = { 'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/111.0.0.0 Safari/537.36', } - if proxy and not connector: - try: - from aiohttp_socks import ProxyConnector - connector = ProxyConnector.from_url(proxy) - except ImportError: - raise RuntimeError('Install "aiohttp_socks" package for proxy support') async with ClientSession( cookies=cookies, headers=headers, - connector=connector + connector=get_connector(connector, proxy) ) as session: async with session.post(f"{cls.url}/conversation", json={"model": cls.get_model(model)}, proxy=proxy) as response: conversation_id = (await response.json())["conversationId"] diff --git a/g4f/Provider/Liaobots.py b/g4f/Provider/Liaobots.py index b16b44d3..e93642ba 100644 --- a/g4f/Provider/Liaobots.py +++ b/g4f/Provider/Liaobots.py @@ -6,6 +6,7 @@ from aiohttp import ClientSession, BaseConnector from ..typing import AsyncResult, Messages from .base_provider import AsyncGeneratorProvider, ProviderModelMixin +from .helper import get_connector models = { "gpt-4": { @@ -101,16 +102,10 @@ class Liaobots(AsyncGeneratorProvider, ProviderModelMixin): "referer": f"{cls.url}/", "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/112.0.0.0 Safari/537.36", } - if proxy and not connector: - try: - from aiohttp_socks import ProxyConnector - connector = ProxyConnector.from_url(proxy) - except ImportError: - raise RuntimeError('Install "aiohttp_socks" package for proxy support') async with ClientSession( headers=headers, cookie_jar=cls._cookie_jar, - connector=connector + connector=get_connector(connector, proxy) ) as session: cls._auth_code = auth if isinstance(auth, str) else cls._auth_code if not cls._auth_code: diff --git a/g4f/Provider/PerplexityLabs.py b/g4f/Provider/PerplexityLabs.py index 9b11b708..5002b39f 100644 --- a/g4f/Provider/PerplexityLabs.py +++ b/g4f/Provider/PerplexityLabs.py @@ -6,6 +6,7 @@ from aiohttp import ClientSession, BaseConnector from ..typing import AsyncResult, Messages from .base_provider import AsyncGeneratorProvider, ProviderModelMixin +from .helper import get_connector API_URL = "https://labs-api.perplexity.ai/socket.io/" WS_URL = "wss://labs-api.perplexity.ai/socket.io/" @@ -48,13 +49,7 @@ class PerplexityLabs(AsyncGeneratorProvider, ProviderModelMixin): "Sec-Fetch-Site": "same-site", "TE": "trailers", } - if proxy and not connector: - try: - from aiohttp_socks import ProxyConnector - connector = ProxyConnector.from_url(proxy) - except ImportError: - raise RuntimeError('Install "aiohttp_socks" package for proxy support') - async with ClientSession(headers=headers, connector=connector) as session: + async with ClientSession(headers=headers, connector=get_connector(connector, proxy)) as session: t = format(random.getrandbits(32), '08x') async with session.get( f"{API_URL}?EIO=4&transport=polling&t={t}", diff --git a/g4f/Provider/bing/create_images.py b/g4f/Provider/bing/create_images.py index 30fb372d..af39ef1e 100644 --- a/g4f/Provider/bing/create_images.py +++ b/g4f/Provider/bing/create_images.py @@ -13,7 +13,7 @@ from urllib.parse import quote from typing import Generator, List, Dict from ..create_images import CreateImagesProvider -from ..helper import get_cookies +from ..helper import get_cookies, get_connector from ...webdriver import WebDriver, get_driver_cookies, get_browser from ...base_provider import ProviderType from ...image import ImageResponse @@ -79,13 +79,7 @@ def create_session(cookies: Dict[str, str], proxy: str = None, connector: BaseCo } if cookies: headers["Cookie"] = "; ".join(f"{k}={v}" for k, v in cookies.items()) - if proxy and not connector: - try: - from aiohttp_socks import ProxyConnector - connector = ProxyConnector.from_url(proxy) - except ImportError: - raise RuntimeError('Install "aiohttp_socks" package for proxy support') - return ClientSession(headers=headers, connector=connector) + return ClientSession(headers=headers, connector=get_connector(connector, proxy)) async def create_images(session: ClientSession, prompt: str, proxy: str = None, timeout: int = TIMEOUT_IMAGE_CREATION) -> List[str]: """ diff --git a/g4f/Provider/helper.py b/g4f/Provider/helper.py index fce1ee6f..cf204e39 100644 --- a/g4f/Provider/helper.py +++ b/g4f/Provider/helper.py @@ -6,13 +6,15 @@ import random import secrets import string from asyncio import AbstractEventLoop, BaseEventLoop +from aiohttp import BaseConnector from platformdirs import user_config_dir from browser_cookie3 import ( chrome, chromium, opera, opera_gx, brave, edge, vivaldi, firefox, _LinuxPasswordManager, BrowserCookieError ) -from ..typing import Dict, Messages +from ..typing import Dict, Messages, Optional +from ..errors import AiohttpSocksError from .. import debug # Global variable to store cookies @@ -147,4 +149,13 @@ def get_random_hex() -> str: Returns: str: A random hexadecimal string of 32 characters (16 bytes). """ - return secrets.token_hex(16).zfill(32) \ No newline at end of file + return secrets.token_hex(16).zfill(32) + +def get_connector(connector: BaseConnector = None, proxy: str = None) -> Optional[BaseConnector]: + if proxy and not connector: + try: + from aiohttp_socks import ProxyConnector + connector = ProxyConnector.from_url(proxy) + except ImportError: + raise AiohttpSocksError('Install "aiohttp_socks" package for proxy support') + return connector \ No newline at end of file diff --git a/g4f/errors.py b/g4f/errors.py index ddfe74db..aaaa6b4b 100644 --- a/g4f/errors.py +++ b/g4f/errors.py @@ -29,4 +29,7 @@ class NestAsyncioError(Exception): pass class ModelNotSupportedError(Exception): + pass + +class AiohttpSocksError(Exception): pass \ No newline at end of file diff --git a/g4f/typing.py b/g4f/typing.py index a6a62e3f..fd4fd047 100644 --- a/g4f/typing.py +++ b/g4f/typing.py @@ -1,5 +1,5 @@ import sys -from typing import Any, AsyncGenerator, Generator, NewType, Tuple, Union, List, Dict, Type, IO +from typing import Any, AsyncGenerator, Generator, NewType, Tuple, Union, List, Dict, Type, IO, Optional from PIL.Image import Image if sys.version_info >= (3, 8): @@ -22,6 +22,8 @@ __all__ = [ 'List', 'Dict', 'Type', + 'IO', + 'Optional', 'TypedDict', 'SHA256', 'CreateResult', -- cgit v1.2.3