QuestionProject/ai_knowledge_base_v2/apps/backend/app/api/chat.py

from __future__ import annotations

import asyncio
import json
import logging
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_runtime import (
    cancel_chat_waiter,
    release_chat_slot,
    request_chat_slot,
    wait_for_chat_slot,
)
from app.services.chat_queue_service import load_chat_queue_config
from app.services.chat_stream_service import ChatStreamService

router = APIRouter()
logger = logging.getLogger(__name__)


@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_lease = await request_chat_slot(config)
    queue_request = queue_lease.request
    acquired = queue_request.status == "acquired"
    should_cancel_waiter = queue_request.status == "queued"

    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_lease = await wait_for_chat_slot(queue_lease, config)
        queue_request = queue_lease.request
        acquired = queue_request.status == "acquired"
        should_cancel_waiter = queue_request.status == "queued"
        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 == "timeout":
            yield _sse_event("error", message="排队等待超时，请稍后再试。")
            yield _sse_done()
            return
        if queue_request.status == "canceled":
            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:
        logger.exception("Chat stream failed")
        yield _sse_event("error", message="AI 回复生成失败，请稍后再试。")
    finally:
        if acquired:
            await release_chat_slot(queue_lease)
        elif should_cancel_waiter:
            await cancel_chat_waiter(queue_lease)

    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"