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)