foxus/backend/app/api/models.py

"""
Models API routes for managing LLM models
"""

from typing import List
from fastapi import APIRouter, HTTPException
from app.models.ai import ModelInfo, ModelListResponse
from app.services.ollama_service import ollama_service
from app.core.config import settings

router = APIRouter()

@router.get("/list", response_model=ModelListResponse)
async def list_models():
    """List all available models from Ollama"""
    try:
        if not await ollama_service.is_available():
            raise HTTPException(
                status_code=503,
                detail="Ollama service is not available. Please ensure Ollama is running."
            )
        
        # Get models from Ollama
        ollama_models = await ollama_service.list_models()
        
        model_infos = []
        for model in ollama_models:
            model_name = model.get("name", "unknown")
            model_size = model.get("size", 0)
            
            # Format size for display
            size_str = format_bytes(model_size)
            
            # Get model capabilities based on name
            capabilities = get_model_capabilities(model_name)
            
            # Get model description
            description = get_model_description(model_name)
            
            model_info = ModelInfo(
                name=model_name,
                size=size_str,
                description=description,
                capabilities=capabilities,
                is_available=True
            )
            
            model_infos.append(model_info)
        
        # Add supported models that aren't installed
        installed_model_names = [m.name for m in model_infos]
        for supported_model in settings.SUPPORTED_MODELS:
            if supported_model not in installed_model_names:
                model_info = ModelInfo(
                    name=supported_model,
                    size="Not installed",
                    description=get_model_description(supported_model),
                    capabilities=get_model_capabilities(supported_model),
                    is_available=False
                )
                model_infos.append(model_info)
        
        return ModelListResponse(
            models=model_infos,
            default_model=settings.DEFAULT_MODEL,
            current_model=settings.DEFAULT_MODEL
        )
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@router.post("/pull/{model_name}")
async def pull_model(model_name: str):
    """Pull/download a model from Ollama"""
    try:
        if not await ollama_service.is_available():
            raise HTTPException(
                status_code=503,
                detail="Ollama service is not available"
            )
        
        if model_name not in settings.SUPPORTED_MODELS:
            raise HTTPException(
                status_code=400,
                detail=f"Model {model_name} is not in the supported models list"
            )
        
        success = await ollama_service.pull_model(model_name)
        
        if success:
            return {"message": f"Model {model_name} pulled successfully"}
        else:
            raise HTTPException(
                status_code=500,
                detail=f"Failed to pull model {model_name}"
            )
            
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@router.get("/info/{model_name}", response_model=ModelInfo)
async def get_model_info(model_name: str):
    """Get detailed information about a specific model"""
    try:
        if not await ollama_service.is_available():
            raise HTTPException(
                status_code=503,
                detail="Ollama service is not available"
            )
        
        # Get all models
        ollama_models = await ollama_service.list_models()
        
        # Find the specific model
        target_model = None
        for model in ollama_models:
            if model.get("name") == model_name:
                target_model = model
                break
        
        if not target_model:
            # Check if it's a supported model that's not installed
            if model_name in settings.SUPPORTED_MODELS:
                return ModelInfo(
                    name=model_name,
                    size="Not installed",
                    description=get_model_description(model_name),
                    capabilities=get_model_capabilities(model_name),
                    is_available=False
                )
            else:
                raise HTTPException(
                    status_code=404,
                    detail=f"Model {model_name} not found"
                )
        
        size_str = format_bytes(target_model.get("size", 0))
        
        return ModelInfo(
            name=model_name,
            size=size_str,
            description=get_model_description(model_name),
            capabilities=get_model_capabilities(model_name),
            is_available=True
        )
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@router.get("/supported")
async def get_supported_models():
    """Get list of supported models"""
    return {
        "supported_models": settings.SUPPORTED_MODELS,
        "default_model": settings.DEFAULT_MODEL,
        "model_descriptions": {
            model: get_model_description(model)
            for model in settings.SUPPORTED_MODELS
        }
    }

@router.get("/current")
async def get_current_model():
    """Get currently selected model"""
    return {
        "current_model": settings.DEFAULT_MODEL,
        "is_available": await ollama_service.is_available()
    }

def format_bytes(size_bytes: int) -> str:
    """Format bytes into human readable format"""
    if size_bytes == 0:
        return "0 B"
    
    size_names = ["B", "KB", "MB", "GB", "TB"]
    import math
    i = int(math.floor(math.log(size_bytes, 1024)))
    p = math.pow(1024, i)
    s = round(size_bytes / p, 2)
    return f"{s} {size_names[i]}"

def get_model_capabilities(model_name: str) -> List[str]:
    """Get capabilities for a specific model"""
    capabilities_map = {
        "codellama": [
            "Code generation",
            "Code completion",
            "Bug fixing",
            "Code explanation",
            "Refactoring",
            "Multi-language support"
        ],
        "deepseek-coder": [
            "Advanced code generation",
            "Code understanding",
            "Bug detection",
            "Code optimization",
            "Documentation generation",
            "Multi-language support"
        ],
        "starcoder": [
            "Code completion",
            "Code generation",
            "Cross-language understanding",
            "Documentation",
            "Code translation"
        ],
        "codegemma": [
            "Code generation",
            "Code explanation",
            "Bug fixing",
            "Refactoring",
            "Test generation"
        ]
    }
    
    # Find matching capabilities
    for key, capabilities in capabilities_map.items():
        if key in model_name.lower():
            return capabilities
    
    # Default capabilities
    return [
        "Code assistance",
        "Text generation",
        "Code completion"
    ]

def get_model_description(model_name: str) -> str:
    """Get description for a specific model"""
    descriptions = {
        "codellama:7b-code": "Meta's CodeLlama 7B optimized for code generation and understanding",
        "codellama:13b-code": "Meta's CodeLlama 13B with enhanced code capabilities",
        "deepseek-coder:6.7b": "DeepSeek's code-specialized model with strong programming abilities",
        "starcoder:7b": "BigCode's StarCoder model for code generation and completion",
        "codegemma:7b": "Google's CodeGemma model for code understanding and generation"
    }
    
    return descriptions.get(
        model_name,
        f"Code-specialized language model: {model_name}"
    )