Maas-Model-platform/backend/app/routers/llm_router.py

"""
AI对话API路由

提供LLM对话的RESTful API端点，支持联网搜索功能
需求: 2.1, 3.1, 3.2, 5.5, 6.4, 10.5
流式搜索需求: 9.1, 9.2, 9.3, 9.4, 9.5
"""
from typing import List, AsyncGenerator
from fastapi import APIRouter, Depends, Request, HTTPException
from fastapi.responses import StreamingResponse
from sqlalchemy.orm import Session

from app.database import get_db, SessionLocal
from app.services.llm_service import LLMService
from app.services.system_config_manager import get_config_bool
from app.schemas.model_schema import ApiResponse, ModelResponse
from app.schemas.llm_schema import (
    ChatRequest, ChatResponse, 
    EnhancedChatRequest, EnhancedChatResponse
)
from app.models.user import User
from app.middleware import get_current_user_from_request

router = APIRouter(prefix="/api/llm", tags=["AI对话"])


async def _stream_with_db(generator: AsyncGenerator) -> AsyncGenerator[str, None]:
    """
    包装流式生成器，确保 db 连接在流完全结束后才关闭。
    生成器本身负责持有 db 引用，此函数只负责透传数据。
    """
    async for chunk in generator:
        yield chunk


@router.post("/chat")
async def chat(
    request: EnhancedChatRequest,
    req: Request,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user_from_request)
):
    """
    统一对话API端点
    
    支持流式和非流式输出，自动检测搜索选项
    需要用户认证，使用用户的apikey调用百炼平台
    需要余额检查，余额不足时返回402错误
    向后兼容：如果没有提供search_options，则使用普通对话模式
    """
    api_key = current_user.apikey
    if not api_key:
        return ApiResponse(code=403, message="未配置API密钥，请在用户设置中配置apikey", data=None)

    # 优先使用模型自带的 api_key（爬虫同步的），没有才 fallback 到用户自己配置的 apikey
    from app.services.crypto_utils import get_effective_api_key
    api_key = get_effective_api_key(db, request.model, api_key)

    search_enabled = (
        hasattr(request, 'search_options') and
        request.search_options and
        request.search_options.enable_search
    )

    if search_enabled and not get_config_bool("enable_search", True):
        return ApiResponse(code=403, message="系统暂未开放联网搜索功能", data=None)

    try:
        if request.stream:
            # 流式请求：手动管理 db 生命周期，确保流结束后才关闭连接
            stream_db = SessionLocal()

            async def stream_and_close():
                try:
                    service = LLMService(stream_db, api_key=api_key, user_id=str(current_user.id))
                    if search_enabled:
                        gen = service.chat_stream_with_search(request, conversation_id=request.conversation_id)
                    else:
                        gen = service.chat_stream(request, conversation_id=request.conversation_id)
                    async for chunk in gen:
                        yield chunk
                finally:
                    stream_db.close()

            return StreamingResponse(
                stream_and_close(),
                media_type="text/event-stream",
                headers={
                    "Cache-Control": "no-cache",
                    "Connection": "keep-alive",
                    "X-Accel-Buffering": "no"
                }
            )
        else:
            service = LLMService(db, api_key=api_key, user_id=str(current_user.id))
            if search_enabled:
                data = service.chat_with_search(request, conversation_id=request.conversation_id)
            else:
                data = service.chat(request, conversation_id=request.conversation_id)
            return ApiResponse(code=200, message="success", data=data)

    except HTTPException:
        raise
    except Exception as e:
        return ApiResponse(code=500, message=f"对话服务异常: {str(e)}", data=None)


@router.post("/chat/search")
async def chat_with_search(
    request: EnhancedChatRequest,
    req: Request,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user_from_request)
):
    """
    支持搜索的对话API端点
    
    支持流式和非流式输出，集成联网搜索功能
    需要用户认证，使用用户的apikey调用百炼平台
    需求: 5.5, 6.4, 9.1, 9.2, 9.3, 9.4, 9.5
    """
    if not get_config_bool("enable_search", True):
        return ApiResponse(code=403, message="系统暂未开放联网搜索功能", data=None)

    api_key = current_user.apikey
    if not api_key:
        return ApiResponse(code=403, message="未配置API密钥，请在用户设置中配置apikey", data=None)

    # 优先使用模型自带的 api_key（爬虫同步的），没有才 fallback 到用户自己配置的 apikey
    from app.services.crypto_utils import get_effective_api_key
    api_key = get_effective_api_key(db, request.model, api_key)

    try:
        if request.stream:
            stream_db = SessionLocal()

            async def stream_and_close():
                try:
                    service = LLMService(stream_db, api_key=api_key, user_id=str(current_user.id))
                    async for chunk in service.chat_stream_with_search(request, conversation_id=request.conversation_id):
                        yield chunk
                finally:
                    stream_db.close()

            return StreamingResponse(
                stream_and_close(),
                media_type="text/event-stream",
                headers={
                    "Cache-Control": "no-cache",
                    "Connection": "keep-alive",
                    "X-Accel-Buffering": "no"
                }
            )
        else:
            service = LLMService(db, api_key=api_key, user_id=str(current_user.id))
            data = service.chat_with_search(request, conversation_id=request.conversation_id)
            return ApiResponse(code=200, message="success", data=data)

    except HTTPException:
        raise
    except Exception as e:
        return ApiResponse(code=500, message=f"搜索增强对话服务异常: {str(e)}", data=None)


@router.get("/search/models", response_model=ApiResponse[List[str]])
def get_search_supported_models(
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user_from_request)
):
    """
    获取支持搜索功能的模型列表
    
    返回支持联网搜索的模型名称列表
    需求: 10.1, 10.4
    """
    api_key = current_user.apikey
    if not api_key:
        return ApiResponse(code=403, message="未配置API密钥，请在用户设置中配置apikey", data=None)
    
    service = LLMService(db, api_key=api_key, user_id=str(current_user.id))
    data = service.get_search_supported_models()
    return ApiResponse(code=200, message="success", data=data)


@router.get("/search/check/{model}")
def check_search_support(
    model: str,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user_from_request)
):
    """
    检查指定模型是否支持搜索功能
    
    Args:
        model: 模型名称
        
    Returns:
        是否支持搜索功能的布尔值
        
    需求: 10.1, 10.4
    """
    api_key = current_user.apikey
    if not api_key:
        return ApiResponse(code=403, message="未配置API密钥，请在用户设置中配置apikey", data=None)
    
    service = LLMService(db, api_key=api_key, user_id=str(current_user.id))
    is_supported = service.is_search_supported(model)
    return ApiResponse(
        code=200, 
        message="success", 
        data={"model": model, "search_supported": is_supported}
    )


@router.get("/models", response_model=ApiResponse[List[ModelResponse]])
def get_llm_models(db: Session = Depends(get_db)):
    """
    获取所有可用的LLM模型列表
    
    返回type=0的语言模型
    """
    service = LLMService(db)
    data = service.get_llm_models()
    return ApiResponse(code=200, message="success", data=data)