diff options
Diffstat (limited to 'g4f/Provider/needs_auth')
| -rw-r--r-- | g4f/Provider/needs_auth/Bard.py | 45 | ||||
| -rw-r--r-- | g4f/Provider/needs_auth/OpenaiChat.py | 274 | ||||
| -rw-r--r-- | g4f/Provider/needs_auth/Poe.py | 93 | ||||
| -rw-r--r-- | g4f/Provider/needs_auth/Theb.py | 92 |
4 files changed, 311 insertions, 193 deletions
diff --git a/g4f/Provider/needs_auth/Bard.py b/g4f/Provider/needs_auth/Bard.py index b1df6909..77c029b8 100644 --- a/g4f/Provider/needs_auth/Bard.py +++ b/g4f/Provider/needs_auth/Bard.py @@ -4,7 +4,7 @@ import time from ...typing import CreateResult, Messages from ..base_provider import BaseProvider -from ..helper import WebDriver, format_prompt, get_browser +from ..helper import WebDriver, WebDriverSession, format_prompt class Bard(BaseProvider): url = "https://bard.google.com" @@ -18,34 +18,32 @@ class Bard(BaseProvider): messages: Messages, stream: bool, proxy: str = None, - browser: WebDriver = None, + web_driver: WebDriver = None, user_data_dir: str = None, headless: bool = True, **kwargs ) -> CreateResult: prompt = format_prompt(messages) - driver = browser if browser else get_browser(user_data_dir, headless, proxy) + session = WebDriverSession(web_driver, user_data_dir, headless, proxy=proxy) + with session as driver: + from selenium.webdriver.common.by import By + from selenium.webdriver.support.ui import WebDriverWait + from selenium.webdriver.support import expected_conditions as EC - from selenium.webdriver.common.by import By - from selenium.webdriver.support.ui import WebDriverWait - from selenium.webdriver.support import expected_conditions as EC - - try: - driver.get(f"{cls.url}/chat") - wait = WebDriverWait(driver, 10 if headless else 240) - wait.until(EC.visibility_of_element_located((By.CSS_SELECTOR, "div.ql-editor.textarea"))) - except: - # Reopen browser for login - if not browser: - driver.quit() - driver = get_browser(None, False, proxy) + try: driver.get(f"{cls.url}/chat") - wait = WebDriverWait(driver, 240) + wait = WebDriverWait(driver, 10 if headless else 240) wait.until(EC.visibility_of_element_located((By.CSS_SELECTOR, "div.ql-editor.textarea"))) - else: - raise RuntimeError("Prompt textarea not found. You may not be logged in.") + except: + # Reopen browser for login + if not web_driver: + driver = session.reopen(headless=False) + driver.get(f"{cls.url}/chat") + wait = WebDriverWait(driver, 240) + wait.until(EC.visibility_of_element_located((By.CSS_SELECTOR, "div.ql-editor.textarea"))) + else: + raise RuntimeError("Prompt textarea not found. You may not be logged in.") - try: # Add hook in XMLHttpRequest script = """ const _http_request_open = XMLHttpRequest.prototype.open; @@ -72,9 +70,4 @@ XMLHttpRequest.prototype.open = function(method, url) { yield chunk return else: - time.sleep(0.1) - finally: - if not browser: - driver.close() - time.sleep(0.1) - driver.quit()
\ No newline at end of file + time.sleep(0.1)
\ No newline at end of file diff --git a/g4f/Provider/needs_auth/OpenaiChat.py b/g4f/Provider/needs_auth/OpenaiChat.py index eccf2bd7..9fd90812 100644 --- a/g4f/Provider/needs_auth/OpenaiChat.py +++ b/g4f/Provider/needs_auth/OpenaiChat.py @@ -1,20 +1,64 @@ from __future__ import annotations -import uuid, json, time, asyncio +import uuid, json, asyncio from py_arkose_generator.arkose import get_values_for_request +from asyncstdlib.itertools import tee +from async_property import async_cached_property from ..base_provider import AsyncGeneratorProvider -from ..helper import get_browser, get_cookies, format_prompt, get_event_loop +from ..helper import get_browser, get_event_loop from ...typing import AsyncResult, Messages from ...requests import StreamSession -from ... import debug + +models = { + "gpt-3.5": "text-davinci-002-render-sha", + "gpt-3.5-turbo": "text-davinci-002-render-sha", + "gpt-4": "gpt-4", + "gpt-4-gizmo": "gpt-4-gizmo" +} class OpenaiChat(AsyncGeneratorProvider): url = "https://chat.openai.com" - needs_auth = True working = True + needs_auth = True supports_gpt_35_turbo = True - _access_token = None + supports_gpt_4 = True + _access_token: str = None + + @classmethod + async def create( + cls, + prompt: str = None, + model: str = "", + messages: Messages = [], + history_disabled: bool = False, + action: str = "next", + conversation_id: str = None, + parent_id: str = None, + **kwargs + ) -> Response: + if prompt: + messages.append({"role": "user", "content": prompt}) + generator = cls.create_async_generator( + model, + messages, + history_disabled=history_disabled, + action=action, + conversation_id=conversation_id, + parent_id=parent_id, + response_fields=True, + **kwargs + ) + fields: ResponseFields = await anext(generator) + if "access_token" not in kwargs: + kwargs["access_token"] = cls._access_token + return Response( + generator, + fields, + action, + messages, + kwargs + ) @classmethod async def create_async_generator( @@ -25,50 +69,56 @@ class OpenaiChat(AsyncGeneratorProvider): timeout: int = 120, access_token: str = None, auto_continue: bool = False, - cookies: dict = None, + history_disabled: bool = True, + action: str = "next", + conversation_id: str = None, + parent_id: str = None, + response_fields: bool = False, **kwargs ) -> AsyncResult: - proxies = {"https": proxy} + if not model: + model = "gpt-3.5" + elif model not in models: + raise ValueError(f"Model are not supported: {model}") + if not parent_id: + parent_id = str(uuid.uuid4()) if not access_token: - access_token = await cls.get_access_token(cookies, proxies) + access_token = await cls.get_access_token(proxy) headers = { "Accept": "text/event-stream", "Authorization": f"Bearer {access_token}", + "Cookie": 'intercom-device-id-dgkjq2bp=0f047573-a750-46c8-be62-6d54b56e7bf0; ajs_user_id=user-iv3vxisaoNodwWpxmNpMfekH; ajs_anonymous_id=fd91be0b-0251-4222-ac1e-84b1071e9ec1; __Host-next-auth.csrf-token=d2b5f67d56f7dd6a0a42ae4becf2d1a6577b820a5edc88ab2018a59b9b506886%7Ce5c33eecc460988a137cbc72d90ee18f1b4e2f672104f368046df58e364376ac; _cfuvid=gt_mA.q6rue1.7d2.AR0KHpbVBS98i_ppfi.amj2._o-1700353424353-0-604800000; cf_clearance=GkHCfPSFU.NXGcHROoe4FantnqmnNcluhTNHz13Tk.M-1700353425-0-1-dfe77f81.816e9bc2.714615da-0.2.1700353425; __Secure-next-auth.callback-url=https%3A%2F%2Fchat.openai.com; intercom-session-dgkjq2bp=UWdrS1hHazk5VXN1c0V5Q1F0VXdCQmsyTU9pVjJMUkNpWnFnU3dKWmtIdGwxTC9wbjZuMk5hcEc0NWZDOGdndS0tSDNiaDNmMEdIL1RHU1dFWDBwOHFJUT09--f754361b91fddcd23a13b288dcb2bf8c7f509e91; _uasid="Z0FBQUFBQmxXVnV0a3dmVno4czRhcDc2ZVcwaUpSNUdZejlDR25YSk5NYTJQQkpyNmRvOGxjTHMyTlAxWmJhaURrMVhjLXZxQXdZeVpBbU1aczA5WUpHT2dwaS1MOWc4MnhyNWFnbGRzeGdJcGFKT0ZRdnBTMVJHcGV2MGNTSnVQY193c0hqUWIycHhQRVF4dENlZ3phcDdZeHgxdVhoalhrZmtZME9NbWhMQjdVR3Vzc3FRRk0ybjJjNWMwTWtIRjdPb19lUkFtRmV2MDVqd1kwWU11QTYtQkdZenEzVHhLMGplY1hZM3FlYUt1cVZaNWFTRldleEJETzJKQjk1VTJScy1GUnMxUVZWMnVxYklxMjdockVZbkZyd1R4U1RtMnA1ZzlSeXphdmVOVk9xeEdrRkVOSjhwTVd1QzFtQjhBcWdDaE92Q1VlM2pwcjFQTXRuLVJNRVlZSGpIdlZ0aGV3PT0="; _dd_s=rum=0&expire=1700356244884; __Secure-next-auth.session-token=eyJhbGciOiJkaXIiLCJlbmMiOiJBMjU2R0NNIn0..3aK6Fbdy2_8f07bf.8eT2xgonrCnz7ySY6qXFsg3kzL6UQfXKAYaw3tyn-6_X9657zy47k9qGvmi9mF0QKozj5jau3_Ca62AQQ7FmeC6Y2F1urtzqrXqwTTsQ2LuzFPIQkx6KKb2DXc8zW2-oyEzJ_EY5yxfLB2RlRkSh3M7bYNZh4_ltEcfkj38s_kIPGMxv34udtPWGWET99MCjkdwQWXylJag4s0fETA0orsBAKnGCyqAUNJbb_D7BYtGSV-MQ925kZMG6Di_QmfO0HQWURDYjmdRNcuy1PT_xJ1DJko8sjL42i4j3RhkNDkhqCIqyYImz2eHFWHW7rYKxTkrBhlCPMS5hRdcCswD7JYPcSBiwnVRYgyOocFGXoFvQgIZ2FX9NiZ3SMEVM1VwIGSE-qH0H2nMa8_iBvsOgOWJgKjVAvzzyzZvRVDUUHzJrikSFPNONVDU3h-04c1kVL4qIu9DfeTPN7n8AvNmYwMbro0L9-IUAeXNo4-pwF0Kt-AtTsamqWvMqnK4O_YOyLnDDlvkmnOvDC2d5uinwlQIxr6APO6qFfGLlHiLZemKoekxEE1Fx70dl-Ouhk1VIzbF3OC6XNNxeBm9BUYUiHdL0wj2H9rHgX4cz6ZmS_3VTgpD6UJh-evu5KJ2gIvjYmVbyzEN0aPNDxfvBaOm-Ezpy4bUJ2bUrOwNn-0knWkDiTvjYmNhCyefPCtCF6rpKNay8PCw_yh79C4SdEP6Q4V7LI0Tvdi5uz7kLCiBC4AT9L0ao1WDX03mkUOpjvzHDvPLmj8chW3lTVm_kA0eYGQY4wT0jzleWlfV0Q8rB2oYECNLWksA3F1zlGfcl4lQjprvTXRePkvAbMpoJEsZD3Ylq7-foLDLk4-M2LYAFZDs282AY04sFjAjQBxTELFCCuDgTIgTXSIskY_XCxpVXDbdLlbCJY7XVK45ybwtfqwlKRp8Mo0B131uQAFc-migHaUaoGujxJJk21bP8F0OmhNYHBo4FQqE1rQm2JH5bNM7txKeh5KXdJgVUVbRSr7OIp_OF5-Bx_v9eRBGAIDkue26E2-O8Rnrp5zQ5TnvecQLDaUzWavCLPwsZ0_gsOLBxNOmauNYZtF8IElCsQSFDdhoiMxXsYUm4ZYKEAy3GWq8HGTAvBhNkh1hvnI7y-d8-DOaZf_D_D98-olZfm-LUkeosLNpPB9rxYMqViCiW3KrXE9Yx0wlFm5ePKaVvR7Ym_EPhSOhJBKFPCvdTdMZSNPUcW0ZJBVByq0A9sxD51lYq3gaFyqh94S4s_ox182AQ3szGzHkdgLcnQmJG9OYvKxAVcd43eg6_gODAYhx02GjbMw-7JTAhyXSeCrlMteHyOXl8hai-3LilC3PmMzi7Vbu49dhF1s4LcVlUowen5ira44rQQaB26mdaOUoQfodgt66M3RTWGPXyK1Nb72AzSXsCKyaQPbzeb6cN0fdGSdG4ktwvR04eFNEkquo_3aKu2GmUKTD0XcRx9dYrfXjgY-X1DDTVs1YND2gRhdx7FFEeBVjtbj2UqmG3Rvd4IcHGe7OnYWw2MHDcol68SsR1KckXWwWREz7YTGUnDB2M1kx_H4W2mjclytnlHOnYU3RflegRPeSTbdzUZJvGKXCCz45luHkQWN_4DExE76D-9YqbFIz-RY5yL4h-Zs-i2xjm2K-4xCMM9nQIOqhLMqixIZQ2ldDAidKoYtbs5ppzbcBLyrZM96bq9DwRBY3aacqWdlRd-TfX0wv5KO4fo0sSh5FsuhuN0zcEV_NNXgqIEM_p14EcPqgbrAvCBQ8os70TRBQLXiF0EniSofGjxwF8kQvUk3C6Wfc8cTTeN-E6GxCVTn91HBwA1iSEZlRLMVb8_BcRJNqwbgnb_07jR6-eo42u88CR3KQdAWwbQRdMxsURFwZ0ujHXVGG0Ll6qCFBcHXWyDO1x1yHdHnw8_8yF26pnA2iPzrFR-8glMgIA-639sLuGAxjO1_ZuvJ9CAB41Az9S_jaZwaWy215Hk4-BRYD-MKmHtonwo3rrxhE67WJgbbu14efsw5nT6ow961pffgwXov5VA1Rg7nv1E8RvQOx7umWW6o8R4W6L8f2COsmPTXfgwIjoJKkjhUqAQ8ceG7cM0ET-38yaC0ObU8EkXfdGGgxI28qTEZWczG66_iM4hw7QEGCY5Cz2kbO6LETAiw9OsSigtBvDS7f0Ou0bZ41pdK7G3FmvdZAnjWPjObnDF4k4uWfn7mzt0fgj3FyqK20JezRDyGuAbUUhOvtZpc9sJpzxR34eXEZTouuALrHcGuNij4z6rx51FrQsaMtiup8QVrhtZbXtKLMYnWYSbkhuTeN2wY-xV1ZUsQlakIZszzGF7kuIG87KKWMpuPMvbXjz6Pp_gWJiIC6aQuk8xl5g0iBPycf_6Q-MtpuYxzNE2TpI1RyR9mHeXmteoRzrFiWp7yEC-QGNFyAJgxTqxM3CjHh1Jt6IddOsmn89rUo1dZM2Smijv_fbIv3avXLkIPX1KZjILeJCtpU0wAdsihDaRiRgDdx8fG__F8zuP0n7ziHas73cwrfg-Ujr6DhC0gTNxyd9dDA_oho9N7CQcy6EFmfNF2te7zpLony0859jtRv2t1TnpzAa1VvMK4u6mXuJ2XDo04_6GzLO3aPHinMdl1BcIAWnqAqWAu3euGFLTHOhXlfijut9N1OCifd_zWjhVtzlR39uFeCQBU5DyQArzQurdoMx8U1ETsnWgElxGSStRW-YQoPsAJ87eg9trqKspFpTVlAVN3t1GtoEAEhcwhe81SDssLmKGLc.7PqS6jRGTIfgTPlO7Ognvg; __cf_bm=VMWoAKEB45hQSwxXtnYXcurPaGZDJS4dMi6dIMFLwdw-1700355394-0-ATVsbq97iCaTaJbtYr8vtg1Zlbs3nLrJLKVBHYa2Jn7hhkGclqAy8Gbyn5ePEhDRqj93MsQmtayfYLqY5n4WiLY=; __cflb=0H28vVfF4aAyg2hkHFH9CkdHRXPsfCUf6VpYf2kz3RX' } - messages = [ - { - "id": str(uuid.uuid4()), - "author": {"role": "user"}, - "content": {"content_type": "text", "parts": [format_prompt(messages)]}, - }, - ] - message_id = str(uuid.uuid4()) - data = { - "action": "next", - "arkose_token": await get_arkose_token(proxy), - "messages": messages, - "conversation_id": None, - "parent_message_id": message_id, - "model": "text-davinci-002-render-sha", - "history_and_training_disabled": not auto_continue, - } - conversation_id = None - end_turn = False - while not end_turn: - if not auto_continue: - end_turn = True - async with StreamSession( - proxies=proxies, - headers=headers, - impersonate="chrome107", - timeout=timeout - ) as session: + async with StreamSession( + proxies={"https": proxy}, + impersonate="chrome110", + headers=headers, + timeout=timeout + ) as session: + data = { + "action": action, + "arkose_token": await get_arkose_token(proxy, timeout), + "conversation_id": conversation_id, + "parent_message_id": parent_id, + "model": models[model], + "history_and_training_disabled": history_disabled and not auto_continue, + } + if action != "continue": + data["messages"] = [{ + "id": str(uuid.uuid4()), + "author": {"role": "user"}, + "content": {"content_type": "text", "parts": [messages[-1]["content"]]}, + }] + first = True + end_turn = EndTurn() + while first or auto_continue and not end_turn.is_end: + first = False async with session.post(f"{cls.url}/backend-api/conversation", json=data) as response: try: response.raise_for_status() except: - raise RuntimeError(f"Response: {await response.text()}") - last_message = "" + raise RuntimeError(f"Error {response.status_code}: {await response.text()}") + last_message = 0 async for line in response.iter_lines(): if line.startswith(b"data: "): line = line[6:] @@ -82,50 +132,52 @@ class OpenaiChat(AsyncGeneratorProvider): continue if "error" in line and line["error"]: raise RuntimeError(line["error"]) - end_turn = line["message"]["end_turn"] - message_id = line["message"]["id"] - if line["conversation_id"]: - conversation_id = line["conversation_id"] if "message_type" not in line["message"]["metadata"]: continue - if line["message"]["metadata"]["message_type"] in ("next", "continue"): + if line["message"]["author"]["role"] != "assistant": + continue + if line["message"]["metadata"]["message_type"] in ("next", "continue", "variant"): + conversation_id = line["conversation_id"] + parent_id = line["message"]["id"] + if response_fields: + response_fields = False + yield ResponseFields(conversation_id, parent_id, end_turn) new_message = line["message"]["content"]["parts"][0] - yield new_message[len(last_message):] - last_message = new_message - if end_turn: - return + yield new_message[last_message:] + last_message = len(new_message) + if "finish_details" in line["message"]["metadata"]: + if line["message"]["metadata"]["finish_details"]["type"] == "max_tokens": + end_turn.end() + data = { "action": "continue", - "arkose_token": await get_arkose_token(proxy), + "arkose_token": await get_arkose_token(proxy, timeout), "conversation_id": conversation_id, - "parent_message_id": message_id, - "model": "text-davinci-002-render-sha", + "parent_message_id": parent_id, + "model": models[model], "history_and_training_disabled": False, } await asyncio.sleep(5) @classmethod - async def browse_access_token(cls) -> str: + async def browse_access_token(cls, proxy: str = None) -> str: def browse() -> str: try: from selenium.webdriver.common.by import By from selenium.webdriver.support.ui import WebDriverWait from selenium.webdriver.support import expected_conditions as EC - driver = get_browser() + driver = get_browser("~/openai", proxy=proxy) except ImportError: return - - driver.get(f"{cls.url}/") try: + driver.get(f"{cls.url}/") WebDriverWait(driver, 1200).until( EC.presence_of_element_located((By.ID, "prompt-textarea")) ) javascript = "return (await (await fetch('/api/auth/session')).json())['accessToken']" return driver.execute_script(javascript) finally: - driver.close() - time.sleep(0.1) driver.quit() loop = get_event_loop() return await loop.run_in_executor( @@ -134,22 +186,9 @@ class OpenaiChat(AsyncGeneratorProvider): ) @classmethod - async def fetch_access_token(cls, cookies: dict, proxies: dict = None) -> str: - async with StreamSession(proxies=proxies, cookies=cookies, impersonate="chrome107") as session: - async with session.get(f"{cls.url}/api/auth/session") as response: - response.raise_for_status() - auth = await response.json() - if "accessToken" in auth: - return auth["accessToken"] - - @classmethod - async def get_access_token(cls, cookies: dict = None, proxies: dict = None) -> str: - if not cls._access_token: - cookies = cookies if cookies else get_cookies("chat.openai.com") - if cookies: - cls._access_token = await cls.fetch_access_token(cookies, proxies) + async def get_access_token(cls, proxy: str = None) -> str: if not cls._access_token: - cls._access_token = await cls.browse_access_token() + cls._access_token = await cls.browse_access_token(proxy) if not cls._access_token: raise RuntimeError("Read access token failed") return cls._access_token @@ -163,12 +202,11 @@ class OpenaiChat(AsyncGeneratorProvider): ("stream", "bool"), ("proxy", "str"), ("access_token", "str"), - ("cookies", "dict[str, str]") ] param = ", ".join([": ".join(p) for p in params]) return f"g4f.provider.{cls.__name__} supports: ({param})" -async def get_arkose_token(proxy: str = None) -> str: +async def get_arkose_token(proxy: str = None, timeout: int = None) -> str: config = { "pkey": "3D86FBBA-9D22-402A-B512-3420086BA6CC", "surl": "https://tcr9i.chat.openai.com", @@ -181,10 +219,98 @@ async def get_arkose_token(proxy: str = None) -> str: async with StreamSession( proxies={"https": proxy}, impersonate="chrome107", + timeout=timeout ) as session: async with session.post(**args_for_request) as response: response.raise_for_status() decoded_json = await response.json() if "token" in decoded_json: return decoded_json["token"] - raise RuntimeError(f"Response: {decoded_json}")
\ No newline at end of file + raise RuntimeError(f"Response: {decoded_json}") + +class EndTurn(): + def __init__(self): + self.is_end = False + + def end(self): + self.is_end = True + +class ResponseFields(): + def __init__( + self, + conversation_id: str, + message_id: str, + end_turn: EndTurn + ): + self.conversation_id = conversation_id + self.message_id = message_id + self._end_turn = end_turn + +class Response(): + def __init__( + self, + generator: AsyncResult, + fields: ResponseFields, + action: str, + messages: Messages, + options: dict + ): + self.aiter, self.copy = tee(generator) + self.fields = fields + self.action = action + self._messages = messages + self._options = options + + def __aiter__(self): + return self.aiter + + @async_cached_property + async def message(self) -> str: + return "".join([chunk async for chunk in self.copy]) + + async def next(self, prompt: str, **kwargs) -> Response: + return await OpenaiChat.create( + **self._options, + prompt=prompt, + messages=await self.messages, + action="next", + conversation_id=self.fields.conversation_id, + parent_id=self.fields.message_id, + **kwargs + ) + + async def do_continue(self, **kwargs) -> Response: + if self.end_turn: + raise RuntimeError("Can't continue message. Message already finished.") + return await OpenaiChat.create( + **self._options, + messages=await self.messages, + action="continue", + conversation_id=self.fields.conversation_id, + parent_id=self.fields.message_id, + **kwargs + ) + + async def variant(self, **kwargs) -> Response: + if self.action != "next": + raise RuntimeError("Can't create variant with continue or variant request.") + return await OpenaiChat.create( + **self._options, + messages=self._messages, + action="variant", + conversation_id=self.fields.conversation_id, + parent_id=self.fields.message_id, + **kwargs + ) + + @async_cached_property + async def messages(self): + messages = self._messages + messages.append({ + "role": "assistant", "content": await self.message + }) + return messages + + @property + def end_turn(self): + return self.fields._end_turn.is_end
\ No newline at end of file diff --git a/g4f/Provider/needs_auth/Poe.py b/g4f/Provider/needs_auth/Poe.py index a894bcb1..1c8c97d7 100644 --- a/g4f/Provider/needs_auth/Poe.py +++ b/g4f/Provider/needs_auth/Poe.py @@ -4,7 +4,7 @@ import time from ...typing import CreateResult, Messages from ..base_provider import BaseProvider -from ..helper import WebDriver, format_prompt, get_browser +from ..helper import WebDriver, WebDriverSession, format_prompt models = { "meta-llama/Llama-2-7b-chat-hf": {"name": "Llama-2-7b"}, @@ -33,7 +33,7 @@ class Poe(BaseProvider): messages: Messages, stream: bool, proxy: str = None, - browser: WebDriver = None, + web_driver: WebDriver = None, user_data_dir: str = None, headless: bool = True, **kwargs @@ -43,56 +43,54 @@ class Poe(BaseProvider): elif model not in models: raise ValueError(f"Model are not supported: {model}") prompt = format_prompt(messages) - driver = browser if browser else get_browser(user_data_dir, headless, proxy) - script = """ -window._message = window._last_message = ""; -window._message_finished = false; -class ProxiedWebSocket extends WebSocket { - constructor(url, options) { - super(url, options); - this.addEventListener("message", (e) => { - const data = JSON.parse(JSON.parse(e.data)["messages"][0])["payload"]["data"]; - if ("messageAdded" in data) { - if (data["messageAdded"]["author"] != "human") { - window._message = data["messageAdded"]["text"]; - if (data["messageAdded"]["state"] == "complete") { - window._message_finished = true; + session = WebDriverSession(web_driver, user_data_dir, headless, proxy=proxy) + with session as driver: + from selenium.webdriver.common.by import By + from selenium.webdriver.support.ui import WebDriverWait + from selenium.webdriver.support import expected_conditions as EC + + driver.execute_cdp_cmd("Page.addScriptToEvaluateOnNewDocument", { + "source": """ + window._message = window._last_message = ""; + window._message_finished = false; + class ProxiedWebSocket extends WebSocket { + constructor(url, options) { + super(url, options); + this.addEventListener("message", (e) => { + const data = JSON.parse(JSON.parse(e.data)["messages"][0])["payload"]["data"]; + if ("messageAdded" in data) { + if (data["messageAdded"]["author"] != "human") { + window._message = data["messageAdded"]["text"]; + if (data["messageAdded"]["state"] == "complete") { + window._message_finished = true; + } } } - } - }); - } -} -window.WebSocket = ProxiedWebSocket; -""" - driver.execute_cdp_cmd("Page.addScriptToEvaluateOnNewDocument", { - "source": script - }) - - from selenium.webdriver.common.by import By - from selenium.webdriver.support.ui import WebDriverWait - from selenium.webdriver.support import expected_conditions as EC + }); + } + } + window.WebSocket = ProxiedWebSocket; + """ + }) - try: - driver.get(f"{cls.url}/{models[model]['name']}") - wait = WebDriverWait(driver, 10 if headless else 240) - wait.until(EC.visibility_of_element_located((By.CSS_SELECTOR, "textarea[class^='GrowingTextArea']"))) - except: - # Reopen browser for login - if not browser: - driver.quit() - driver = get_browser(None, False, proxy) + try: driver.get(f"{cls.url}/{models[model]['name']}") - wait = WebDriverWait(driver, 240) + wait = WebDriverWait(driver, 10 if headless else 240) wait.until(EC.visibility_of_element_located((By.CSS_SELECTOR, "textarea[class^='GrowingTextArea']"))) - else: - raise RuntimeError("Prompt textarea not found. You may not be logged in.") + except: + # Reopen browser for login + if not web_driver: + driver = session.reopen(headless=False) + driver.get(f"{cls.url}/{models[model]['name']}") + wait = WebDriverWait(driver, 240) + wait.until(EC.visibility_of_element_located((By.CSS_SELECTOR, "textarea[class^='GrowingTextArea']"))) + else: + raise RuntimeError("Prompt textarea not found. You may not be logged in.") - driver.find_element(By.CSS_SELECTOR, "footer textarea[class^='GrowingTextArea']").send_keys(prompt) - driver.find_element(By.CSS_SELECTOR, "footer button[class*='ChatMessageSendButton']").click() + driver.find_element(By.CSS_SELECTOR, "footer textarea[class^='GrowingTextArea']").send_keys(prompt) + driver.find_element(By.CSS_SELECTOR, "footer button[class*='ChatMessageSendButton']").click() - try: script = """ if(window._message && window._message != window._last_message) { try { @@ -113,9 +111,4 @@ if(window._message && window._message != window._last_message) { elif chunk != "": break else: - time.sleep(0.1) - finally: - if not browser: - driver.close() - time.sleep(0.1) - driver.quit()
\ No newline at end of file + time.sleep(0.1)
\ No newline at end of file diff --git a/g4f/Provider/needs_auth/Theb.py b/g4f/Provider/needs_auth/Theb.py index 89c69727..cf33f0c6 100644 --- a/g4f/Provider/needs_auth/Theb.py +++ b/g4f/Provider/needs_auth/Theb.py @@ -4,7 +4,7 @@ import time from ...typing import CreateResult, Messages from ..base_provider import BaseProvider -from ..helper import WebDriver, format_prompt, get_browser +from ..helper import WebDriver, WebDriverSession, format_prompt models = { "theb-ai": "TheB.AI", @@ -44,26 +44,60 @@ class Theb(BaseProvider): messages: Messages, stream: bool, proxy: str = None, - browser: WebDriver = None, - headless: bool = True, + web_driver: WebDriver = None, + virtual_display: bool = True, **kwargs ) -> CreateResult: if model in models: model = models[model] prompt = format_prompt(messages) - driver = browser if browser else get_browser(None, headless, proxy) + web_session = WebDriverSession(web_driver, virtual_display=virtual_display, proxy=proxy) + with web_session as driver: + from selenium.webdriver.common.by import By + from selenium.webdriver.support.ui import WebDriverWait + from selenium.webdriver.support import expected_conditions as EC + from selenium.webdriver.common.keys import Keys - from selenium.webdriver.common.by import By - from selenium.webdriver.support.ui import WebDriverWait - from selenium.webdriver.support import expected_conditions as EC - from selenium.webdriver.common.keys import Keys + # Register fetch hook + script = """ +window._fetch = window.fetch; +window.fetch = (url, options) => { + // Call parent fetch method + const result = window._fetch(url, options); + if (!url.startsWith("/api/conversation")) { + return result; + } + // Load response reader + result.then((response) => { + if (!response.body.locked) { + window._reader = response.body.getReader(); + } + }); + // Return dummy response + return new Promise((resolve, reject) => { + resolve(new Response(new ReadableStream())) + }); +} +window._last_message = ""; +""" + driver.execute_cdp_cmd("Page.addScriptToEvaluateOnNewDocument", { + "source": script + }) + + try: + driver.get(f"{cls.url}/home") + wait = WebDriverWait(driver, 5) + wait.until(EC.visibility_of_element_located((By.ID, "textareaAutosize"))) + except: + driver = web_session.reopen() + driver.execute_cdp_cmd("Page.addScriptToEvaluateOnNewDocument", { + "source": script + }) + driver.get(f"{cls.url}/home") + wait = WebDriverWait(driver, 240) + wait.until(EC.visibility_of_element_located((By.ID, "textareaAutosize"))) - - try: - driver.get(f"{cls.url}/home") - wait = WebDriverWait(driver, 10 if headless else 240) - wait.until(EC.visibility_of_element_located((By.TAG_NAME, "body"))) - time.sleep(0.1) + time.sleep(200) try: driver.find_element(By.CSS_SELECTOR, ".driver-overlay").click() driver.find_element(By.CSS_SELECTOR, ".driver-overlay").click() @@ -87,29 +121,6 @@ class Theb(BaseProvider): button = container.find_element(By.CSS_SELECTOR, "button.btn-blue.btn-small.border") button.click() - # Register fetch hook - script = """ -window._fetch = window.fetch; -window.fetch = (url, options) => { - // Call parent fetch method - const result = window._fetch(url, options); - if (!url.startsWith("/api/conversation")) { - return result; - } - // Load response reader - result.then((response) => { - if (!response.body.locked) { - window._reader = response.body.getReader(); - } - }); - // Return dummy response - return new Promise((resolve, reject) => { - resolve(new Response(new ReadableStream())) - }); -} -window._last_message = ""; -""" - driver.execute_script(script) # Submit prompt wait.until(EC.visibility_of_element_located((By.ID, "textareaAutosize"))) @@ -150,9 +161,4 @@ return ''; elif chunk != "": break else: - time.sleep(0.1) - finally: - if not browser: - driver.close() - time.sleep(0.1) - driver.quit()
\ No newline at end of file + time.sleep(0.1)
\ No newline at end of file |