feat: add proxy support (#122)

requirements.txt

@@ -1,3 +1,3 @@
 asyncio
 requests
-websockets
+aiohttp

src/EdgeGPT.py

@@ -9,14 +9,16 @@
 import random
 import sys
 import uuid
+import urllib.parse
 from enum import Enum
 from typing import Generator
 from typing import Literal
 from typing import Optional
 from typing import Union
 
 import httpx
-import websockets.client as websockets
+import aiohttp
+from aiohttp import ClientWebSocketResponse
 
 DELIMITER = "\x1e"
 
@@ -171,14 +173,18 @@ class Conversation:
     Conversation API
     """
 
-    def __init__(self, cookiePath: str = "", cookies: Optional[dict] = None) -> None:
+    def __init__(self, cookiePath: str = "", cookies: Optional[dict] = None, proxy: Optional[str] = None) -> None:
         self.struct: dict = {
             "conversationId": None,
             "clientId": None,
             "conversationSignature": None,
             "result": {"value": "Success", "message": None},
         }
-        self.session = httpx.Client()
+        self.proxy = proxy
+        proxies = {
+            "all://": proxy,
+        }
+        self.session = httpx.Client(proxies=proxies)
         self.session.headers.update(
             {
                 "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/108.0.0.0 Safari/537.36",
@@ -222,7 +228,8 @@ class ChatHub:
     """
 
     def __init__(self, conversation: Conversation) -> None:
-        self.wss: Optional[websockets.WebSocketClientProtocol] = None
+        self.session: Optional[aiohttp.ClientSession] = None
+        self.wss: Optional[ClientWebSocketResponse] = None
         self.request: ChatHubRequest
         self.loop: bool
         self.task: asyncio.Task
@@ -231,6 +238,7 @@ def __init__(self, conversation: Conversation) -> None:
             client_id=conversation.struct["clientId"],
             conversation_id=conversation.struct["conversationId"],
         )
+        self.proxy = conversation.proxy
 
     async def ask_stream(
         self,
@@ -241,20 +249,23 @@ async def ask_stream(
         Ask a question to the bot
         """
         # Check if websocket is closed
-        if self.wss and self.wss.closed or not self.wss:
-            self.wss = await websockets.connect(
-                "wss://sydney.bing.com/sydney/ChatHub",
-                extra_headers=HEADERS,
-                max_size=None,
+        if not self.session:
+            self.session = aiohttp.ClientSession()
+        if not self.wss:
+            self.wss = await self.session.ws_connect(
+                url="wss://sydney.bing.com/sydney/ChatHub",
+                proxy=self.proxy
             )
             await self.__initial_handshake()
+
         # Construct a ChatHub request
         self.request.update(prompt=prompt, conversation_style=conversation_style)
         # Send request
-        await self.wss.send(append_identifier(self.request.struct))
+        await self.wss.send_str(append_identifier(self.request.struct))
         final = False
         while not final:
-            objects = str(await self.wss.recv()).split(DELIMITER)
+            msg = await self.wss.receive()
+            objects = msg.data.split(DELIMITER)
             for obj in objects:
                 if obj is None or obj == "":
                     continue
@@ -268,26 +279,29 @@ async def ask_stream(
                     yield True, response
 
     async def __initial_handshake(self):
-        await self.wss.send(append_identifier({"protocol": "json", "version": 1}))
-        await self.wss.recv()
+        await self.wss.send_str(append_identifier({"protocol": "json", "version": 1}))
+        await self.wss.receive()
 
     async def close(self):
         """
         Close the connection
         """
         if self.wss and not self.wss.closed:
             await self.wss.close()
+        if self.session and not self.session:
+            await self.session.close()
 
 
 class Chatbot:
     """
     Combines everything to make it seamless
     """
 
-    def __init__(self, cookiePath: str = "", cookies: Optional[dict] = None) -> None:
+    def __init__(self, cookiePath: str = "", cookies: Optional[dict] = None, proxy: Optional[str] = None) -> None:
         self.cookiePath: str = cookiePath
         self.cookies: Optional[dict] = cookies
-        self.chat_hub: ChatHub = ChatHub(Conversation(self.cookiePath, self.cookies))
+        self.chat_hub: ChatHub = ChatHub(Conversation(self.cookiePath, self.cookies, proxy))
+        self.proxy = proxy
 
     async def ask(
         self,
@@ -303,7 +317,7 @@ async def ask(
         ):
             if final:
                 return response
-        self.chat_hub.wss.close()
+        await self.chat_hub.wss.close()
 
     async def ask_stream(
         self,
@@ -330,7 +344,7 @@ async def reset(self):
         Reset the conversation
         """
         await self.close()
-        self.chat_hub = ChatHub(Conversation(self.cookiePath, self.cookies))
+        self.chat_hub = ChatHub(Conversation(self.cookiePath, self.cookies, self.proxy))
 
 
 def get_input(prompt):
@@ -362,12 +376,12 @@ def get_input(prompt):
     return user_input
 
 
-async def main():
+async def main(args: argparse.Namespace):
     """
     Main function
     """
     print("Initializing...")
-    bot = Chatbot()
+    bot = Chatbot(proxy=args.proxy)
     while True:
         prompt = get_input("\nYou:\n")
         if prompt == "!exit":
@@ -427,6 +441,11 @@ async def main():
         choices=["creative", "balanced", "precise"],
         default="balanced",
     )
+    parser.add_argument(
+        "--proxy",
+        type=str,
+        required=False
+    )
     parser.add_argument(
         "--cookie-file",
         type=str,
@@ -435,5 +454,4 @@ async def main():
     )
     args = parser.parse_args()
     os.environ["COOKIE_FILE"] = args.cookie_file
-    args = parser.parse_args()
-    asyncio.run(main())
+    asyncio.run(main(args))