110 lines
3.5 KiB
Python
110 lines
3.5 KiB
Python
"""Model management endpoints."""
|
|
import logging
|
|
|
|
from fastapi import APIRouter, HTTPException
|
|
|
|
from app.schemas.models import ModelDownloadRequest, ModelDownloadResponse, ModelListResponse
|
|
from app.services import model_manager
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
router = APIRouter(prefix='/api/v1', tags=['models'])
|
|
|
|
|
|
@router.get('/models')
|
|
async def list_models() -> ModelListResponse:
|
|
"""List all available models."""
|
|
try:
|
|
available = model_manager.get_available_models()
|
|
local_count = sum(1 for m in available if m['available'])
|
|
|
|
return ModelListResponse(
|
|
available_models=available,
|
|
total_models=len(available),
|
|
local_models=local_count,
|
|
)
|
|
except Exception as e:
|
|
logger.error(f'Failed to list models: {e}', exc_info=True)
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@router.post('/models/download')
|
|
async def download_models(request: ModelDownloadRequest) -> ModelDownloadResponse:
|
|
"""Download one or more models."""
|
|
if not request.models:
|
|
raise HTTPException(status_code=400, detail='No models specified')
|
|
|
|
try:
|
|
logger.info(f'Downloading models: {request.models}')
|
|
results = await model_manager.download_models(request.models)
|
|
|
|
downloaded = []
|
|
failed = []
|
|
errors = {}
|
|
|
|
for model_name, (success, message) in results.items():
|
|
if success:
|
|
downloaded.append(model_name)
|
|
else:
|
|
failed.append(model_name)
|
|
errors[model_name] = message
|
|
|
|
return ModelDownloadResponse(
|
|
success=len(failed) == 0,
|
|
message=f'Downloaded {len(downloaded)} model(s)',
|
|
downloaded=downloaded,
|
|
failed=failed,
|
|
errors=errors,
|
|
)
|
|
except Exception as e:
|
|
logger.error(f'Model download failed: {e}', exc_info=True)
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@router.get('/models/{model_name}')
|
|
async def get_model_info(model_name: str):
|
|
"""Get information about a specific model."""
|
|
models = model_manager.get_available_models()
|
|
|
|
for model in models:
|
|
if model['name'] == model_name:
|
|
return model
|
|
|
|
raise HTTPException(status_code=404, detail=f'Model not found: {model_name}')
|
|
|
|
|
|
@router.post('/models/{model_name}/download')
|
|
async def download_model(model_name: str):
|
|
"""Download a specific model."""
|
|
try:
|
|
success, message = await model_manager.download_model(model_name)
|
|
|
|
if not success:
|
|
raise HTTPException(status_code=500, detail=message)
|
|
|
|
return {
|
|
'success': True,
|
|
'message': message,
|
|
'model': model_name,
|
|
}
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f'Failed to download model {model_name}: {e}', exc_info=True)
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@router.get('/models-info')
|
|
async def get_models_directory_info():
|
|
"""Get information about the models directory."""
|
|
try:
|
|
info = model_manager.get_models_directory_info()
|
|
return {
|
|
'models_directory': info['path'],
|
|
'total_size_mb': round(info['size_mb'], 2),
|
|
'model_count': info['model_count'],
|
|
}
|
|
except Exception as e:
|
|
logger.error(f'Failed to get models directory info: {e}', exc_info=True)
|
|
raise HTTPException(status_code=500, detail=str(e))
|