Hansimov commited on
Commit
c706328
1 Parent(s): 391cdfe

:gem: [Feature] HuggingchatStreamer: New log request and response

Browse files
Files changed (1) hide show
  1. networks/huggingchat_streamer.py +52 -0
networks/huggingchat_streamer.py CHANGED
@@ -94,6 +94,58 @@ class HuggingchatStreamer:
94
  return conversation_id
95
 
96
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
97
  def chat_response(
98
  self,
99
  prompt: str = None,
 
94
  return conversation_id
95
 
96
 
97
+ def log_request(self, url, method="GET"):
98
+ logger.note(f"> {method}:", end=" ")
99
+ logger.mesg(f"{url}", end=" ")
100
+
101
+ def log_response(
102
+ self, res: requests.Response, stream=False, iter_lines=False, verbose=False
103
+ ):
104
+ status_code = res.status_code
105
+ status_code_str = f"[{status_code}]"
106
+
107
+ if status_code == 200:
108
+ logger_func = logger.success
109
+ else:
110
+ logger_func = logger.warn
111
+
112
+ logger_func(status_code_str)
113
+
114
+ logger.enter_quiet(not verbose)
115
+
116
+ if status_code != 200:
117
+ logger_func(res.text)
118
+
119
+ if stream:
120
+ if not iter_lines:
121
+ return
122
+
123
+ for line in res.iter_lines():
124
+ line = line.decode("utf-8")
125
+ line = re.sub(r"^data:\s*", "", line)
126
+ line = line.strip()
127
+ if line:
128
+ try:
129
+ data = json.loads(line, strict=False)
130
+ msg_type = data.get("type")
131
+ if msg_type == "status":
132
+ msg_status = data.get("status")
133
+ elif msg_type == "stream":
134
+ content = data.get("token", "")
135
+ logger_func(content, end="")
136
+ elif msg_type == "finalAnswer":
137
+ full_content = data.get("text")
138
+ logger.success("\n[Finished]")
139
+ break
140
+ else:
141
+ pass
142
+ except Exception as e:
143
+ logger.warn(e)
144
+ else:
145
+ logger_func(res.json())
146
+
147
+ logger.exit_quiet(not verbose)
148
+
149
  def chat_response(
150
  self,
151
  prompt: str = None,