Add: Supports the chatgpt automatic ws mode

re_gpt/async_chatgpt.py

@@ -3,6 +3,9 @@
 import inspect
 import json
 import uuid
+import re
+import websockets
+import base64
 from typing import AsyncGenerator, Callable, Optional
 
 from curl_cffi.requests import AsyncSession
@@ -20,6 +23,8 @@
 USER_AGENT = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/110.0.0.0 Safari/537.36"
 CHATGPT_API = "https://chat.openai.com/backend-api/{}"
 BACKUP_ARKOSE_TOKEN_GENERATOR = "https://arkose-token-generator.zaieem.repl.co/token"
+WS_REGISTER_URL = CHATGPT_API.format("register-websocket")
+
 MODELS = {
     "gpt-4": {"slug": "gpt-4", "needs_arkose_token": True},
     "gpt-3.5": {"slug": "text-davinci-002-render-sha", "needs_arkose_token": False},
@@ -92,10 +97,10 @@ async def chat(self, user_input: str) -> AsyncGenerator[dict, None]:
         try:
             full_message = None
             while True:
-                response = self.send_message(payload=payload)
+                response = self.send_message(payload=payload) if not self.chatgpt.websocket_mode else self.send_websocket_message(payload=payload)
                 async for chunk in response:
-                    decoded_chunk = chunk.decode()
-
+                    decoded_chunk = chunk.decode() if isinstance(chunk, bytes) else chunk
+                    
                     server_response += decoded_chunk
                     for line in decoded_chunk.splitlines():
                         if not line.startswith("data: "):
@@ -168,6 +173,47 @@ def content_callback(chunk):
             if chunk is None:
                 break
             yield chunk
+
+    async def send_websocket_message(self, payload: dict) -> AsyncGenerator[str, None]:
+        """
+        Send a message payload via WebSocket and receive the response.
+
+        Args:
+            payload (dict): Payload containing message information.
+
+        Yields:
+            str: Chunk of data received as a response.
+        """
+        await self.chatgpt.ensure_websocket()
+
+        response_queue = asyncio.Queue()
+        websocket_request_id = None
+
+        async def perform_request():
+            nonlocal websocket_request_id
+
+            url = CHATGPT_API.format("conversation")
+            response = (await self.chatgpt.session.post(
+                url=url,
+                headers=self.chatgpt.build_request_headers(),
+                json=payload,
+            )).json()
+
+            websocket_request_id = response.get("websocket_request_id")
+
+            if websocket_request_id not in self.chatgpt.ws_conversation_map:
+                self.chatgpt.ws_conversation_map[websocket_request_id] = response_queue
+
+        asyncio.create_task(perform_request())
+
+        while True:
+            chunk = await response_queue.get()
+            if chunk is None:
+                break
+            yield chunk
+
+        del self.chatgpt.ws_conversation_map[websocket_request_id]
+
 
     async def build_message_payload(self, user_input: str) -> dict:
         """
@@ -201,6 +247,9 @@ async def build_message_payload(self, user_input: str) -> dict:
             "parent_message_id": str(uuid.uuid4())
             if not self.parent_id
             else self.parent_id,
+            "websocket_request_id": str(uuid.uuid4())
+            if self.chatgpt.websocket_mode
+            else None,
         }
 
         return payload
@@ -310,6 +359,7 @@ def __init__(
         exit_callback_function: Optional[Callable] = None,
         auth_token: Optional[str] = None,
         generate_arkose_token: Optional[bool] = False,
+        websocket_mode: Optional[bool] = False,
     ):
         """
         Initializes an instance of the class.
@@ -320,6 +370,7 @@ def __init__(
             exit_callback_function (Optional[callable]): A function to be called on exit. Defaults to None.
             auth_token (Optional[str]): An authentication token. Defaults to None.
             generate_arkose_token (Optional[bool]): Toggle whether to generate and send arkose-token in the payload. Defaults to False.
+            websocket_mode (Optional[bool]): Toggle whether to use WebSocket for chat. Defaults to False.
         """
         self.proxies = proxies
         self.exit_callback_function = exit_callback_function
@@ -332,6 +383,10 @@ def __init__(
         self.session_token = session_token
         self.auth_token = auth_token
         self.session = None
+
+        self.websocket_mode = websocket_mode
+        self.ws_loop = None
+        self.ws_conversation_map = {}
 
     async def __aenter__(self):
         self.session = AsyncSession(
@@ -351,6 +406,12 @@ async def __aenter__(self):
                 raise TokenNotProvided
             self.auth_token = await self.fetch_auth_token()
 
+        if not self.websocket_mode:
+           self.websocket_mode = await self.check_websocket_availability()
+
+        if self.websocket_mode:
+            await self.ensure_websocket()
+
         return self
 
     async def __aexit__(self, *_):
@@ -499,3 +560,53 @@ async def retrieve_chats(
         )
 
         return response.json()
+
+    async def check_websocket_availability(self) -> bool:
+        """
+        Check if WebSocket is available.
+
+        Returns:
+            bool: True if WebSocket is available, otherwise False.
+        """
+        url = CHATGPT_API.format("accounts/check/v4-2023-04-27")
+        response = (await self.session.get(
+            url=url, headers=self.build_request_headers()
+        )).json()
+
+        if 'account_ordering' in response and 'accounts' in response:
+            account_id = response['account_ordering'][0]
+            if account_id in response['accounts']:
+                return 'shared_websocket' in response['accounts'][account_id]['features']
+
+        return False
+
+    async def ensure_websocket(self):
+        if not self.ws_loop:
+            ws_url_rsp = (await self.session.post(WS_REGISTER_URL, headers=self.build_request_headers())).json()
+            ws_url = ws_url_rsp['wss_url']
+            access_token = self.extract_access_token(ws_url)
+            self.ws_loop = asyncio.create_task(self.listen_to_websocket(ws_url, access_token))
+
+    def extract_access_token(self, url):
+        match = re.search(r'access_token=([^&]*)', url)
+        if match:
+            return match.group(1)
+        else:
+            return None
+
+    async def listen_to_websocket(self, ws_url: str, access_token: str):
+        headers = {'Authorization': f'Bearer {access_token}'}
+        async with websockets.connect(ws_url, extra_headers=headers) as websocket:
+            while True:
+                message = await websocket.recv()
+                message_data = json.loads(message)
+                body_encoded = message_data.get("body", "")
+                ws_id = message_data.get("websocket_request_id", "")
+                decoded_body = base64.b64decode(body_encoded).decode('utf-8')
+                response_queue = self.ws_conversation_map.get(ws_id)
+                if response_queue is None:
+                    continue
+                response_queue.put_nowait(decoded_body)
+                if '[DONE]' in decoded_body or '[ERROR]' in decoded_body:
+                    await response_queue.put(None)
+                    continue

re_gpt/sync_chatgpt.py

@@ -2,6 +2,11 @@
 import inspect
 import time
 import uuid
+import websockets
+from websockets.exceptions import ConnectionClosed
+import json
+import base64
+import asyncio
 from queue import Queue
 from threading import Thread
 from typing import Callable, Generator, Optional
@@ -15,6 +20,7 @@
     AsyncChatGPT,
     AsyncConversation,
     MODELS,
+    WS_REGISTER_URL,
 )
 from .errors import (
     BackendError,
@@ -90,9 +96,9 @@ def chat(self, user_input: str) -> Generator[dict, None, None]:
         try:
             full_message = None
             while True:
-                response = self.send_message(payload=payload)
+                response = self.send_message(payload=payload) if not self.chatgpt.websocket_mode else self.send_websocket_message(payload=payload)
                 for chunk in response:
-                    decoded_chunk = chunk.decode()
+                    decoded_chunk = chunk.decode() if not self.chatgpt.websocket_mode else chunk
 
                     server_response += decoded_chunk
                     for line in decoded_chunk.splitlines():
@@ -166,6 +172,45 @@ def content_callback(chunk):
             if chunk is None:
                 break
             yield chunk
+
+    def send_websocket_message(self, payload: dict) -> Generator[str, None, None]:
+        """
+        Send a message payload via WebSocket and receive the response.
+
+        Args:
+            payload (dict): Payload containing message information.
+
+        Yields:
+            str: Chunk of data received as a response.
+        """
+
+        response_queue = Queue()
+        websocket_request_id = None
+
+        def perform_request():
+            nonlocal websocket_request_id
+
+            url = CHATGPT_API.format("conversation")
+            response = (self.chatgpt.session.post(
+                url=url,
+                headers=self.chatgpt.build_request_headers(),
+                json=payload,
+            )).json()
+
+            websocket_request_id = response.get("websocket_request_id")
+
+            if websocket_request_id not in self.chatgpt.ws_conversation_map:
+                self.chatgpt.ws_conversation_map[websocket_request_id] = response_queue
+
+        Thread(target=perform_request).start()
+
+        while True:
+            chunk = response_queue.get()
+            if chunk is None:
+                break
+            yield chunk
+
+        del self.chatgpt.ws_conversation_map[websocket_request_id]
 
     def build_message_payload(self, user_input: str) -> dict:
         """
@@ -279,6 +324,7 @@ def __init__(
         session_token: Optional[str] = None,
         exit_callback_function: Optional[Callable] = None,
         auth_token: Optional[str] = None,
+        websocket_mode: Optional[bool] = False,
     ):
         """
         Initializes an instance of the class.
@@ -288,14 +334,19 @@ def __init__(
             session_token (Optional[str]): A session token. Defaults to None.
             exit_callback_function (Optional[callable]): A function to be called on exit. Defaults to None.
             auth_token (Optional[str]): An authentication token. Defaults to None.
+            websocket_mode (Optional[bool]): Toggle whether to use WebSocket for chat. Defaults to False.
         """
         super().__init__(
             proxies=proxies,
             session_token=session_token,
             exit_callback_function=exit_callback_function,
             auth_token=auth_token,
+            websocket_mode=websocket_mode,
         )
 
+        self.stop_websocket_flag = False
+        self.stop_websocket = None
+
     def __enter__(self):
         self.session = Session(
             impersonate="chrome110", timeout=99999, proxies=self.proxies
@@ -314,6 +365,17 @@ def __enter__(self):
             if self.session_token is None:
                 raise TokenNotProvided
             self.auth_token = self.fetch_auth_token()
+
+        # automaticly check the status of websocket_mode
+        if not self.websocket_mode:
+            self.websocket_mode = self.check_websocket_availability()
+            print(f"WebSocket mode is {'enabled' if self.websocket_mode else 'disabled'}")
+
+        if self.websocket_mode:
+            def run_websocket():
+                asyncio.run(self.ensure_websocket())
+            self.ws_loop = Thread(target=run_websocket)
+            self.ws_loop.start()
 
         return self
 
@@ -325,6 +387,10 @@ def __exit__(self, *args):
         finally:
             self.session.close()
 
+        if self.websocket_mode:
+            self.stop_websocket_flag = True
+            self.ws_loop.join()
+
     def get_conversation(self, conversation_id: str) -> SyncConversation:
         """
         Makes an instance of class Conversation and return it.
@@ -442,3 +508,61 @@ def retrieve_chats(
         )
 
         return response.json()
+
+    def check_websocket_availability(self) -> bool:
+        """
+        Check if WebSocket is available.
+
+        Returns:
+            bool: True if WebSocket is available, otherwise False.
+        """
+        url = CHATGPT_API.format("accounts/check/v4-2023-04-27")
+        response = (self.session.get(
+            url=url, headers=self.build_request_headers()
+        )).json()
+
+        if 'account_ordering' in response and 'accounts' in response:
+            account_id = response['account_ordering'][0]
+            if account_id in response['accounts']:
+                return 'shared_websocket' in response['accounts'][account_id]['features']
+
+        return False
+
+    async def ensure_websocket(self):
+        ws_url_rsp = self.session.post(WS_REGISTER_URL, headers=self.build_request_headers()).json()
+        ws_url = ws_url_rsp['wss_url']
+        access_token = self.extract_access_token(ws_url)
+        asyncio.create_task(self.ensure_close_websocket())
+        await self.listen_to_websocket(ws_url, access_token)
+
+    async def ensure_close_websocket(self):
+        while True:
+            if self.stop_websocket_flag:
+                break
+            await asyncio.sleep(1)
+        await self.stop_websocket()
+
+    async def listen_to_websocket(self, ws_url: str, access_token: str):
+        headers = {'Authorization': f'Bearer {access_token}'}
+        async with websockets.connect(ws_url, extra_headers=headers) as websocket:
+            async def stop_websocket():
+                await websocket.close()
+            self.stop_websocket = stop_websocket
+
+            while True:
+                message = None
+                try:
+                    message = await websocket.recv()
+                except ConnectionClosed:
+                    break
+                message_data = json.loads(message)
+                body_encoded = message_data.get("body", "")
+                ws_id = message_data.get("websocket_request_id", "")
+                decoded_body = base64.b64decode(body_encoded).decode('utf-8')
+                response_queue = self.ws_conversation_map.get(ws_id)
+                if response_queue is None:
+                    continue
+                response_queue.put_nowait(decoded_body)
+                if '[DONE]' in decoded_body or '[ERROR]' in decoded_body:
+                    response_queue.put(None)
+                    continue

requirements.txt

@@ -1 +1,2 @@
 curl_cffi==0.5.9
+websockets==12.0

setup.py

@@ -12,5 +12,8 @@
         "Bug Tracker": "https://github.com/Zai-Kun/reverse-engineered-chatgpt/issues",
     },
     packages=find_packages(),
-    install_requires=["curl_cffi==0.5.9"],
+    install_requires=[
+        "curl_cffi==0.5.9",
+        "websockets==12.0"
+    ],
 )