yym68686 commited on
Commit
27b2cca
·
1 Parent(s): dd27448

🐛 Bug: Fix the bug where FastAPI internally used different types of streaming responses.

Browse files
Files changed (1) hide show
  1. main.py +7 -4
main.py CHANGED
@@ -9,7 +9,9 @@ from starlette.middleware.base import BaseHTTPMiddleware
9
 
10
  from fastapi.middleware.cors import CORSMiddleware
11
  from fastapi import FastAPI, HTTPException, Depends, Request
12
- from fastapi.responses import StreamingResponse, JSONResponse
 
 
13
  from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
14
  from fastapi.exceptions import RequestValidationError
15
 
@@ -390,7 +392,7 @@ class StatsMiddleware(BaseHTTPMiddleware):
390
  try:
391
  response = await call_next(request)
392
 
393
- if isinstance(response, StreamingResponse):
394
  response = LoggingStreamingResponse(
395
  content=response.body_iterator,
396
  status_code=response.status_code,
@@ -527,14 +529,15 @@ async def process_request(request: Union[RequestModel, ImageGenerationRequest, A
527
  model = provider['model'][request.model]
528
  generator = fetch_response_stream(app.state.client, url, headers, payload, engine, model)
529
  wrapped_generator, first_response_time = await error_handling_wrapper(generator)
530
- response = StreamingResponse(wrapped_generator, media_type="text/event-stream")
531
  else:
532
  generator = fetch_response(app.state.client, url, headers, payload)
533
  wrapped_generator, first_response_time = await error_handling_wrapper(generator)
534
  first_element = await anext(wrapped_generator)
535
  first_element = first_element.lstrip("data: ")
536
  first_element = json.loads(first_element)
537
- response = JSONResponse(first_element)
 
538
 
539
  # 更新成功计数和首次响应时间
540
  await app.middleware_stack.app.update_channel_stats(current_info["request_id"], provider['provider'], request.model, token, success=True)
 
9
 
10
  from fastapi.middleware.cors import CORSMiddleware
11
  from fastapi import FastAPI, HTTPException, Depends, Request
12
+ from fastapi.responses import JSONResponse
13
+ from fastapi.responses import StreamingResponse as FastAPIStreamingResponse
14
+ from starlette.responses import StreamingResponse as StarletteStreamingResponse
15
  from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
16
  from fastapi.exceptions import RequestValidationError
17
 
 
392
  try:
393
  response = await call_next(request)
394
 
395
+ if isinstance(response, (FastAPIStreamingResponse, StarletteStreamingResponse)):
396
  response = LoggingStreamingResponse(
397
  content=response.body_iterator,
398
  status_code=response.status_code,
 
529
  model = provider['model'][request.model]
530
  generator = fetch_response_stream(app.state.client, url, headers, payload, engine, model)
531
  wrapped_generator, first_response_time = await error_handling_wrapper(generator)
532
+ response = StarletteStreamingResponse(wrapped_generator, media_type="text/event-stream")
533
  else:
534
  generator = fetch_response(app.state.client, url, headers, payload)
535
  wrapped_generator, first_response_time = await error_handling_wrapper(generator)
536
  first_element = await anext(wrapped_generator)
537
  first_element = first_element.lstrip("data: ")
538
  first_element = json.loads(first_element)
539
+ response = StarletteStreamingResponse(iter([json.dumps(first_element)]), media_type="application/json")
540
+ # response = JSONResponse(first_element)
541
 
542
  # 更新成功计数和首次响应时间
543
  await app.middleware_stack.app.update_channel_stats(current_info["request_id"], provider['provider'], request.model, token, success=True)