models.py

#!/usr/bin/env python3
"""
Model handling for dejavu2-cli.

This module handles loading and selecting LLM models from Models.json,
as well as listing and displaying model information.
"""

import json
import logging
from pathlib import Path
from typing import Any

import click

# Import custom exceptions
from errors import ConfigurationError, ModelError

# Configure module logger
logger = logging.getLogger(__name__)

# Module-level cache for models.json to avoid redundant file I/O
_models_cache: dict[str, dict[str, Any]] = {}
_cache_mtime: dict[str, float] = {}


def load_models_json(json_file: str, force_reload: bool = False) -> dict[str, Any]:
  """
  Load models from JSON file with caching.

  Uses module-level caching with mtime checking to avoid redundant file I/O.
  The cache is automatically invalidated when the file is modified.

  Args:
      json_file: Path to the Models.json file
      force_reload: If True, bypass cache and reload from disk

  Returns:
      Dictionary of models loaded from the JSON file

  Raises:
      ConfigurationError: If the models file cannot be found or parsed
  """
  global _models_cache, _cache_mtime

  # Load from file
  logger.debug(f"Loading models from: {json_file}")

  try:
    # Check if file has been modified
    current_mtime = Path(json_file).stat().st_mtime

    # Use cached version if available and file hasn't changed
    if not force_reload and json_file in _models_cache:
      if _cache_mtime.get(json_file) == current_mtime:
        logger.debug(f"Using cached models from {json_file}")
        return _models_cache[json_file]

    # Load from file
    with open(json_file, encoding="utf-8") as file:
      models = json.load(file)

    # Update cache
    _models_cache[json_file] = models
    _cache_mtime[json_file] = current_mtime

    logger.debug(f"Successfully loaded and cached {len(models)} models from file")
    return models

  except FileNotFoundError as e:
    error_msg = f"Models file not found: {json_file}"
    logger.error(error_msg)
    raise ConfigurationError(error_msg) from e
  except json.JSONDecodeError as e:
    error_msg = f"Invalid JSON in models file '{json_file}': {str(e)}"
    logger.error(error_msg)
    raise ConfigurationError(error_msg) from e
  except Exception as e:
    error_msg = f"Error loading models file '{json_file}': {str(e)}"
    logger.error(error_msg)
    raise ConfigurationError(error_msg) from e


def list_available_canonical_models(json_file: str) -> list[str]:
  """
  List all available model names from Models.json.

  Args:
      json_file: Path to the Models.json file

  Returns:
      List of available model names sorted alphabetically

  Raises:
      ConfigurationError: If the models file cannot be found or parsed
  """
  # Use cached loading
  models = load_models_json(json_file)

  # Extract canonical model names where 'available' is not 0
  canonical_names = [name for name, details in models.items() if details.get("available") != 0]

  # Sort the canonical model names
  canonical_names.sort()

  logger.debug(f"Found {len(canonical_names)} available models out of {len(models)} total")
  return canonical_names


def list_available_canonical_models_with_details(json_file: str) -> dict[str, Any]:
  """
  Get a dictionary of all available models with their complete details.

  Args:
      json_file: Path to the Models.json file

  Returns:
      Dictionary mapping model names to their complete details for available models

  Raises:
      ConfigurationError: If the models file cannot be found or parsed
  """
  # Use cached loading
  models = load_models_json(json_file)

  # Extract models where 'available' is greater than 0
  available_models = {name: details for name, details in models.items() if details.get("available") > 0}
  return available_models


def get_canonical_model(model_name: str, json_file: str) -> tuple[str, dict[str, Any]]:
  """
  Get canonical model name and load model parameters.

  Searches Models.json for either an exact match on model name or a match
  on the alias field. When found, loads parameters.

  Args:
      model_name: The model name or alias to look up
      json_file: Path to the Models.json file

  Returns:
      Tuple of (canonical_model_name, model_parameters)

  Raises:
      ConfigurationError: If Models.json cannot be found or contains invalid JSON
      ModelError: If the requested model is not found
  """
  logger.debug(f"Looking up model: '{model_name}' in {json_file}")
  model_parameters = {}

  # Use cached loading
  models = load_models_json(json_file)
  logger.debug(f"Successfully loaded {len(models)} models from {json_file}")

  # Check if the model name is a canonical name
  canonical_name = None
  if model_name in models:
    canonical_name = model_name
    logger.debug(f"Found exact match for model name: '{model_name}'")
  else:
    logger.debug(f"Model '{model_name}' not found directly, searching aliases")
    # Search through aliases
    for name, details in models.items():
      if details.get("alias") == model_name:
        logger.debug(f"Found alias match: '{model_name}' → '{name}'")

        if details.get("available") == 0:
          error_msg = f"Alias '{model_name}' was found but is unavailable"
          logger.warning(error_msg)
          click.echo(f"Warning: {error_msg}", err=True)
          return None, {}

        if details.get("enabled") == 0:
          error_msg = f"Alias '{model_name}' was found but is not enabled"
          logger.warning(error_msg)
          click.echo(f"Warning: {error_msg}", err=True)
          return None, {}

        canonical_name = name
        break

  if not canonical_name:
    # Model name not found
    error_msg = f"Model '{model_name}' not found in models file"
    logger.warning(error_msg)
    raise ModelError(error_msg)

  # Initialize model_parameters from Models.json - include all model information
  model_info = models.get(canonical_name, {})
  logger.debug(f"Found model info for '{canonical_name}' with {len(model_info)} parameters")

  # Essential fields that must be present
  required_fields = ["model", "series", "url", "apikey", "context_window", "max_output_tokens", "available", "enabled"]

  # Copy all model info to model_parameters
  for field, value in model_info.items():
    model_parameters[field] = value

  # Verify required fields exist (set to None if missing)
  missing_fields = []
  for field in required_fields:
    if field not in model_parameters:
      model_parameters[field] = None
      missing_fields.append(field)

  if missing_fields:
    logger.warning(f"Model '{canonical_name}' is missing required fields: {', '.join(missing_fields)}")

  logger.info(f"Selected model: {canonical_name} ({model_parameters.get('provider', 'Unknown provider')})")
  return canonical_name, model_parameters


def list_models(json_file: str, details: bool = False) -> None:
  """
  List all available LLM models as defined in Models.json.

  Args:
      json_file: Path to the Models.json file
      details: If True, shows detailed information about each model.
              If False, shows only the model names.
  """
  if details:
    models = list_available_canonical_models_with_details(json_file)
    for name, details in sorted(models.items()):
      click.echo(f"Model: {name}")
      for key, value in details.items():
        click.echo(f"  {key}: {value}")
      click.echo()
  else:
    models = list_available_canonical_models(json_file)
    for name in models:
      click.echo(f"{name}")