Hansimov commited on
Commit
3a1b095
1 Parent(s): 5095a2b

:recycle: [Refactor] Separate init_wss_connection() and send_chathub_request()

Browse files
Files changed (1) hide show
  1. 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 construct_chathub_request_payload(self, prompt):
 
 
 
 
 
 
 
 
 
 
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.aiohttp_session = aiohttp.ClientSession(cookies=self.cookies)
76
- request_headers_constructor = ConversationConnectRequestHeadersConstructor()
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: