yym68686 commited on
Commit
e6f7928
·
1 Parent(s): 1ad99a8

Fix a possible bug that occurs when attempting to encode a dictionary (dict) object as a string in a FastAPI application.

Browse files
Files changed (1) hide show
  1. utils.py +22 -4
utils.py CHANGED
@@ -31,6 +31,20 @@ def load_config():
31
 
32
  config, api_keys_db, api_list = load_config()
33
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
34
  async def error_handling_wrapper(generator, status_code=200):
35
  try:
36
  first_item = await generator.__anext__()
@@ -42,21 +56,25 @@ async def error_handling_wrapper(generator, status_code=200):
42
  first_item_str = first_item_str[6:]
43
  elif first_item_str.startswith("data:"):
44
  first_item_str = first_item_str[5:]
45
- first_item_str = json.loads(first_item_str)
 
 
 
 
46
  if isinstance(first_item_str, dict) and 'error' in first_item_str:
47
  # 如果第一个 yield 的项是错误信息,抛出 HTTPException
48
  raise HTTPException(status_code=status_code, detail=f"{first_item_str}"[:300])
49
 
50
  # 如果不是错误,创建一个新的生成器,首先yield第一个项,然后yield剩余的项
51
  async def new_generator():
52
- yield first_item
53
  async for item in generator:
54
- yield item
55
 
56
  return new_generator()
57
  except StopAsyncIteration:
58
  # 处理生成器为空的情况
59
- return "data: {'error': 'No data returned'}\n\n"
60
 
61
  def post_all_models(token):
62
  all_models = []
 
31
 
32
  config, api_keys_db, api_list = load_config()
33
 
34
+ def ensure_string(item):
35
+ if isinstance(item, (bytes, bytearray)):
36
+ return item.decode("utf-8")
37
+ elif isinstance(item, str):
38
+ return item
39
+ elif isinstance(item, dict):
40
+ return f"data: {json.dumps(item)}\n\n"
41
+ else:
42
+ return str(item)
43
+
44
+ async def async_generator(items):
45
+ for item in items:
46
+ yield item
47
+
48
  async def error_handling_wrapper(generator, status_code=200):
49
  try:
50
  first_item = await generator.__anext__()
 
56
  first_item_str = first_item_str[6:]
57
  elif first_item_str.startswith("data:"):
58
  first_item_str = first_item_str[5:]
59
+ try:
60
+ first_item_str = json.loads(first_item_str)
61
+ except json.JSONDecodeError:
62
+ print("error_handling_wrapper JSONDecodeError!")
63
+ pass # 如果不是有效的JSON,保持原样
64
  if isinstance(first_item_str, dict) and 'error' in first_item_str:
65
  # 如果第一个 yield 的项是错误信息,抛出 HTTPException
66
  raise HTTPException(status_code=status_code, detail=f"{first_item_str}"[:300])
67
 
68
  # 如果不是错误,创建一个新的生成器,首先yield第一个项,然后yield剩余的项
69
  async def new_generator():
70
+ yield ensure_string(first_item)
71
  async for item in generator:
72
+ yield ensure_string(item)
73
 
74
  return new_generator()
75
  except StopAsyncIteration:
76
  # 处理生成器为空的情况
77
+ return async_generator(["data: {'error': 'No data returned'}\n\n"])
78
 
79
  def post_all_models(token):
80
  all_models = []