yym68686 commited on
Commit
55a4247
·
1 Parent(s): 809165c

🐛 Bug: 1. Fix the bug of httpx reading streamed output timeout.

Browse files
Files changed (2) hide show
  1. main.py +2 -2
  2. utils.py +58 -12
main.py CHANGED
@@ -20,7 +20,7 @@ from urllib.parse import urlparse
20
  @asynccontextmanager
21
  async def lifespan(app: FastAPI):
22
  # 启动时的代码
23
- timeout = httpx.Timeout(connect=15.0, read=10.0, write=30.0, pool=30.0)
24
  app.state.client = httpx.AsyncClient(timeout=timeout)
25
  app.state.config, app.state.api_keys_db, app.state.api_list = await load_config(app)
26
  yield
@@ -74,7 +74,7 @@ async def process_request(request: RequestModel, provider: Dict):
74
  if request.stream:
75
  model = provider['model'][request.model]
76
  generator = fetch_response_stream(app.state.client, url, headers, payload, engine, model)
77
- wrapped_generator = await error_handling_wrapper(generator, status_code=500)
78
  return StreamingResponse(wrapped_generator, media_type="text/event-stream")
79
  else:
80
  return await fetch_response(app.state.client, url, headers, payload)
 
20
  @asynccontextmanager
21
  async def lifespan(app: FastAPI):
22
  # 启动时的代码
23
+ timeout = httpx.Timeout(connect=15.0, read=20.0, write=30.0, pool=30.0)
24
  app.state.client = httpx.AsyncClient(timeout=timeout)
25
  app.state.config, app.state.api_keys_db, app.state.api_list = await load_config(app)
26
  yield
 
74
  if request.stream:
75
  model = provider['model'][request.model]
76
  generator = fetch_response_stream(app.state.client, url, headers, payload, engine, model)
77
+ wrapped_generator = error_handling_wrapper(generator, status_code=500)
78
  return StreamingResponse(wrapped_generator, media_type="text/event-stream")
79
  else:
80
  return await fetch_response(app.state.client, url, headers, payload)
utils.py CHANGED
@@ -73,11 +73,16 @@ def ensure_string(item):
73
  else:
74
  return str(item)
75
 
76
- async def async_generator(items):
77
- for item in items:
78
- yield item
79
-
80
  async def error_handling_wrapper(generator, status_code=200):
 
 
 
 
 
 
 
 
 
81
  try:
82
  first_item = await generator.__anext__()
83
  first_item_str = first_item
@@ -97,19 +102,60 @@ async def error_handling_wrapper(generator, status_code=200):
97
  logger.error("error_handling_wrapper JSONDecodeError!" + repr(first_item_str))
98
  raise StopAsyncIteration
99
  if isinstance(first_item_str, dict) and 'error' in first_item_str:
100
- # 如果第一个 yield 的项是错误信息,抛出 HTTPException
101
  raise HTTPException(status_code=status_code, detail=f"{first_item_str}"[:300])
102
 
103
- # 如果不是错误,创建一个新的生成器,首先yield第一个项,然后yield剩余的项
104
- async def new_generator():
105
- yield ensure_string(first_item)
106
- async for item in generator:
107
- yield ensure_string(item)
108
-
109
- return new_generator()
 
 
 
110
 
111
  except StopAsyncIteration:
112
  raise HTTPException(status_code=status_code, detail="data: {'error': 'No data returned'}")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
113
 
114
  def post_all_models(token, config, api_list):
115
  all_models = []
 
73
  else:
74
  return str(item)
75
 
 
 
 
 
76
  async def error_handling_wrapper(generator, status_code=200):
77
+ async def new_generator():
78
+ try:
79
+ yield ensure_string(first_item)
80
+ async for item in generator:
81
+ yield ensure_string(item)
82
+ except Exception as e:
83
+ logger.exception(f"Error in new_generator: {e}")
84
+ raise HTTPException(status_code=status_code, detail=f"Stream error: {str(e)}")
85
+
86
  try:
87
  first_item = await generator.__anext__()
88
  first_item_str = first_item
 
102
  logger.error("error_handling_wrapper JSONDecodeError!" + repr(first_item_str))
103
  raise StopAsyncIteration
104
  if isinstance(first_item_str, dict) and 'error' in first_item_str:
 
105
  raise HTTPException(status_code=status_code, detail=f"{first_item_str}"[:300])
106
 
107
+ # 创建新的生成器并包装在 try-except 块中
108
+ wrapped_generator = new_generator()
109
+ try:
110
+ async for item in wrapped_generator:
111
+ yield item
112
+ except HTTPException as http_exc:
113
+ raise http_exc
114
+ except Exception as e:
115
+ logger.exception(f"Unexpected error in error_handling_wrapper: {e}")
116
+ raise HTTPException(status_code=status_code, detail=f"Unexpected error: {str(e)}")
117
 
118
  except StopAsyncIteration:
119
  raise HTTPException(status_code=status_code, detail="data: {'error': 'No data returned'}")
120
+ except HTTPException as http_exc:
121
+ raise http_exc
122
+ except Exception as e:
123
+ logger.exception(f"Error in error_handling_wrapper: {e}")
124
+ raise HTTPException(status_code=status_code, detail=f"Wrapper error: {str(e)}")
125
+
126
+ # async def error_handling_wrapper(generator, status_code=200):
127
+ # try:
128
+ # first_item = await generator.__anext__()
129
+ # first_item_str = first_item
130
+ # if isinstance(first_item_str, (bytes, bytearray)):
131
+ # first_item_str = first_item_str.decode("utf-8")
132
+ # if isinstance(first_item_str, str):
133
+ # if first_item_str.startswith("data: "):
134
+ # first_item_str = first_item_str[6:]
135
+ # elif first_item_str.startswith("data:"):
136
+ # first_item_str = first_item_str[5:]
137
+ # if first_item_str.startswith("[DONE]"):
138
+ # logger.error("error_handling_wrapper [DONE]!")
139
+ # raise StopAsyncIteration
140
+ # try:
141
+ # first_item_str = json.loads(first_item_str)
142
+ # except json.JSONDecodeError:
143
+ # logger.error("error_handling_wrapper JSONDecodeError!" + repr(first_item_str))
144
+ # raise StopAsyncIteration
145
+ # if isinstance(first_item_str, dict) and 'error' in first_item_str:
146
+ # # 如果第一个 yield 的项是错误信息,抛出 HTTPException
147
+ # raise HTTPException(status_code=status_code, detail=f"{first_item_str}"[:300])
148
+
149
+ # # 如果不是错误,创建一个新的生成器,首先yield第一个项,然后yield剩余的项
150
+ # async def new_generator():
151
+ # yield ensure_string(first_item)
152
+ # async for item in generator:
153
+ # yield ensure_string(item)
154
+
155
+ # return new_generator()
156
+
157
+ # except StopAsyncIteration:
158
+ # raise HTTPException(status_code=status_code, detail="data: {'error': 'No data returned'}")
159
 
160
  def post_all_models(token, config, api_list):
161
  all_models = []