yym68686 commited on
Commit
bdfeca1
·
1 Parent(s): 983694e

🐛 Bug: Fix the bug where the next channel cannot be automatically switched when an error occurs in non-streaming output for o1-mini and o1-preview.

Browse files
Files changed (2) hide show
  1. main.py +6 -1
  2. utils.py +2 -4
main.py CHANGED
@@ -241,7 +241,12 @@ async def process_request(request: Union[RequestModel, ImageGenerationRequest],
241
  wrapped_generator = await error_handling_wrapper(generator)
242
  response = StreamingResponse(wrapped_generator, media_type="text/event-stream")
243
  else:
244
- response = await anext(fetch_response(app.state.client, url, headers, payload))
 
 
 
 
 
245
 
246
  # 更新成功计数
247
  async with app.middleware_stack.app.lock:
 
241
  wrapped_generator = await error_handling_wrapper(generator)
242
  response = StreamingResponse(wrapped_generator, media_type="text/event-stream")
243
  else:
244
+ generator = fetch_response(app.state.client, url, headers, payload)
245
+ wrapped_generator = await error_handling_wrapper(generator)
246
+ first_element = await anext(wrapped_generator)
247
+ first_element = first_element.lstrip("data: ")
248
+ first_element = json.loads(first_element)
249
+ response = JSONResponse(first_element)
250
 
251
  # 更新成功计数
252
  async with app.middleware_stack.app.lock:
utils.py CHANGED
@@ -121,10 +121,8 @@ async def error_handling_wrapper(generator):
121
  if isinstance(first_item_str, (bytes, bytearray)):
122
  first_item_str = first_item_str.decode("utf-8")
123
  if isinstance(first_item_str, str):
124
- if first_item_str.startswith("data: "):
125
- first_item_str = first_item_str[6:]
126
- elif first_item_str.startswith("data:"):
127
- first_item_str = first_item_str[5:]
128
  if first_item_str.startswith("[DONE]"):
129
  logger.error("error_handling_wrapper [DONE]!")
130
  raise StopAsyncIteration
 
121
  if isinstance(first_item_str, (bytes, bytearray)):
122
  first_item_str = first_item_str.decode("utf-8")
123
  if isinstance(first_item_str, str):
124
+ if first_item_str.startswith("data:"):
125
+ first_item_str = first_item_str.lstrip("data: ")
 
 
126
  if first_item_str.startswith("[DONE]"):
127
  logger.error("error_handling_wrapper [DONE]!")
128
  raise StopAsyncIteration