src/nim_fallback.py

"""Centralized V1 model fallback with health-checking (NVIDIA NIM / OpenAI-compat).

The configured model can be dead (e.g. z-ai/glm-5.1 → 410 EOL). This module
health-checks a chain of models with a small STRUCTURED task and returns the
first model that actually works end-to-end (HTTP 200, valid JSON, schema-valid,
JD-traceable, no hallucinated phrase, within timeout). If none pass, callers get
`None` and must preserve the résumé and report `live_model_unavailable`.

Never combines partial outputs from different models: a failover re-runs the
whole operation on the replacement model.
"""
from __future__ import annotations

import time
from typing import List, Optional, Tuple

# Ordered preference. Kept centralized so every V1 route shares the same chain.
try:
    from config import V1_MODEL_CHAIN as _CFG_CHAIN
except Exception:
    _CFG_CHAIN = None

DEFAULT_MODEL_CHAIN: List[str] = _CFG_CHAIN or [
    "z-ai/glm-5.2",
    "mistralai/mistral-small-4-119b-2603",
    "nvidia/nemotron-3-super-120b-a12b",
]

# A tiny probe JD with a known exact phrase and a decoy that must NOT be invented.
_PROBE_JD = (
    "About the Role\nWe are hiring a Product Manager.\n"
    "Requirements: strong stakeholder management and SQL. 5+ years experience.\n"
    "Responsibilities: roadmap prioritization and A/B testing."
)
_PROBE_MUST_TRACE = "stakeholder management"


def health_check(model: str, timeout: float = 25.0) -> dict:
    """Probe one model with the real structured-extraction task and validate the
    output the same way production does. Returns a log dict (never raises)."""
    from .llm_client import LLMClient
    from .keyword_schema import validate_and_repair

    log = {
        "requested_model": model, "selected_model": None, "status": "unknown",
        "latency": 0.0, "retry_count": 0, "failure_reason": "",
        "schema_valid": False, "evidence_test_passed": False,
    }
    t0 = time.monotonic()
    try:
        # Rebuild the OpenAI client with a SHORT per-request timeout so a hung
        # model fails fast instead of blocking on the default 90s × retries.
        from openai import OpenAI
        from config import NVIDIA_API_KEY, GLM_BASE_URL
        client = LLMClient(model=model)
        client.client = OpenAI(base_url=GLM_BASE_URL, api_key=NVIDIA_API_KEY,
                               timeout=timeout, max_retries=0)
        sys_p = ("Extract ATS keywords. Return ONLY a JSON array of objects with "
                 "keys exact_phrase, normalized_concept, category, requirement_type, "
                 "importance, source_text, semantic_variants, confidence, "
                 "requires_resume_evidence. No prose.")
        rawtext = client._call(sys_p, f"JD:\n{_PROBE_JD}", max_tokens=1800, retries=2)
        raw = client._extract_json(rawtext)
        log["latency"] = round(time.monotonic() - t0, 2)
        if not isinstance(raw, list) or not raw:
            log["status"] = "empty_or_invalid"
            log["failure_reason"] = "no structured items returned"
            return log
        valid, rejected = validate_and_repair(raw, _PROBE_JD)
        log["schema_valid"] = bool(valid)
        # traceability: the known phrase must be extracted; nothing untraceable kept
        traced = any(_PROBE_MUST_TRACE in (v.get("exact_phrase", "").lower()
                                           + v.get("normalized_concept", "").lower())
                     for v in valid)
        log["evidence_test_passed"] = bool(traced and valid)
        if valid and traced:
            log["status"] = "healthy"
            log["selected_model"] = model
        else:
            log["status"] = "evidence_failure"
            log["failure_reason"] = "probe phrase not traceably extracted"
    except Exception as e:
        log["latency"] = round(time.monotonic() - t0, 2)
        msg = str(e)
        log["failure_reason"] = msg[:160]
        # classify a few common transient/terminal signals
        if "410" in msg:
            log["status"] = "gone_410"
        elif "404" in msg:
            log["status"] = "not_found_404"
        elif "429" in msg:
            log["status"] = "rate_limited_429"
        elif "timeout" in msg.lower():
            log["status"] = "timeout"
        else:
            log["status"] = "error"
    return log


# Module cache: once a model is confirmed healthy, reuse it across requests rather
# than re-probing (which would pay the hung-model timeout every call). If that
# model later fails during real use, callers fall back to deterministic (safe);
# call reset_cache() to force a fresh probe.
_CACHE: dict = {"model": None, "logs": []}


def reset_cache() -> None:
    _CACHE["model"] = None
    _CACHE["logs"] = []


def extraction_gaps(clean_jd: str, valid_items: list) -> list:
    """Deterministic completeness check: obvious source-grounded requirements the
    extraction should have captured but didn't. Returns a list of missed phrases
    (named tools present in the JD, an explicit years-of-experience requirement).
    Used to trigger a correction retry / failover before trusting an extraction."""
    import re
    try:
        from .jd_analyzer import _TOOLS
    except Exception:
        _TOOLS = set()
    low = (clean_jd or "").lower()
    got = " ".join((v.get("exact_phrase", "") + " " + v.get("normalized_concept", ""))
                   for v in (valid_items or [])).lower()
    missed = []
    # named tools/technologies present verbatim in the JD but absent from criteria
    for tool in _TOOLS:
        if re.search(r"(?<![a-z0-9])" + re.escape(tool) + r"(?![a-z0-9])", low) \
                and tool not in got:
            missed.append(tool)
    # explicit years-of-experience requirement (e.g. "10+ years")
    ym = re.search(r"\b(\d{1,2}\+?\s*years?)\b", low)
    if ym and "year" not in got:
        missed.append(ym.group(1))
    return missed[:12]


def select_model(chain: Optional[List[str]] = None,
                 timeout: float = 25.0, use_cache: bool = True) -> Tuple[Optional[str], List[dict]]:
    """Return (first_healthy_model | None, per-model health logs)."""
    if use_cache and _CACHE["model"]:
        return _CACHE["model"], _CACHE["logs"]
    chain = chain or DEFAULT_MODEL_CHAIN
    logs: List[dict] = []
    for model in chain:
        log = health_check(model, timeout=timeout)
        logs.append(log)
        if log["status"] == "healthy":
            if use_cache:
                _CACHE["model"], _CACHE["logs"] = model, logs
            return model, logs
    return None, logs


def build_llm(chain: Optional[List[str]] = None, timeout: float = 25.0,
              use_cache: bool = True):
    """Return (LLMClient bound to a healthy model | None, health logs). None means
    every model failed → caller must preserve résumé + report live_model_unavailable."""
    from .llm_client import LLMClient
    model, logs = select_model(chain, timeout=timeout, use_cache=use_cache)
    if model is None:
        return None, logs
    return LLMClient(model=model), logs


if __name__ == "__main__":  # live health-check (honest: reports real availability)
    import json
    model, logs = select_model()
    print("SELECTED:", model)
    for lg in logs:
        print(json.dumps({k: lg[k] for k in
              ("requested_model", "status", "latency", "schema_valid",
               "evidence_test_passed", "failure_reason")}, ensure_ascii=False))