diff --git a/pyproject.toml b/pyproject.toml index d245307..c822ba9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ packages = [{ include = "sydney" }] [tool.poetry.dependencies] python = "^3.9" aiohttp = "^3.9.3" -websockets = "^12.0" +websocket-client = "^1.7.0" brotli = "^1.1.0" [tool.poetry.group.dev.dependencies] diff --git a/sydney/sydney.py b/sydney/sydney.py index 18b7ccf..3567b66 100644 --- a/sydney/sydney.py +++ b/sydney/sydney.py @@ -7,9 +7,8 @@ from typing import AsyncGenerator from urllib import parse -import websockets.client as websockets +from websocket import WebSocket from aiohttp import ClientSession, FormData, TCPConnector -from websockets.client import WebSocketClientProtocol from sydney.constants import ( BING_BLOB_URL, @@ -57,6 +56,9 @@ def __init__( style: str = "balanced", bing_cookies: str | None = None, use_proxy: bool = False, + http_proxy: str | None = None, + http_proxy_port: str | None = None, + http_proxy_auth: tuple | None = None ) -> None: """ Client for Copilot (formerly named Bing Chat), also known as Sydney. @@ -76,6 +78,9 @@ def __init__( """ self.bing_cookies = bing_cookies if bing_cookies else getenv("BING_COOKIES") self.use_proxy = use_proxy + self.http_proxy = http_proxy + self.http_proxy_port = http_proxy_port + self.http_proxy_auth = http_proxy_auth self.conversation_style: ConversationStyle = getattr(ConversationStyle, style.upper()) self.conversation_style_option_sets: ConversationStyleOptionSets = getattr( ConversationStyleOptionSets, style.upper() @@ -87,7 +92,7 @@ def __init__( self.invocation_id: int | None = None self.number_of_messages: int | None = None self.max_messages: int | None = None - self.wss_client: WebSocketClientProtocol | None = None + self.wss_client: WebSocket() | None = None self.session: ClientSession | None = None async def __aenter__(self) -> SydneyClient: @@ -320,16 +325,19 @@ async def _ask( if self.encrypted_conversation_signature: bing_chathub_url += f"?sec_access_token={parse.quote(self.encrypted_conversation_signature)}" + ws = WebSocket() # Create a websocket connection with Copilot for sending and receiving messages. try: - self.wss_client = await websockets.connect( - bing_chathub_url, extra_headers=CHATHUB_HEADERS, max_size=None + ws.connect( + bing_chathub_url, extra_headers=CHATHUB_HEADERS, max_size=None, http_proxy_host=self.http_proxy, http_proxy_port=self.http_proxy_port, http_proxy_auth=self.http_proxy_auth ) + except TimeoutError: - raise ConnectionTimeoutException("Failed to connect to Copilot, connection timed out") from None - await self.wss_client.send(as_json({"protocol": "json", "version": 1})) - await self.wss_client.recv() + raise ConnectionTimeoutException("Failed to connect to Copilot, connection timed out") from None + ws.send(as_json({"protocol": "json", "version": 1})) + ws.recv() + attachment_info = None if attachment: attachment_info = await self._upload_attachment(attachment) @@ -340,11 +348,11 @@ async def _ask( request = self._build_ask_arguments(prompt, search, attachment_info, context) self.invocation_id += 1 - await self.wss_client.send(as_json(request)) + ws.send(as_json(request)) streaming = True while streaming: - objects = str(await self.wss_client.recv()).split(DELIMETER) + objects = str(ws.recv()).split(DELIMETER) for obj in objects: if not obj: continue @@ -427,7 +435,7 @@ async def _ask( # Exit, type 2 is the last message. streaming = False - await self.wss_client.close() + ws.close() async def start_conversation(self) -> None: """