Skip to content

Commit

Permalink
🐛 Bug: 1. Fix the bug of httpx reading streamed output timeout.
Browse files Browse the repository at this point in the history
2. Increase read response timeout.
  • Loading branch information
yym68686 committed Aug 25, 2024
1 parent 809165c commit 55a4247
Show file tree
Hide file tree
Showing 2 changed files with 60 additions and 14 deletions.
4 changes: 2 additions & 2 deletions main.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
@asynccontextmanager
async def lifespan(app: FastAPI):
# 启动时的代码
timeout = httpx.Timeout(connect=15.0, read=10.0, write=30.0, pool=30.0)
timeout = httpx.Timeout(connect=15.0, read=20.0, write=30.0, pool=30.0)
app.state.client = httpx.AsyncClient(timeout=timeout)
app.state.config, app.state.api_keys_db, app.state.api_list = await load_config(app)
yield
Expand Down Expand Up @@ -74,7 +74,7 @@ async def process_request(request: RequestModel, provider: Dict):
if request.stream:
model = provider['model'][request.model]
generator = fetch_response_stream(app.state.client, url, headers, payload, engine, model)
wrapped_generator = await error_handling_wrapper(generator, status_code=500)
wrapped_generator = error_handling_wrapper(generator, status_code=500)
return StreamingResponse(wrapped_generator, media_type="text/event-stream")
else:
return await fetch_response(app.state.client, url, headers, payload)
Expand Down
70 changes: 58 additions & 12 deletions utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -73,11 +73,16 @@ def ensure_string(item):
else:
return str(item)

async def async_generator(items):
for item in items:
yield item

async def error_handling_wrapper(generator, status_code=200):
async def new_generator():
try:
yield ensure_string(first_item)
async for item in generator:
yield ensure_string(item)
except Exception as e:
logger.exception(f"Error in new_generator: {e}")
raise HTTPException(status_code=status_code, detail=f"Stream error: {str(e)}")

try:
first_item = await generator.__anext__()
first_item_str = first_item
Expand All @@ -97,19 +102,60 @@ async def error_handling_wrapper(generator, status_code=200):
logger.error("error_handling_wrapper JSONDecodeError!" + repr(first_item_str))
raise StopAsyncIteration
if isinstance(first_item_str, dict) and 'error' in first_item_str:
# 如果第一个 yield 的项是错误信息,抛出 HTTPException
raise HTTPException(status_code=status_code, detail=f"{first_item_str}"[:300])

# 如果不是错误,创建一个新的生成器,首先yield第一个项,然后yield剩余的项
async def new_generator():
yield ensure_string(first_item)
async for item in generator:
yield ensure_string(item)

return new_generator()
# 创建新的生成器并包装在 try-except 块中
wrapped_generator = new_generator()
try:
async for item in wrapped_generator:
yield item
except HTTPException as http_exc:
raise http_exc
except Exception as e:
logger.exception(f"Unexpected error in error_handling_wrapper: {e}")
raise HTTPException(status_code=status_code, detail=f"Unexpected error: {str(e)}")

except StopAsyncIteration:
raise HTTPException(status_code=status_code, detail="data: {'error': 'No data returned'}")
except HTTPException as http_exc:
raise http_exc
except Exception as e:
logger.exception(f"Error in error_handling_wrapper: {e}")
raise HTTPException(status_code=status_code, detail=f"Wrapper error: {str(e)}")

# async def error_handling_wrapper(generator, status_code=200):
# try:
# first_item = await generator.__anext__()
# first_item_str = first_item
# if isinstance(first_item_str, (bytes, bytearray)):
# first_item_str = first_item_str.decode("utf-8")
# if isinstance(first_item_str, str):
# if first_item_str.startswith("data: "):
# first_item_str = first_item_str[6:]
# elif first_item_str.startswith("data:"):
# first_item_str = first_item_str[5:]
# if first_item_str.startswith("[DONE]"):
# logger.error("error_handling_wrapper [DONE]!")
# raise StopAsyncIteration
# try:
# first_item_str = json.loads(first_item_str)
# except json.JSONDecodeError:
# logger.error("error_handling_wrapper JSONDecodeError!" + repr(first_item_str))
# raise StopAsyncIteration
# if isinstance(first_item_str, dict) and 'error' in first_item_str:
# # 如果第一个 yield 的项是错误信息,抛出 HTTPException
# raise HTTPException(status_code=status_code, detail=f"{first_item_str}"[:300])

# # 如果不是错误,创建一个新的生成器,首先yield第一个项,然后yield剩余的项
# async def new_generator():
# yield ensure_string(first_item)
# async for item in generator:
# yield ensure_string(item)

# return new_generator()

# except StopAsyncIteration:
# raise HTTPException(status_code=status_code, detail="data: {'error': 'No data returned'}")

def post_all_models(token, config, api_list):
all_models = []
Expand Down

0 comments on commit 55a4247

Please sign in to comment.