src/routers/models.py - master - Anna_brv/ASR-Service

src/routers/models.py
138 строк4 KB
Anna
first_commit
30 апр 2026, 21:52
30 апр 2026, 21:52283885c
"""Роутер для управления моделями."""

from fastapi import APIRouter, Depends, HTTPException, status
from loguru import logger

from ..auth import verify_api_key
from ..config import settings
from ..models import (
    ModelListResponse,
    ModelInfo,
    LoadModelRequest,
    LoadModelResponse,
    UnloadModelResponse,
)
from ..services.model_manager import ModelManager

router = APIRouter()

DIARIZATION_MODEL = "pyannote/speaker-diarization-3.1"


def get_model_manager() -> ModelManager:
    return ModelManager(
        cache_dir=settings.model_cache_dir,
        default_model_name=settings.default_model,
        diarization_model_name=DIARIZATION_MODEL,
        hf_token=settings.hf_token.get_secret_value(),
    )


@router.get(
    "",
    response_model=ModelListResponse,
    summary="Список моделей",
    description="Получить список доступных моделей.",
)
async def list_models(
    model_manager: ModelManager = Depends(get_model_manager),
):
    models = [
        ModelInfo(
            name=settings.default_model,
            loaded=model_manager.is_model_loaded(settings.default_model),
            is_default=True,
            supports_diarization=False,
        ),
        ModelInfo(
            name=DIARIZATION_MODEL,
            loaded=model_manager.is_model_loaded(DIARIZATION_MODEL),
            is_default=False,
            supports_diarization=True,
        ),
    ]

    return ModelListResponse(
        models=models,
        default_model=settings.default_model,
    )


@router.post(
    "/load",
    response_model=LoadModelResponse,
    summary="Загрузить модель",
    description="Загрузить модель в кэш.",
)
async def load_model(
    request: LoadModelRequest,
    model_manager: ModelManager = Depends(get_model_manager),
):
    try:
        model_manager.load_model(request.model_name)
        logger.info(f"Модель {request.model_name} загружена")

        return LoadModelResponse(
            model_name=request.model_name,
            status="success",
            message="Модель загружена в кэш",
        )

    except Exception as e:
        logger.error(f"Ошибка загрузки модели {request.model_name}: {e}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Ошибка загрузки модели: {e}",
        )


@router.post(
    "/unload",
    response_model=UnloadModelResponse,
    summary="Выгрузить модель",
    description="Выгрузить модель из кэша.",
)
async def unload_model(
    model_name: str,
    model_manager: ModelManager = Depends(get_model_manager),
):
    if model_name == settings.default_model:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Нельзя выгрузить модель по умолчанию",
        )

    try:
        model_manager.unload_model(model_name)
        logger.info(f"Модель {model_name} выгружена")

        return UnloadModelResponse(
            model_name=model_name,
            status="success",
            message="Модель выгружена из кэша",
        )

    except Exception as e:
        logger.error(f"Ошибка выгрузки модели {model_name}: {e}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Ошибка выгрузки модели: {e}",
        )


@router.get(
    "/{model_name}/status",
    response_model=ModelInfo,
    summary="Статус модели",
    description="Получить статус загрузки модели.",
)
async def get_model_status(
    model_name: str,
    model_manager: ModelManager = Depends(get_model_manager),
):
    return ModelInfo(
        name=model_name,
        loaded=model_manager.is_model_loaded(model_name),
        is_default=model_name == settings.default_model,
        supports_diarization=model_name == DIARIZATION_MODEL,
    )