Spaces:
Paused
Paused
:recycle: [Refactor] Separate init_wss_connection() and send_chathub_request()
Browse files- conversation_connector.py +14 -12
conversation_connector.py
CHANGED
@@ -61,7 +61,17 @@ class ConversationConnector:
|
|
61 |
await self.wss.receive_str()
|
62 |
await self.wss_send({"type": 6})
|
63 |
|
64 |
-
def
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
65 |
chathub_request_constructor = ChathubRequestConstructor(
|
66 |
prompt=prompt,
|
67 |
conversation_style="precise",
|
@@ -70,19 +80,11 @@ class ConversationConnector:
|
|
70 |
invocation_id=self.invocation_id,
|
71 |
)
|
72 |
self.connect_request_payload = chathub_request_constructor.request_payload
|
|
|
73 |
|
74 |
async def stream_chat(self, prompt=""):
|
75 |
-
self.
|
76 |
-
|
77 |
-
self.wss = await self.aiohttp_session.ws_connect(
|
78 |
-
self.ws_url,
|
79 |
-
headers=request_headers_constructor.request_headers,
|
80 |
-
proxy=http_proxy,
|
81 |
-
)
|
82 |
-
|
83 |
-
await self.init_handshake()
|
84 |
-
self.construct_chathub_request_payload(prompt)
|
85 |
-
await self.wss_send(self.connect_request_payload)
|
86 |
|
87 |
delta_content_pointer = 0
|
88 |
while not self.wss.closed:
|
|
|
61 |
await self.wss.receive_str()
|
62 |
await self.wss_send({"type": 6})
|
63 |
|
64 |
+
async def init_wss_connection(self):
|
65 |
+
self.aiohttp_session = aiohttp.ClientSession(cookies=self.cookies)
|
66 |
+
request_headers_constructor = ConversationConnectRequestHeadersConstructor()
|
67 |
+
self.wss = await self.aiohttp_session.ws_connect(
|
68 |
+
self.ws_url,
|
69 |
+
headers=request_headers_constructor.request_headers,
|
70 |
+
proxy=http_proxy,
|
71 |
+
)
|
72 |
+
await self.init_handshake()
|
73 |
+
|
74 |
+
async def send_chathub_request(self, prompt):
|
75 |
chathub_request_constructor = ChathubRequestConstructor(
|
76 |
prompt=prompt,
|
77 |
conversation_style="precise",
|
|
|
80 |
invocation_id=self.invocation_id,
|
81 |
)
|
82 |
self.connect_request_payload = chathub_request_constructor.request_payload
|
83 |
+
await self.wss_send(self.connect_request_payload)
|
84 |
|
85 |
async def stream_chat(self, prompt=""):
|
86 |
+
await self.init_wss_connection()
|
87 |
+
await self.send_chathub_request(prompt)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
88 |
|
89 |
delta_content_pointer = 0
|
90 |
while not self.wss.closed:
|