core/ollama_models.py

"""Helper functions to query Ollama API for available models."""

import logging
from typing import Any

import requests

logger = logging.getLogger(__name__)


def fetch_ollama_models(api_base: str) -> list[str]:
    """
    Fetch available models from Ollama API.
    
    Args:
        api_base: Ollama API base URL (e.g., "http://localhost:11434")
        
    Returns:
        List of model names available on the Ollama server
    """
    try:
        # Ollama API endpoint for listing models
        url = f"{api_base.rstrip('/')}/api/tags"
        response = requests.get(url, timeout=10)
        response.raise_for_status()
        
        data = response.json()
        models = data.get("models", [])
        
        # Extract model names
        model_names = [model.get("name", "") for model in models if model.get("name")]
        
        logger.info(f"Fetched {len(model_names)} models from Ollama at {api_base}")
        return model_names
        
    except requests.exceptions.RequestException as e:
        logger.error(f"Error fetching models from Ollama: {e}")
        return []
    except Exception as e:
        logger.error(f"Unexpected error fetching models from Ollama: {e}")
        return []