151 lines
5.2 KiB
Python
151 lines
5.2 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
from collections.abc import AsyncIterator
|
|
|
|
from fastapi import APIRouter, Depends, HTTPException, Query
|
|
from fastapi.responses import StreamingResponse
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.core.database import get_db
|
|
from app.core.dependencies import get_current_user
|
|
from app.core.responses import api_success
|
|
from app.models.user import User
|
|
from app.schemas.chat import (
|
|
ChatCompletionRequest,
|
|
ChatMessageRead,
|
|
ChatSessionRead,
|
|
CreateSessionResponse,
|
|
StopChatRequest,
|
|
UpdateSessionTitleRequest,
|
|
)
|
|
from app.services.chat_service import ChatService
|
|
from app.services.chat_queue_service import chat_queue_manager, load_chat_queue_config
|
|
from app.services.chat_stream_service import ChatStreamService
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
@router.post("/session")
|
|
def create_session(db: Session = Depends(get_db), current_user: User = Depends(get_current_user)) -> dict:
|
|
session = ChatService.create_session(db, current_user)
|
|
return api_success(CreateSessionResponse(sessionId=session.id).model_dump())
|
|
|
|
|
|
@router.get("/session/list")
|
|
def list_sessions(db: Session = Depends(get_db), current_user: User = Depends(get_current_user)) -> dict:
|
|
sessions = ChatService.list_sessions(db, current_user)
|
|
return api_success([ChatSessionRead.model_validate(session).model_dump() for session in sessions])
|
|
|
|
|
|
@router.get("/history")
|
|
def history(
|
|
sessionId: int = Query(...),
|
|
db: Session = Depends(get_db),
|
|
current_user: User = Depends(get_current_user),
|
|
) -> dict:
|
|
messages = ChatService.get_history(db, current_user, sessionId)
|
|
return api_success([ChatMessageRead.model_validate(message).model_dump(mode="json") for message in messages])
|
|
|
|
|
|
@router.put("/session/title")
|
|
def update_title(
|
|
payload: UpdateSessionTitleRequest,
|
|
db: Session = Depends(get_db),
|
|
current_user: User = Depends(get_current_user),
|
|
) -> dict:
|
|
session = ChatService.update_title(db, current_user, payload.sessionId, payload.title)
|
|
return api_success(ChatSessionRead.model_validate(session).model_dump())
|
|
|
|
|
|
@router.delete("/session/{session_id}")
|
|
def delete_session(
|
|
session_id: int,
|
|
db: Session = Depends(get_db),
|
|
current_user: User = Depends(get_current_user),
|
|
) -> dict:
|
|
ChatService.delete_session(db, current_user, session_id)
|
|
return api_success()
|
|
|
|
|
|
@router.post("/completions")
|
|
def completions(
|
|
payload: ChatCompletionRequest,
|
|
db: Session = Depends(get_db),
|
|
current_user: User = Depends(get_current_user),
|
|
) -> StreamingResponse:
|
|
return StreamingResponse(_chat_stream(payload, db, current_user), media_type="text/event-stream")
|
|
|
|
|
|
@router.post("/stop")
|
|
def stop(
|
|
payload: StopChatRequest,
|
|
db: Session = Depends(get_db),
|
|
current_user: User = Depends(get_current_user),
|
|
) -> dict:
|
|
ChatService.stop_generation(db, current_user, payload.sessionId)
|
|
return api_success()
|
|
|
|
|
|
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"
|
|
queued_token = queue_request.token
|
|
|
|
if queue_request.status == "rejected":
|
|
yield _sse_event(
|
|
"error",
|
|
message="当前请求过多，请稍后再试。",
|
|
activeCount=queue_request.active_count,
|
|
waitingCount=queue_request.waiting_count,
|
|
)
|
|
yield _sse_done()
|
|
return
|
|
|
|
if queue_request.status == "queued" and queued_token is not None:
|
|
yield _sse_event(
|
|
"queued",
|
|
message=f"当前请求较多，正在排队中，前方约 {max(queue_request.position - 1, 0)} 个请求。",
|
|
position=queue_request.position,
|
|
activeCount=queue_request.active_count,
|
|
waitingCount=queue_request.waiting_count,
|
|
)
|
|
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="排队等待超时，请稍后再试。")
|
|
yield _sse_done()
|
|
return
|
|
|
|
try:
|
|
yield _sse_event(
|
|
"generating",
|
|
message="已进入生成队列，正在生成回答。",
|
|
activeCount=queue_request.active_count,
|
|
waitingCount=queue_request.waiting_count,
|
|
)
|
|
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:
|
|
if acquired:
|
|
chat_queue_manager.release_slot()
|
|
elif queued_token is not None:
|
|
chat_queue_manager.cancel_waiter(queued_token)
|
|
|
|
yield _sse_done()
|
|
|
|
def _sse_event(event_type: str, **payload: object) -> str:
|
|
return f"data: {json.dumps({'type': event_type, **payload}, ensure_ascii=False)}\n\n"
|
|
|
|
|
|
def _sse_done() -> str:
|
|
return "data: [DONE]\n\n"
|