Use async model streaming for chat
This commit is contained in:
@@ -1,7 +1,8 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from collections.abc import Iterator
|
||||
from collections.abc import AsyncIterator
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
from fastapi.responses import StreamingResponse
|
||||
@@ -87,7 +88,7 @@ def stop(
|
||||
return api_success()
|
||||
|
||||
|
||||
def _chat_stream(payload: ChatCompletionRequest, db: Session, current_user: User) -> Iterator[str]:
|
||||
async def _chat_stream(payload: ChatCompletionRequest, db: Session, current_user: User) -> AsyncIterator[str]:
|
||||
config = load_chat_queue_config(db)
|
||||
queue_request = chat_queue_manager.request_slot(config)
|
||||
acquired = queue_request.status == "acquired"
|
||||
@@ -111,7 +112,7 @@ def _chat_stream(payload: ChatCompletionRequest, db: Session, current_user: User
|
||||
activeCount=queue_request.active_count,
|
||||
waitingCount=queue_request.waiting_count,
|
||||
)
|
||||
queue_request = chat_queue_manager.wait_for_slot(queued_token, config)
|
||||
queue_request = await asyncio.to_thread(chat_queue_manager.wait_for_slot, queued_token, config)
|
||||
acquired = queue_request.status == "acquired"
|
||||
if queue_request.status == "timeout":
|
||||
yield _sse_event("error", message="排队等待超时,请稍后再试。")
|
||||
@@ -125,10 +126,12 @@ def _chat_stream(payload: ChatCompletionRequest, db: Session, current_user: User
|
||||
activeCount=queue_request.active_count,
|
||||
waitingCount=queue_request.waiting_count,
|
||||
)
|
||||
for chunk in ChatStreamService.stream_answer(db, current_user, payload.sessionId, payload.message):
|
||||
async for chunk in ChatStreamService.stream_answer_async(db, current_user, payload.sessionId, payload.message):
|
||||
yield _sse_event("content", content=chunk)
|
||||
except HTTPException as exc:
|
||||
yield _sse_event("error", message=str(exc.detail))
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception:
|
||||
yield _sse_event("error", message="AI 回复生成失败,请稍后再试。")
|
||||
finally:
|
||||
|
||||
Reference in New Issue
Block a user