yym68686 commited on
Commit
ea47d28
·
1 Parent(s): d7169b0

🐛 Bug: Fix the bug where there is no error handling for response JSON parsing errors.

Browse files
Files changed (1) hide show
  1. main.py +13 -12
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
- resp: dict = json.loads(line)
272
- input_tokens = safe_get(resp, "message", "usage", "input_tokens", default=0)
273
- input_tokens = safe_get(resp, "usage", "prompt_tokens", default=0)
274
- output_tokens = safe_get(resp, "usage", "completion_tokens", default=0)
275
- total_tokens = input_tokens + output_tokens
276
-
277
- model = self.current_info.get("model", "")
278
- # total_cost = calculate_cost(model, input_tokens, output_tokens)
279
- self.current_info["prompt_tokens"] = input_tokens
280
- self.current_info["completion_tokens"] = output_tokens
281
- self.current_info["total_tokens"] = total_tokens
282
- # self.current_info["cost"] = total_cost
 
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