ai-agent-admin/backend-fastapi/ai_platform/api/chat_api.py

"""
对话 API
"""
import json
import logging
from typing import List, Optional

from fastapi import APIRouter, Depends, HTTPException, Query
from fastapi.responses import StreamingResponse
from sqlalchemy import select, func
from sqlalchemy.ext.asyncio import AsyncSession

from app.database import get_db
from app.base_schema import PaginatedResponse, ResponseModel
from utils.context import get_current_user_id_from_context
from ai_platform.models import Conversation, Message, AIApp
from ai_platform.schemas.chat_schema import (
    ConversationCreate,
    ConversationUpdate,
    ConversationResponse,
    ConversationListResponse,
    MessageResponse,
    SendMessageInput,
    MessageFeedbackInput,
)

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/chat", tags=["AI-对话"])


@router.get("/conversations", response_model=PaginatedResponse[ConversationListResponse], summary="对话列表")
async def list_conversations(
    app_id: str = Query(..., description="应用 ID"),
    page: int = Query(1, ge=1, description="页码"),
    page_size: int = Query(20, ge=1, le=100, alias="pageSize", description="每页数量"),
    db: AsyncSession = Depends(get_db),
):
    """获取对话列表"""
    query = select(Conversation).where(
        Conversation.app_id == app_id,
        Conversation.is_deleted == False
    )
    
    # 获取总数
    count_result = await db.execute(select(func.count()).select_from(query.subquery()))
    total = count_result.scalar() or 0
    
    # 分页
    offset = (page - 1) * page_size
    query = query.order_by(
        Conversation.is_pinned.desc(),
        Conversation.sys_update_datetime.desc()
    )
    query = query.offset(offset).limit(page_size)
    
    result = await db.execute(query)
    conversations = result.scalars().all()
    
    items = [
        {
            "id": c.id,
            "title": c.title or "",
            "message_count": c.message_count or 0,
            "is_pinned": c.is_pinned or False,
            "sys_update_datetime": c.sys_update_datetime,
        }
        for c in conversations
    ]
    
    return PaginatedResponse(items=items, total=total)


@router.post("/conversations", response_model=ConversationResponse, summary="创建对话")
async def create_conversation(data: ConversationCreate, db: AsyncSession = Depends(get_db)):
    """创建对话"""
    # 验证应用
    app_result = await db.execute(
        select(AIApp).where(AIApp.id == data.app_id, AIApp.is_deleted == False)
    )
    app = app_result.scalar_one_or_none()
    if not app:
        raise HTTPException(status_code=400, detail="应用不存在")

    user_id = get_current_user_id_from_context()
    if not user_id:
        raise HTTPException(status_code=401, detail="未提供认证凭据")
    
    conversation = Conversation(
        app_id=data.app_id,
        user_id=user_id,
        title=data.title or "新对话",
    )
    db.add(conversation)
    await db.commit()
    await db.refresh(conversation)
    
    return await _build_conversation_response(conversation, db)


@router.get("/conversations/{conversation_id}", response_model=ConversationResponse, summary="对话详情")
async def get_conversation(conversation_id: str, db: AsyncSession = Depends(get_db)):
    """获取对话详情"""
    result = await db.execute(
        select(Conversation).where(
            Conversation.id == conversation_id,
            Conversation.is_deleted == False
        )
    )
    conversation = result.scalar_one_or_none()
    if not conversation:
        raise HTTPException(status_code=404, detail="对话不存在")
    
    return await _build_conversation_response(conversation, db)


@router.put("/conversations/{conversation_id}", response_model=ConversationResponse, summary="更新对话")
async def update_conversation(
    conversation_id: str,
    data: ConversationUpdate,
    db: AsyncSession = Depends(get_db),
):
    """更新对话"""
    result = await db.execute(
        select(Conversation).where(
            Conversation.id == conversation_id,
            Conversation.is_deleted == False
        )
    )
    conversation = result.scalar_one_or_none()
    if not conversation:
        raise HTTPException(status_code=404, detail="对话不存在")
    
    if data.title is not None:
        conversation.title = data.title
    if data.is_pinned is not None:
        conversation.is_pinned = data.is_pinned
    
    await db.commit()
    await db.refresh(conversation)
    
    return await _build_conversation_response(conversation, db)


@router.delete("/conversations/{conversation_id}", response_model=ResponseModel, summary="删除对话")
async def delete_conversation(conversation_id: str, db: AsyncSession = Depends(get_db)):
    """删除对话"""
    result = await db.execute(
        select(Conversation).where(
            Conversation.id == conversation_id,
            Conversation.is_deleted == False
        )
    )
    conversation = result.scalar_one_or_none()
    if not conversation:
        raise HTTPException(status_code=404, detail="对话不存在")
    
    conversation.is_deleted = True
    await db.commit()
    
    return ResponseModel(message="删除成功")


@router.get("/conversations/{conversation_id}/messages", response_model=List[MessageResponse], summary="获取消息列表")
async def get_messages(
    conversation_id: str,
    limit: int = Query(50, ge=1, le=200, description="限制数量"),
    db: AsyncSession = Depends(get_db),
):
    """获取对话消息"""
    # 验证对话存在
    conv_result = await db.execute(
        select(Conversation).where(
            Conversation.id == conversation_id,
            Conversation.is_deleted == False
        )
    )
    if not conv_result.scalar_one_or_none():
        raise HTTPException(status_code=404, detail="对话不存在")
    
    query = select(Message).where(
        Message.conversation_id == conversation_id,
        Message.is_deleted == False
    ).order_by(Message.sys_create_datetime).limit(limit)
    
    result = await db.execute(query)
    messages = result.scalars().all()
    
    return [
        {
            "id": m.id,
            "role": m.role,
            "content": m.content or "",
            "status": m.status or "completed",
            "prompt_tokens": m.prompt_tokens or 0,
            "completion_tokens": m.completion_tokens or 0,
            "total_tokens": m.total_tokens or 0,
            "model_name": m.model_name or "",
            "latency": m.latency or 0,
            "error_message": m.error_message or "",
            "feedback": m.feedback or "",
            "sys_create_datetime": m.sys_create_datetime,
        }
        for m in messages
    ]


@router.post("/conversations/{conversation_id}/messages", response_model=MessageResponse, summary="发送消息")
async def send_message(
    conversation_id: str,
    data: SendMessageInput,
    db: AsyncSession = Depends(get_db),
):
    """发送消息并获取 AI 回复"""
    from ai_platform.services.chat_service import ChatService
    
    # 验证对话存在
    conv_result = await db.execute(
        select(Conversation).where(
            Conversation.id == conversation_id,
            Conversation.is_deleted == False
        )
    )
    conversation = conv_result.scalar_one_or_none()
    if not conversation:
        raise HTTPException(status_code=404, detail="对话不存在")
    
    # 使用ChatService发送消息
    chat_service = ChatService(db)
    _, assistant_message = await chat_service.send_message(
        conversation_id=conversation_id,
        content=data.content,
    )
    
    return {
        "id": assistant_message.id,
        "role": assistant_message.role,
        "content": assistant_message.content or "",
        "status": assistant_message.status or "completed",
        "prompt_tokens": assistant_message.prompt_tokens or 0,
        "completion_tokens": assistant_message.completion_tokens or 0,
        "total_tokens": assistant_message.total_tokens or 0,
        "model_name": assistant_message.model_name or "",
        "latency": assistant_message.latency or 0,
        "error_message": assistant_message.error_message or "",
        "feedback": assistant_message.feedback or "",
        "sys_create_datetime": assistant_message.sys_create_datetime,
    }


@router.post("/conversations/{conversation_id}/messages/stream", summary="流式发送消息")
async def send_message_stream(
    conversation_id: str,
    data: SendMessageInput,
    db: AsyncSession = Depends(get_db),
):
    """发送消息并获取 AI 流式回复（SSE）"""
    from ai_platform.services.chat_service import ChatService
    
    # 验证对话存在
    conv_result = await db.execute(
        select(Conversation).where(
            Conversation.id == conversation_id,
            Conversation.is_deleted == False
        )
    )
    conversation = conv_result.scalar_one_or_none()
    if not conversation:
        raise HTTPException(status_code=404, detail="对话不存在")
    
    async def generate():
        chat_service = ChatService(db)
        async for event in chat_service.send_message_stream(
            conversation_id=conversation_id,
            content=data.content,
        ):
            yield f"data: {json.dumps(event, ensure_ascii=False)}\n\n"
        yield "data: [DONE]\n\n"
    
    return StreamingResponse(
        generate(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
        }
    )


@router.post("/messages/{message_id}/feedback", response_model=ResponseModel, summary="消息反馈")
async def message_feedback(
    message_id: str,
    data: MessageFeedbackInput,
    db: AsyncSession = Depends(get_db),
):
    """对消息进行反馈"""
    result = await db.execute(
        select(Message).where(
            Message.id == message_id,
            Message.is_deleted == False
        )
    )
    message = result.scalar_one_or_none()
    if not message:
        raise HTTPException(status_code=404, detail="消息不存在")
    
    if data.feedback not in ("like", "dislike", ""):
        raise HTTPException(status_code=400, detail="无效的反馈类型")
    
    message.feedback = data.feedback
    await db.commit()
    
    return ResponseModel(message="反馈成功")


async def _build_conversation_response(conversation: Conversation, db: AsyncSession) -> dict:
    """构建对话输出"""
    app_name = ""
    if conversation.app_id:
        app_result = await db.execute(
            select(AIApp).where(AIApp.id == conversation.app_id)
        )
        app = app_result.scalar_one_or_none()
        app_name = app.name if app else ""
    
    return {
        "id": conversation.id,
        "app_id": conversation.app_id,
        "app_name": app_name,
        "title": conversation.title or "",
        "message_count": conversation.message_count or 0,
        "total_tokens": conversation.total_tokens or 0,
        "is_pinned": conversation.is_pinned or False,
        "sort": conversation.sort or 0,
        "sys_create_datetime": conversation.sys_create_datetime,
        "sys_update_datetime": conversation.sys_update_datetime,
    }