🐛 Bug: Fix the bug where there is no error handling for response JSON parsing errors.
Browse files
main.py
CHANGED
@@ -268,18 +268,19 @@ class LoggingStreamingResponse(Response):
|
|
268 |
if line.startswith("data:"):
|
269 |
line = line.lstrip("data: ")
|
270 |
if not line.startswith("[DONE]"):
|
271 |
-
|
272 |
-
|
273 |
-
|
274 |
-
|
275 |
-
|
276 |
-
|
277 |
-
|
278 |
-
|
279 |
-
|
280 |
-
|
281 |
-
|
282 |
-
|
|
|
283 |
yield chunk
|
284 |
except Exception as e:
|
285 |
raise
|
|
|
268 |
if line.startswith("data:"):
|
269 |
line = line.lstrip("data: ")
|
270 |
if not line.startswith("[DONE]"):
|
271 |
+
try:
|
272 |
+
resp: dict = json.loads(line)
|
273 |
+
input_tokens = safe_get(resp, "message", "usage", "input_tokens", default=0)
|
274 |
+
input_tokens = safe_get(resp, "usage", "prompt_tokens", default=0)
|
275 |
+
output_tokens = safe_get(resp, "usage", "completion_tokens", default=0)
|
276 |
+
total_tokens = input_tokens + output_tokens
|
277 |
+
|
278 |
+
self.current_info["prompt_tokens"] = input_tokens
|
279 |
+
self.current_info["completion_tokens"] = output_tokens
|
280 |
+
self.current_info["total_tokens"] = total_tokens
|
281 |
+
except Exception as e:
|
282 |
+
logger.error(f"Error parsing response: {str(e)}, line: {repr(line)}")
|
283 |
+
continue
|
284 |
yield chunk
|
285 |
except Exception as e:
|
286 |
raise
|