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
# Which server answers behind OLLAMA_HOST: "boilerroom", "llama-server"
# or "ollama". They disagree on tool_choice semantics and on which
# extension fields do anything (see get_tool_choice_settings and
# phase_extra_body), so the health check probes for the wrapper's name
# on /health, then for /props — served by llama-server only — and
# caches the answer here.
_local_flavor: str | 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, _local_flavor

    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", [])]

            # The boilerroom wrapper names itself on /health; a bare
            # llama-server answers /health without a service field and
            # serves /props (its own surface); Ollama answers neither.
            is_wrapper = False
            try:
                health = await client.get(f"{host}/health")
                is_wrapper = (
                    health.status_code == 200 and health.json().get("service") == "boilerroom"
                )
            except (httpx.HTTPError, ValueError):
                is_wrapper = False

            if is_wrapper:
                _local_flavor = "boilerroom"
            else:
                try:
                    props = await client.get(f"{host}/props")
                    _local_flavor = "llama-server" if props.status_code == 200 else "ollama"
                except httpx.HTTPError:
                    _local_flavor = "ollama"

        if model in names or f"{model}:latest" in names:
            _ollama_available = True
            logger.info(
                "ollama_health_check_passed",
                host=host,
                model=model,
                flavor=_local_flavor,
            )
            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 tool calling on the orchestration phase.

    Claude: PydanticAI handles tool_choice natively — no extra_body.

    Ollama: tool_choice="required" via extra_body. Advisory there (Ollama
    ignores it), but it nudges gemma4 to actually call tools, which the
    persona-suppression gotcha made necessary.

    llama-server: NO tool_choice. It enforces "required" on every request
    in the run, so after a tool returns, the next generation is again
    forced to call a tool — an unbreakable tool loop (~80 s turns,
    observed at cutover). Its --jinja template renders tool definitions
    the way gemma4 was trained, and the model calls tools reliably
    unforced (probed 3/3).

    boilerroom: same as llama-server — the wrapper forwards to it, and
    its /props passthrough is what lets this detection keep working.
    """
    from pydantic_ai.settings import ModelSettings

    if resolve_backend() == "claude":
        # PydanticAI's Anthropic model handles tool_choice internally
        return ModelSettings()
    if _local_flavor in ("llama-server", "boilerroom"):
        return ModelSettings()
    # Ollama needs explicit tool_choice via extra_body
    return ModelSettings(extra_body={"tool_choice": "required"})


# Phase bindings for the local backend (this repo's T-4, decided
# 2026-09-12): through the boilerroom wrapper each phase is a named
# session with an eviction rank — lower parks sooner, so the most-used
# context keeps its KV warm. The pipeline phases outrank the librarian,
# which outranks webber (20, filed in its own repo); Open WebUI stays
# session-less and can never evict anyone (wrapper D-4).
# phase -> (direct llama-server slot, wrapper session name, eviction rank)
_PHASE_BINDINGS: dict[str, tuple[int | None, str, int]] = {
    "steward": (0, "tatlock-steward", 40),
    "orchestrate": (1, "tatlock-orchestrate", 40),
    "synthesize": (2, "tatlock-synthesize", 40),
    # Experts (biographer, housekeeper) share one session between the
    # librarian and the phases: the two least-used consumers trade the
    # spare slot by rank instead of anyone hitting an empty pool — a
    # session-less call 503s once four sessions are resident (D-4's
    # idle-only rule), so every generation call here must carry a name.
    "experts": (None, "tatlock-experts", 35),
    "librarian": (None, "librarian", 30),
}


def phase_extra_body(phase: str) -> dict[str, str | int]:
    """
    The extension fields a pipeline phase sends to the local backend.

    Through the boilerroom wrapper: a named session with its eviction
    rank — the wrapper owns the name-to-slot map and reports parks and
    restores on the response. Against a bare llama-server: the raw
    id_slot pin this replaces (the librarian floats there, as it always
    did). Ollama and Claude get nothing, and an unprobed flavor sends
    nothing rather than guessing. Gated by BACKEND_SLOT_PINNING like
    the pinning it grew out of.
    """
    if not config.BACKEND_SLOT_PINNING or resolve_backend() == "claude":
        return {}
    slot, session, rank = _PHASE_BINDINGS[phase]
    if _local_flavor == "boilerroom":
        return {"session": session, "eviction_order": rank}
    if _local_flavor == "llama-server" and slot is not None:
        return {"id_slot": slot}
    return {}


def with_phase_binding(settings: ModelSettings | None, phase: str) -> ModelSettings | None:
    """
    Merge a phase's local-backend binding into model settings.

    Session fields through the wrapper, id_slot against a bare
    llama-server, settings untouched everywhere else — one call site
    stays correct on any backend (see phase_extra_body).
    """
    extra = phase_extra_body(phase)
    if not extra:
        return settings

    from pydantic_ai.settings import ModelSettings

    merged = dict(settings or {})
    extra_body = dict(merged.get("extra_body") or {})
    extra_body.update(extra)
    merged["extra_body"] = extra_body
    return ModelSettings(**merged)


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,
        "local_flavor": _local_flavor,
        "prefer_cloud": config.PREFER_CLOUD_BACKEND,
    }