Hansimov commited on
Commit
b4870cf
1 Parent(s): a2e9c80

:recycle: [Refactor] Encapsulate wss_send, and remove underscores of private functions

Browse files
chathub_request_constructor.py CHANGED
@@ -16,13 +16,13 @@ class ChathubRequestConstructor:
16
  self.conversation_id = conversation_id
17
  self.invocation_id = invocation_id
18
  self.conversation_style = conversation_style
19
- self.message_id = self._generate_random_uuid()
20
  self.construct()
21
 
22
- def _generate_random_uuid(self):
23
  return str(uuid.uuid4())
24
 
25
- def _generate_random_hex_str(self, length: int = 32) -> str:
26
  return "".join(random.choice("0123456789abcdef") for _ in range(length))
27
 
28
  def construct(self):
@@ -97,7 +97,7 @@ class ChathubRequestConstructor:
97
  "plugins": [
98
  {"id": "c310c353-b9f0-4d76-ab0d-1dd5e979cf68"},
99
  ],
100
- "traceId": self._generate_random_hex_str(),
101
  "conversationHistoryOptionsSets": [
102
  "autosave",
103
  "savemem",
 
16
  self.conversation_id = conversation_id
17
  self.invocation_id = invocation_id
18
  self.conversation_style = conversation_style
19
+ self.message_id = self.generate_random_uuid()
20
  self.construct()
21
 
22
+ def generate_random_uuid(self):
23
  return str(uuid.uuid4())
24
 
25
+ def generate_random_hex_str(self, length: int = 32) -> str:
26
  return "".join(random.choice("0123456789abcdef") for _ in range(length))
27
 
28
  def construct(self):
 
97
  "plugins": [
98
  {"id": "c310c353-b9f0-4d76-ab0d-1dd5e979cf68"},
99
  ],
100
+ "traceId": self.generate_random_hex_str(),
101
  "conversationHistoryOptionsSets": [
102
  "autosave",
103
  "savemem",
conversation_connector.py CHANGED
@@ -13,10 +13,6 @@ from logger.logger import logger
13
  http_proxy = "http://localhost:11111" # Replace with yours
14
 
15
 
16
- def serialize_websocket_message(msg: dict) -> str:
17
- return json.dumps(msg, ensure_ascii=False) + "\x1e"
18
-
19
-
20
  class ConversationConnector:
21
  def __init__(
22
  self,
@@ -36,12 +32,14 @@ class ConversationConnector:
36
  + f"?sec_access_token={urllib.parse.quote(self.sec_access_token)}"
37
  )
38
 
39
- async def _init_handshake(self):
40
- await self.wss.send_str(
41
- serialize_websocket_message({"protocol": "json", "version": 1})
42
- )
 
 
43
  await self.wss.receive_str()
44
- await self.wss.send_str(serialize_websocket_message({"type": 6}))
45
 
46
  async def stream_chat(self, prompt=""):
47
  self.aiohttp_session = aiohttp.ClientSession(cookies=self.cookies)
@@ -64,7 +62,7 @@ class ConversationConnector:
64
  proxy=http_proxy,
65
  )
66
 
67
- await self._init_handshake()
68
  chathub_request_constructor = ChathubRequestConstructor(
69
  prompt=prompt,
70
  conversation_style="precise",
@@ -74,9 +72,7 @@ class ConversationConnector:
74
  )
75
  chathub_request_constructor.construct()
76
 
77
- await self.wss.send_str(
78
- serialize_websocket_message(chathub_request_constructor.request_payload)
79
- )
80
 
81
  delta_content_pointer = 0
82
  while not self.wss.closed:
@@ -105,7 +101,7 @@ class ConversationConnector:
105
  if message_type is None:
106
  content = message["adaptiveCards"][0]["body"][0]["text"]
107
  delta_content = content[delta_content_pointer:]
108
- logger.mesg(delta_content, end="")
109
  delta_content_pointer = len(content)
110
  # Message: Suggested Questions
111
  if message.get("suggestedResponses"):
 
13
  http_proxy = "http://localhost:11111" # Replace with yours
14
 
15
 
 
 
 
 
16
  class ConversationConnector:
17
  def __init__(
18
  self,
 
32
  + f"?sec_access_token={urllib.parse.quote(self.sec_access_token)}"
33
  )
34
 
35
+ async def wss_send(self, message):
36
+ serialized_websocket_message = json.dumps(message, ensure_ascii=False) + "\x1e"
37
+ await self.wss.send_str(serialized_websocket_message)
38
+
39
+ async def init_handshake(self):
40
+ await self.wss_send({"protocol": "json", "version": 1})
41
  await self.wss.receive_str()
42
+ await self.wss_send({"type": 6})
43
 
44
  async def stream_chat(self, prompt=""):
45
  self.aiohttp_session = aiohttp.ClientSession(cookies=self.cookies)
 
62
  proxy=http_proxy,
63
  )
64
 
65
+ await self.init_handshake()
66
  chathub_request_constructor = ChathubRequestConstructor(
67
  prompt=prompt,
68
  conversation_style="precise",
 
72
  )
73
  chathub_request_constructor.construct()
74
 
75
+ await self.wss_send(chathub_request_constructor.request_payload)
 
 
76
 
77
  delta_content_pointer = 0
78
  while not self.wss.closed:
 
101
  if message_type is None:
102
  content = message["adaptiveCards"][0]["body"][0]["text"]
103
  delta_content = content[delta_content_pointer:]
104
+ logger.line(delta_content, end="")
105
  delta_content_pointer = len(content)
106
  # Message: Suggested Questions
107
  if message.get("suggestedResponses"):