src/llm.py

"""LLM factory with access tiers and provider fallback.

Tiers (cost-control layer for public deployment):
  community — open models (Gemma 4 / DeepSeek V4 / GLM 5.2) served via the
              Hugging Face Inference Providers router. Needs HF_TOKEN. This is
              the default for anonymous visitors, so premium API keys are
              never spent by the public.
  premium   — the frontier pipeline (Gemini / GPT / Claude) from AGENT_MODELS,
              unlocked per-session with OWNER_PASSCODE.

The active tier is carried in a contextvar set by the request handler, so the
agents just call get_llm(role) and inherit the caller's tier. Each tier
degrades gracefully into the other when its credentials are missing.
"""

from __future__ import annotations

import logging
from contextlib import contextmanager
from contextvars import ContextVar
from functools import lru_cache

from config import (
    AGENT_MODELS,
    COMMUNITY_MODELS,
    FALLBACK_MODELS,
    HF_ROUTER_BASE_URL,
    HF_TOKEN,
    PROVIDER_KEYS,
    available_providers,
    community_available,
)

logger = logging.getLogger(__name__)

_TIER: ContextVar[str] = ContextVar("llm_tier", default="community")


def current_tier() -> str:
    return _TIER.get()


@contextmanager
def use_tier(tier: str):
    """Scope all get_llm() calls inside the block to the given tier."""
    token = _TIER.set(tier if tier in ("community", "premium") else "community")
    try:
        yield
    finally:
        _TIER.reset(token)


def _build_premium(provider: str, model: str, temperature: float | None):
    kwargs = {} if temperature is None else {"temperature": temperature}
    if provider == "openai":
        from langchain_openai import ChatOpenAI
        return ChatOpenAI(model=model, **kwargs)
    if provider == "google":
        from langchain_google_genai import ChatGoogleGenerativeAI
        return ChatGoogleGenerativeAI(model=model, **kwargs)
    if provider == "anthropic":
        from langchain_anthropic import ChatAnthropic
        return ChatAnthropic(model=model, max_tokens=8192, **kwargs)
    raise ValueError(f"Unknown provider: {provider}")


def _build_community(model: str, temperature: float | None):
    """Open model via the HF Inference Providers router (OpenAI-compatible)."""
    from langchain_openai import ChatOpenAI
    kwargs = {} if temperature is None else {"temperature": temperature}
    return ChatOpenAI(model=model, api_key=HF_TOKEN,
                      base_url=HF_ROUTER_BASE_URL, max_tokens=4096, **kwargs)


@lru_cache(maxsize=None)
def _get_llm(role: str, tier: str):
    spec = AGENT_MODELS[role]

    if tier == "community":
        if community_available():
            return _build_community(COMMUNITY_MODELS[role], spec.temperature)
        if available_providers():  # private use without HF_TOKEN: use premium keys
            logger.warning("Community tier requested but HF_TOKEN missing; "
                           "using premium keys for role %r", role)
            return _get_llm(role, "premium")
        raise RuntimeError(
            "No credentials for the community tier. Set HF_TOKEN (Hugging Face "
            "Inference Providers) or a premium provider key in the environment."
        )

    # premium tier
    provider, model = spec.provider, spec.model
    if not PROVIDER_KEYS.get(provider):
        candidates = available_providers()
        if candidates:
            provider = candidates[0]
            model = FALLBACK_MODELS[provider]
            logger.warning("Role %r: provider %r unavailable, falling back to %s/%s",
                           role, spec.provider, provider, model)
        elif community_available():
            logger.warning("Premium tier requested but no provider keys; "
                           "serving role %r from the community tier", role)
            return _build_community(COMMUNITY_MODELS[role], spec.temperature)
        else:
            raise RuntimeError(
                "No LLM credentials found. Set at least one of OPENAI_API_KEY, "
                "GOOGLE_API_KEY, ANTHROPIC_API_KEY, or HF_TOKEN."
            )
    return _build_premium(provider, model, spec.temperature)


def get_llm(role: str):
    """Chat model for an agent role, resolved under the caller's active tier."""
    return _get_llm(role, current_tier())


def describe_routing(tier: str = "community") -> dict[str, str]:
    """Human-readable role -> model map for the UI header."""
    out = {}
    for role, spec in AGENT_MODELS.items():
        if tier == "premium":
            if PROVIDER_KEYS.get(spec.provider):
                out[role] = f"{spec.provider}/{spec.model}"
            elif available_providers():
                p = available_providers()[0]
                out[role] = f"{p}/{FALLBACK_MODELS[p]} (fallback)"
            elif community_available():
                out[role] = f"hf/{COMMUNITY_MODELS[role]} (fallback)"
            else:
                out[role] = f"{spec.provider}/{spec.model} · key needed"
        else:
            if community_available():
                out[role] = f"hf/{COMMUNITY_MODELS[role]}"
            elif available_providers():
                p = available_providers()[0]
                out[role] = f"{p}/{FALLBACK_MODELS[p]} (fallback)"
            else:
                out[role] = f"hf/{COMMUNITY_MODELS[role]} · HF_TOKEN needed"
    return out