llm-monitor/app/api/models.py

"""
Models endpoints - Gestione dei modelli Ollama
"""

from fastapi import APIRouter, HTTPException
from pydantic import BaseModel
from typing import List, Optional
from datetime import datetime
import requests
import logging
from app.config import settings

logger = logging.getLogger(__name__)
router = APIRouter()


def ensure_rw_api_enabled() -> None:
    """Blocca le API di scrittura se non abilitate esplicitamente."""
    if not settings.ENABLE_MODEL_RW_API:
        raise HTTPException(
            status_code=404,
            detail="Endpoint non disponibile"
        )

class ModelInfo(BaseModel):
    """Informazioni su un modello"""
    name: str
    digest: str
    size: int
    modified_at: datetime
    
    class Config:
        json_schema_extra = {
            "example": {
                "name": "llama2",
                "digest": "abc123def456...",
                "size": 3825922048,
                "modified_at": "2024-01-15T10:30:00Z"
            }
        }

class ModelsResponse(BaseModel):
    """Risposta con lista di modelli"""
    models: List[ModelInfo]
    total: int
    
    class Config:
        json_schema_extra = {
            "example": {
                "models": [
                    {
                        "name": "llama2",
                        "digest": "abc123def456...",
                        "size": 3825922048,
                        "modified_at": "2024-01-15T10:30:00Z"
                    }
                ],
                "total": 1
            }
        }

@router.get("/models", response_model=ModelsResponse)
async def get_models():
    """
    Recupera l'elenco di tutti i modelli caricati in Ollama
    
    Returns:
        ModelsResponse: Lista dei modelli disponibili
        
    Raises:
        HTTPException: Se Ollama non è disponibile
    """
    try:
        response = requests.get(
            f"{settings.OLLAMA_HOST}/api/tags",
            timeout=settings.OLLAMA_TIMEOUT
        )
        
        if response.status_code != 200:
            raise HTTPException(
                status_code=502,
                detail="Ollama non disponibile"
            )
        
        data = response.json()
        models_data = data.get("models", [])
        
        models = [
            ModelInfo(
                name=model.get("name", "unknown"),
                digest=model.get("digest", ""),
                size=model.get("size", 0),
                modified_at=datetime.fromisoformat(
                    model.get("modified_at", "").replace("Z", "+00:00")
                ) if model.get("modified_at") else datetime.utcnow()
            )
            for model in models_data
        ]
        
        return ModelsResponse(
            models=models,
            total=len(models)
        )
        
    except requests.exceptions.Timeout:
        raise HTTPException(
            status_code=504,
            detail="Timeout: Ollama non ha risposto in tempo"
        )
    except requests.exceptions.ConnectionError:
        raise HTTPException(
            status_code=502,
            detail="Impossible connettersi a Ollama"
        )
    except Exception as e:
        logger.error(f"Error fetching models: {e}")
        raise HTTPException(
            status_code=500,
            detail="Errore nel recupero dei modelli"
        )

@router.get("/models/{model_name}", response_model=ModelInfo)
async def get_model(model_name: str):
    """
    Recupera le informazioni di un modello specifico
    
    Args:
        model_name: Nome del modello da cercare
        
    Returns:
        ModelInfo: Informazioni del modello
        
    Raises:
        HTTPException: Se il modello non esiste o Ollama non è disponibile
    """
    try:
        response = requests.get(
            f"{settings.OLLAMA_HOST}/api/tags",
            timeout=settings.OLLAMA_TIMEOUT
        )
        
        if response.status_code != 200:
            raise HTTPException(
                status_code=502,
                detail="Ollama non disponibile"
            )
        
        data = response.json()
        models_data = data.get("models", [])
        
        # Cercare il modello
        for model in models_data:
            if model.get("name") == model_name:
                return ModelInfo(
                    name=model.get("name", "unknown"),
                    digest=model.get("digest", ""),
                    size=model.get("size", 0),
                    modified_at=datetime.fromisoformat(
                        model.get("modified_at", "").replace("Z", "+00:00")
                    ) if model.get("modified_at") else datetime.utcnow()
                )
        
        raise HTTPException(
            status_code=404,
            detail=f"Modello '{model_name}' non trovato"
        )
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error fetching model: {e}")
        raise HTTPException(
            status_code=500,
            detail="Errore nel recupero del modello"
        )

@router.post(
    "/models/{model_name}/pull",
    include_in_schema=settings.ENABLE_MODEL_RW_API
)
async def pull_model(model_name: str):
    """
    Scarica/carica un modello in Ollama
    
    Args:
        model_name: Nome del modello da scaricare
        
    Returns:
        dict: Status del download
    """
    ensure_rw_api_enabled()
    try:
        response = requests.post(
            f"{settings.OLLAMA_HOST}/api/pull",
            json={"name": model_name},
            timeout=None  # Pull può essere lungo
        )
        
        if response.status_code not in [200, 201]:
            raise HTTPException(
                status_code=502,
                detail="Errore nel pull del modello"
            )
        
        return {"status": "pulling", "model": model_name}
        
    except Exception as e:
        logger.error(f"Error pulling model: {e}")
        raise HTTPException(
            status_code=500,
            detail="Errore nel pull del modello"
        )

@router.delete(
    "/models/{model_name}",
    include_in_schema=settings.ENABLE_MODEL_RW_API
)
async def delete_model(model_name: str):
    """
    Elimina un modello da Ollama
    
    Args:
        model_name: Nome del modello da eliminare
        
    Returns:
        dict: Confirmazione eliminazione
    """
    ensure_rw_api_enabled()
    try:
        response = requests.delete(
            f"{settings.OLLAMA_HOST}/api/delete",
            json={"name": model_name},
            timeout=settings.OLLAMA_TIMEOUT
        )
        
        if response.status_code not in [200, 204]:
            raise HTTPException(
                status_code=502,
                detail="Errore nell'eliminazione del modello"
            )
        
        return {"status": "deleted", "model": model_name}
        
    except Exception as e:
        logger.error(f"Error deleting model: {e}")
        raise HTTPException(
            status_code=500,
            detail="Errore nell'eliminazione del modello"
        )