bing-chat-api / conversations /conversation_connector.py
Hansimov's picture
:recycle: [Refactor] Remove unused imports
82b42e7
raw
history blame
5.03 kB
import aiohttp
import json
import urllib
from networks import (
ChathubRequestPayloadConstructor,
ConversationRequestHeadersConstructor,
MessageParser,
OpenaiStreamOutputer,
)
from utils.logger import logger
http_proxy = "http://localhost:11111" # Replace with yours
class ConversationConnector:
"""
Input params:
- `sec_access_token`, `client_id`, `conversation_id`
- Generated by `ConversationCreator`
- `invocation_id` (int):
- For 1st request, this value must be `0`.
- For all requests after, any integer is valid.
- To make it simple, use `1` for all requests after the 1st one.
"""
def __init__(
self,
conversation_style: str = "precise",
sec_access_token: str = "",
client_id: str = "",
conversation_id: str = "",
invocation_id: int = 0,
cookies={},
):
self.conversation_style = conversation_style
self.sec_access_token = sec_access_token
self.client_id = client_id
self.conversation_id = conversation_id
self.invocation_id = invocation_id
self.cookies = cookies
async def wss_send(self, message):
serialized_websocket_message = json.dumps(message, ensure_ascii=False) + "\x1e"
await self.wss.send_str(serialized_websocket_message)
async def init_handshake(self):
await self.wss_send({"protocol": "json", "version": 1})
await self.wss.receive_str()
await self.wss_send({"type": 6})
async def connect(self):
self.quotelized_sec_access_token = urllib.parse.quote(self.sec_access_token)
self.ws_url = (
f"wss://sydney.bing.com/sydney/ChatHub"
f"?sec_access_token={self.quotelized_sec_access_token}"
)
self.aiohttp_session = aiohttp.ClientSession(cookies=self.cookies)
headers_constructor = ConversationRequestHeadersConstructor()
self.wss = await self.aiohttp_session.ws_connect(
self.ws_url,
headers=headers_constructor.request_headers,
proxy=http_proxy,
)
await self.init_handshake()
async def send_chathub_request(self, prompt):
payload_constructor = ChathubRequestPayloadConstructor(
prompt=prompt,
conversation_style=self.conversation_style,
client_id=self.client_id,
conversation_id=self.conversation_id,
invocation_id=self.invocation_id,
)
self.connect_request_payload = payload_constructor.request_payload
await self.wss_send(self.connect_request_payload)
async def stream_chat(self, prompt="", yield_output=False):
await self.connect()
await self.send_chathub_request(prompt)
message_parser = MessageParser(outputer=OpenaiStreamOutputer())
has_output_role_message = False
if yield_output and not has_output_role_message:
has_output_role_message = True
yield message_parser.outputer.output(content="", content_type="Role")
while not self.wss.closed:
response_lines_str = await self.wss.receive_str()
if isinstance(response_lines_str, str):
response_lines = response_lines_str.split("\x1e")
else:
continue
for line in response_lines:
if not line:
continue
data = json.loads(line)
# Stream: Meaningful Messages
if data.get("type") == 1:
if yield_output:
output = message_parser.parse(data, return_output=True)
if isinstance(output, list):
for item in output:
yield item
else:
yield output
else:
message_parser.parse(data)
# Stream: List of all messages in the whole conversation
elif data.get("type") == 2:
if data.get("item"):
# item = data.get("item")
# logger.note("\n[Saving chat messages ...]")
pass
# Stream: End of Conversation
elif data.get("type") == 3:
finished_str = "\n[Finished]"
logger.success(finished_str)
self.invocation_id += 1
await self.wss.close()
await self.aiohttp_session.close()
if yield_output:
yield message_parser.outputer.output(
content=finished_str, content_type="Finished"
)
break
# Stream: Heartbeat Signal
elif data.get("type") == 6:
continue
# Stream: Not Implemented
else:
continue