NyaHome/src/nyahome/router/aii_router.py

from typing import Annotated, Sequence

from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy.exc import NoResultFound
from sqlmodel import Session, select

from nyahome.database import (
    AiiModel,
    AiiModelPublic,
    AiiProvider,
    AiiProviderPublic,
    AiiProviderPublicWithoutKey,
    ModelUser,
    get_session,
)
from nyahome.service.aii_service import apply_get_models, s_check_remote_model, s_list_remote_provider_models

from .auth import verify_token
from .response_model import ReturnDto

aii_router = APIRouter(tags=["Aii"], prefix="/aii")


@aii_router.get("/model/", name="获取模型列表")
async def get_all_model(session: Annotated[Session, Depends(get_session)]) -> list[dict]:
    """
    获取 AI 模型列表。
    此接口无需用户登录即可访问。

    Returns:
        AiiModel 列表
    """
    return apply_get_models(session)


@aii_router.post("/model/", name="添加模型")
async def add_model(
    model: AiiModelPublic,
    user: Annotated[ModelUser, Depends(verify_token)],
    session: Annotated[Session, Depends(get_session)],
) -> AiiModel:
    """
    添加新的 AI 模型。需要基于已添加的模型提供商。
    此接口需要管理员访问。
    不会进行可用性检查，因此 WebUI 在前端实现了检查按钮。此端点不会负责检查。

    Raises:
        HTTPException: 401 用户无权限管理模型（未登录或非管理员）
        HTTPException: 404 模型提供商不存在

    Returns:
        AiiModel
    """
    if not user.is_admin:
        raise HTTPException(status_code=401, detail="用户无权限管理模型。") from None

    try:
        ap: AiiProvider = session.exec(select(AiiProvider).where(AiiProvider.id == model.aii_provider_id)).one()
    except NoResultFound:
        raise HTTPException(status_code=404, detail="Provider 不存在。") from None
    am = AiiModel(
        model_name=model.model_name,
        max_context_length=model.max_context_length,
        aii_provider_id=model.aii_provider_id,
        aii_provider=ap,
    )
    session.add(am)
    session.commit()
    session.refresh(am)
    return am


@aii_router.post("/model/{id_}", name="修改模型")
async def edit_model(
    id_: int,
    model: AiiModelPublic,
    user: Annotated[ModelUser, Depends(verify_token)],
    session: Annotated[Session, Depends(get_session)],
) -> AiiModel:
    """
    修改已添加的 AI 模型。
    此接口需要管理员访问。
    不会进行可用性检查，因此 WebUI 在前端实现了检查按钮。此端点不会负责检查。
    **只允许修改模型的名称、最大上下文长度和是否支持思考。**

    Raises:
        HTTPException: 400 模型提供商 ID 不匹配
        HTTPException: 401 用户无权限管理模型（未登录或非管理员）
        HTTPException: 404 模型提供商不存在

    Returns:
        AiiModel
    """
    if not user.is_admin:
        raise HTTPException(status_code=401, detail="用户无权限管理模型。") from None

    try:
        ap: AiiProvider = session.exec(select(AiiProvider).where(AiiProvider.id == model.aii_provider_id)).one()
    except NoResultFound:
        raise HTTPException(status_code=404, detail="Provider 不存在。") from None

    try:
        am: AiiModel = session.exec(select(AiiModel).where(AiiModel.id == id_)).one()
    except NoResultFound:
        raise HTTPException(status_code=404, detail="模型不存在。") from None

    if ap.id != am.aii_provider_id:
        raise HTTPException(status_code=400, detail="模型提供商 ID 不匹配。") from None

    am.model_name = model.model_name
    am.max_context_length = model.max_context_length
    am.reasonable = model.reasonable
    session.add(am)
    session.commit()
    session.refresh(am)
    return am


@aii_router.get("/provider/", name="获取提供商列表")
async def get_all_provider(session: Annotated[Session, Depends(get_session)]) -> Sequence[AiiProvider]:
    """
    获取 AI 模型提供商列表。
    此接口无需用户登录即可访问。

    Returns:
        被 ReturnDto 包裹的 AiiProvider 列表
    """
    return session.exec(select(AiiProvider)).all()


@aii_router.post("/provider/", name="添加提供商")
async def add_provider(
    provider: AiiProviderPublic,
    user: Annotated[ModelUser, Depends(verify_token)],
    session: Annotated[Session, Depends(get_session)],
) -> AiiProvider:
    """
    添加新的 AI 模型提供商。
    此接口需要管理员才能访问。
    不会进行可用性检查，因此 WebUI 在前端实现了检查按钮。此端点不会负责检查。

    Raises:
        HTTPException: 401 表示用户未登录或非管理员。

    Returns:
        被 ReturnDto 包裹的、添加的 AiiProvider
    """
    if not user.is_admin:
        raise HTTPException(status_code=401, detail="用户无权限管理模型。") from None
    ap = AiiProvider(name=provider.name, base_url=provider.base_url, api_key=provider.api_key)
    session.add(ap)
    session.commit()
    session.refresh(ap)
    return ap


@aii_router.post("/provider/{id_}/", name="修改提供商")
async def edit_provider(
    id_: int,
    provider: AiiProviderPublicWithoutKey,
    user: Annotated[ModelUser, Depends(verify_token)],
    session: Annotated[Session, Depends(get_session)],
) -> AiiProvider:
    """
    修改 AI 模型提供商。
    此接口需要管理员才能访问。
    不会进行可用性检查，因此 WebUI 在前端实现了检查按钮。此端点不会负责检查。
    **只允许修改模型提供商的名称和 Base URL。**

    Raises:
        HTTPException: 400 模型提供商 ID 不匹配。
        HTTPException: 401 表示用户未登录或非管理员。
        HTTPException: 404 提供商不存在。

    Returns:
        被 ReturnDto 包裹的、添加的 AiiProvider
    """
    if not user.is_admin:
        raise HTTPException(status_code=401, detail="用户无权限管理模型。") from None

    if provider.id != id_:
        raise HTTPException(status_code=400, detail="模型提供商 ID 不匹配。") from None

    try:
        ap: AiiProvider = session.exec(select(AiiProvider).where(AiiProvider.id == id_)).one()
    except NoResultFound:
        raise HTTPException(status_code=404, detail="提供商不存在。") from None

    ap.name = provider.name
    ap.base_url = provider.base_url
    session.add(ap)
    session.commit()
    session.refresh(ap)
    return ap


@aii_router.get("/provider/{id_}/remote/models/", name="获取提供商远端模型")
async def get_provider_remote_models(
    id_: int, user: Annotated[ModelUser, Depends(verify_token)], session: Annotated[Session, Depends(get_session)]
) -> ReturnDto:
    """
    查看指定模型提供商提供的远端模型列表。并非添加到 NyaHome 的模型列表。
    此接口需要管理员才能访问。

    Raises:
        HTTPException: 401 表示用户未登录或非管理员。

    Returns:
        被 ReturnDto 包裹的、模型名称字符串列表
    """
    if not user.is_admin:
        raise HTTPException(status_code=401, detail="用户无权限管理模型。") from None
    try:
        ap: AiiProvider = session.exec(select(AiiProvider).where(AiiProvider.id == id_)).one()
    except NoResultFound:
        raise HTTPException(status_code=404, detail="Provider 不存在。") from None
    models = await s_list_remote_provider_models(ap.base_url, ap.api_key)
    # 只返回模型名称列表，方便前端填入表单
    return ReturnDto(result=[m["id"] for m in models])


@aii_router.get("/provider/{id_}/remote/model/{model_name}/", name="检查指定远端模型可用性")
async def check_remote_provider_model(
    id_: int, model_name: str, session: Annotated[Session, Depends(get_session)]
) -> ReturnDto:
    """
    检测指定提供商的指定名称远端模型是否可用。

    Raises:
        HTTPException: 404 表明提供商 ID 未找到。

    Returns:
        ReturnDto，其中 result 字段为布尔值，表明指定名称模型的可用状态。
    """
    try:
        ap: AiiProvider = session.exec(select(AiiProvider).where(AiiProvider.id == id_)).one()
    except NoResultFound:
        raise HTTPException(status_code=404, detail="Provider 不存在。") from None
    return ReturnDto(result=await s_check_remote_model(model_name, ap.base_url, ap.api_key))


@aii_router.post("/remote/provider/check/", name="检查指定提供商可用性")
async def check_remote_provider(provider: AiiProviderPublic) -> ReturnDto:
    """
    检查指定提供商是否可用。会返回提供商提供的模型数量作为测试。

    Returns:
        ReturnDto，其中 success 字段为布尔值，表明可用状态；如果为真，result 字段是整型模型数量。
    """
    try:
        count = len(await s_list_remote_provider_models(provider.base_url, provider.api_key))
        return ReturnDto(result=count)
    except TypeError:
        return ReturnDto(success=False)