tatlock/src/anthropic/model_selector.py

"""
Model selector for Ollama/Claude backend switching.

Provides automatic model selection with Ollama as the primary local backend
and Claude as the cloud fallback. Claude is used when PREFER_CLOUD_BACKEND
is enabled, or automatically when Ollama is unavailable at startup.

The Anthropic SDK is imported lazily so a missing or broken `anthropic`
package degrades to Ollama-only operation instead of crashing the app.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import httpx

from src.core.config import config
from src.core.logging_config import get_logger

if TYPE_CHECKING:
    from pydantic_ai.models.anthropic import AnthropicModel
    from pydantic_ai.models.openai import OpenAIChatModel
    from pydantic_ai.settings import ModelSettings

logger = get_logger(__name__)

# Cached health check results (set once at startup)
_claude_available: bool | None = None
_ollama_available: bool | None = None


async def check_ollama_health() -> bool:
    """
    Check if the Ollama server is reachable and has the configured model.

    This should be called once at application startup.
    The result is cached in `_ollama_available`.

    Returns:
        True if Ollama is reachable and OLLAMA_DEFAULT_MODEL is pulled.
    """
    global _ollama_available

    host = str(config.OLLAMA_HOST).rstrip("/")
    model = config.OLLAMA_DEFAULT_MODEL

    try:
        async with httpx.AsyncClient(timeout=5.0) as client:
            # OpenAI-compat surface, so the check holds for any backend
            # behind OLLAMA_HOST (Ollama, llama-server).
            response = await client.get(f"{host}/v1/models")
            response.raise_for_status()
            names = [m.get("id", "") for m in response.json().get("data", [])]

        if model in names or f"{model}:latest" in names:
            _ollama_available = True
            logger.info(
                "ollama_health_check_passed",
                host=host,
                model=model,
            )
            return True

        _ollama_available = False
        logger.warning(
            "ollama_health_check_failed",
            reason="model_not_served",
            host=host,
            model=model,
            hint=f"backend does not list '{model}' — pull it (ollama) or serve it under that alias (llama-server)",
        )
        return False

    except Exception as e:
        _ollama_available = False
        logger.warning(
            "ollama_health_check_failed",
            reason="server_unreachable",
            host=host,
            error=str(e),
        )
        return False


async def check_claude_health() -> bool:
    """
    Check if Claude API is reachable and working.

    This should be called once at application startup.
    The result is cached in `_claude_available`.

    Returns:
        True if Claude API is accessible, False otherwise.
    """
    global _claude_available

    # No API key configured - Claude not available
    if not config.ANTHROPIC_API_KEY:
        logger.info(
            "claude_health_check_skipped",
            reason="no_api_key",
        )
        _claude_available = False
        return False

    try:
        from anthropic import AsyncAnthropic

        client = AsyncAnthropic(api_key=config.ANTHROPIC_API_KEY)

        # Minimal API call to verify connectivity
        # Using a tiny max_tokens to minimize cost
        await client.messages.create(
            model=config.ANTHROPIC_MODEL,
            max_tokens=1,
            messages=[{"role": "user", "content": "hi"}],
        )

        _claude_available = True
        logger.info(
            "claude_health_check_passed",
            model=config.ANTHROPIC_MODEL,
        )
        return True

    except Exception as e:
        _claude_available = False
        logger.warning(
            "claude_health_check_failed",
            error=str(e),
            model=config.ANTHROPIC_MODEL,
        )
        return False


def is_claude_available() -> bool:
    """
    Check if Claude is available (from cached health check result).

    Returns:
        True if Claude API was reachable at startup, False otherwise.

    Note:
        Returns False if health check hasn't been run yet.
        Call `check_claude_health()` at startup first.
    """
    return _claude_available is True


def is_ollama_available() -> bool:
    """
    Check if Ollama is available (from cached health check result).

    Returns:
        False only if the startup health check confirmed Ollama is down.
        Unknown (check not run yet) counts as available so that contexts
        without lifespan events keep the local-first behavior.
    """
    return _ollama_available is not False


def resolve_backend(prefer_cloud: bool | None = None) -> str:
    """
    Resolve which backend should serve requests.

    Ollama is the primary backend. Claude is used when explicitly
    preferred via PREFER_CLOUD_BACKEND, or as automatic fallback
    when the startup health check found Ollama down.

    Args:
        prefer_cloud: Override config.PREFER_CLOUD_BACKEND for this call.

    Returns:
        "claude" or "ollama".
    """
    use_cloud = prefer_cloud if prefer_cloud is not None else config.PREFER_CLOUD_BACKEND

    if use_cloud and is_claude_available():
        return "claude"

    if not is_ollama_available() and is_claude_available():
        logger.warning(
            "backend_fallback_to_claude",
            reason="ollama_unavailable",
        )
        return "claude"

    return "ollama"


def get_model(prefer_cloud: bool | None = None) -> AnthropicModel | OpenAIChatModel:
    """
    Get the best available model.

    Returns Ollama unless Claude is preferred (or Ollama is down).

    Args:
        prefer_cloud: Override config.PREFER_CLOUD_BACKEND for this call.
                     If None, uses the config value.

    Returns:
        PydanticAI model instance (OpenAIChatModel or AnthropicModel).

    Example:
        >>> model = get_model()
        >>> agent = Agent(model, system_prompt="...")
    """
    if resolve_backend(prefer_cloud) == "claude":
        try:
            from pydantic_ai.models.anthropic import AnthropicModel
            from pydantic_ai.providers.anthropic import AnthropicProvider

            logger.debug(
                "model_selected",
                backend="claude",
                model=config.ANTHROPIC_MODEL,
            )
            return AnthropicModel(
                model_name=config.ANTHROPIC_MODEL,
                provider=AnthropicProvider(api_key=config.ANTHROPIC_API_KEY),
            )
        except ImportError as e:
            logger.error(
                "claude_backend_import_failed",
                error=str(e),
                hint="anthropic package missing or incompatible; using Ollama",
            )

    from pydantic_ai.models.openai import OpenAIChatModel

    from src.ollama.provider import get_ollama_provider

    logger.debug(
        "model_selected",
        backend="ollama",
        model=config.OLLAMA_DEFAULT_MODEL,
    )
    return OpenAIChatModel(
        model_name=config.OLLAMA_DEFAULT_MODEL,
        provider=get_ollama_provider(),
    )


def get_tool_choice_settings() -> ModelSettings:
    """
    Get model_settings for forcing tool calls on the first request.

    For Claude: PydanticAI handles tool_choice natively, so no extra_body needed.
    For Ollama: Pass tool_choice="required" via extra_body to force tool calling.
    """
    from pydantic_ai.settings import ModelSettings

    if resolve_backend() == "claude":
        # PydanticAI's Anthropic model handles tool_choice internally
        return ModelSettings()
    else:
        # Ollama needs explicit tool_choice via extra_body
        return ModelSettings(extra_body={"tool_choice": "required"})


def get_sampling_settings(temperature: float) -> ModelSettings:
    """
    Get model_settings with a sampling temperature where the backend allows it.

    Ollama accepts a temperature; Claude Sonnet 5+ rejects sampling
    parameters, so the Claude backend gets empty settings.
    """
    from pydantic_ai.settings import ModelSettings

    if resolve_backend() == "claude":
        return ModelSettings()
    return ModelSettings(temperature=temperature)


def get_model_info() -> dict:
    """
    Get information about the current model configuration.

    Useful for health checks and debugging.

    Returns:
        Dict with backend, model name, and availability info.
    """
    backend = resolve_backend()

    return {
        "backend": backend,
        "model": config.ANTHROPIC_MODEL if backend == "claude" else config.OLLAMA_DEFAULT_MODEL,
        "claude_available": is_claude_available(),
        "claude_configured": bool(config.ANTHROPIC_API_KEY),
        "ollama_available": is_ollama_available(),
        "ollama_model": config.OLLAMA_DEFAULT_MODEL,
        "prefer_cloud": config.PREFER_CLOUD_BACKEND,
    }