app_kit/model_runtime.py

from __future__ import annotations

import json
import os
import time
from dataclasses import dataclass
from functools import lru_cache
from pathlib import Path
from typing import Any, Sequence

REPO_ROOT = Path(__file__).resolve().parents[1]
MODEL_REGISTRY_PATH = REPO_ROOT / "configs" / "model_registry.yaml"
DEFAULT_COMPONENTS = ("reasoning_llm", "vision_llm", "manual_parser")

_COMPONENT_ENV_VARS: dict[str, tuple[str, ...]] = {
    "reasoning_llm": ("P3_REASONING_MODEL_PATH", "P3_MODEL_PATH"),
    "vision_llm": ("P3_VISION_MODEL_PATH", "P3_MODEL_PATH"),
    "manual_parser": ("P3_MANUAL_PARSER_MODEL_PATH", "P3_MODEL_PATH"),
}


def _json_safe(value: Any) -> Any:
    if isinstance(value, Path):
        return str(value)
    if isinstance(value, dict):
        return {str(key): _json_safe(item) for key, item in value.items()}
    if isinstance(value, (list, tuple)):
        return [_json_safe(item) for item in value]
    return value


@lru_cache(maxsize=1)
def load_model_registry() -> dict[str, Any]:
    try:
        import yaml
    except Exception as exc:  # pragma: no cover - dependency issue is environment-specific
        raise RuntimeError("PyYAML is required to read configs/model_registry.yaml") from exc

    if not MODEL_REGISTRY_PATH.exists():
        raise FileNotFoundError(f"Missing model registry: {MODEL_REGISTRY_PATH}")
    loaded = yaml.safe_load(MODEL_REGISTRY_PATH.read_text(encoding="utf-8"))
    if not isinstance(loaded, dict):
        raise ValueError(f"Invalid model registry format: {MODEL_REGISTRY_PATH}")
    return loaded


@dataclass(frozen=True)
class ComponentSpec:
    component: str
    expected_model_id: str
    backend: str | None
    runtime: str | None


@dataclass(frozen=True)
class ModelAvailability:
    component: str
    expected_model_id: str
    backend: str | None
    runtime: str | None
    available: bool
    resolved_path: Path | None
    checked_paths: tuple[str, ...]
    problem: str

    def to_blocker(self) -> dict[str, Any]:
        return {
            "component": self.component,
            "expected_model_id": self.expected_model_id,
            "backend": self.backend,
            "runtime": self.runtime,
            "available": self.available,
            "resolved_path": str(self.resolved_path) if self.resolved_path else None,
            "checked_paths": list(self.checked_paths),
            "problem": self.problem,
        }


class ModelUnavailableError(RuntimeError):
    def __init__(self, availability: ModelAvailability, *, reason: str | None = None):
        self.availability = availability
        self.reason = reason or availability.problem
        super().__init__(self.reason)

    def to_blocker(self) -> dict[str, Any]:
        payload = self.availability.to_blocker()
        payload["problem"] = self.reason
        return payload


@lru_cache(maxsize=None)
def get_component_spec(component: str) -> ComponentSpec:
    registry = load_model_registry()
    projects = registry.get("projects", {})
    if not isinstance(projects, dict):
        raise ValueError("model_registry.yaml is missing the projects mapping")
    p3 = projects.get("p3", {})
    if not isinstance(p3, dict):
        raise ValueError("model_registry.yaml is missing the projects.p3 mapping")
    spec = p3.get(component)
    if not isinstance(spec, dict):
        raise KeyError(f"Unknown P3 component: {component}")
    model_id = spec.get("model_id")
    if not isinstance(model_id, str) or not model_id.strip():
        raise ValueError(f"Invalid model_id for {component}")
    backend = spec.get("backend")
    runtime = spec.get("runtime")
    return ComponentSpec(
        component=component,
        expected_model_id=model_id.strip(),
        backend=backend.strip() if isinstance(backend, str) and backend.strip() else None,
        runtime=runtime.strip() if isinstance(runtime, str) and runtime.strip() else None,
    )


def _candidate_roots() -> list[Path]:
    roots: list[Path] = []
    for env_var in ("P3_MODEL_CACHE_DIR", "MODEL_CACHE_DIR"):
        value = os.environ.get(env_var)
        if value:
            roots.append(Path(value).expanduser())
    roots.extend([Path("/opt/data/workspace/model-cache"), REPO_ROOT / "models"])
    unique: list[Path] = []
    seen: set[str] = set()
    for root in roots:
        key = str(root)
        if key in seen:
            continue
        seen.add(key)
        unique.append(root)
    return unique


def _candidate_model_paths(component: str, model_id: str) -> list[Path]:
    candidates: list[Path] = []
    env_vars = _COMPONENT_ENV_VARS.get(component, ())
    for env_var in env_vars:
        value = os.environ.get(env_var)
        if value:
            candidates.append(Path(value).expanduser())
    model_name = model_id.strip()
    model_tail = model_name.split("/")[-1]
    for root in _candidate_roots():
        candidates.extend(
            [
                root / model_name,
                root / model_tail,
                root / model_tail.replace("-", "_"),
                root / model_name.replace("/", "-"),
            ]
        )
    unique: list[Path] = []
    seen: set[str] = set()
    for candidate in candidates:
        key = str(candidate)
        if key in seen:
            continue
        seen.add(key)
        unique.append(candidate)
    return unique


@lru_cache(maxsize=None)
def check_component_availability(component: str) -> ModelAvailability:
    spec = get_component_spec(component)
    candidates = _candidate_model_paths(component, spec.expected_model_id)
    for candidate in candidates:
        if candidate.exists():
            return ModelAvailability(
                component=spec.component,
                expected_model_id=spec.expected_model_id,
                backend=spec.backend,
                runtime=spec.runtime,
                available=True,
                resolved_path=candidate,
                checked_paths=tuple(str(path) for path in candidates),
                problem="available",
            )
    problem = f"required model not mounted: {spec.expected_model_id}"
    return ModelAvailability(
        component=spec.component,
        expected_model_id=spec.expected_model_id,
        backend=spec.backend,
        runtime=spec.runtime,
        available=False,
        resolved_path=None,
        checked_paths=tuple(str(path) for path in candidates),
        problem=problem,
    )


def check_required_components(components: Sequence[str] = DEFAULT_COMPONENTS) -> list[ModelAvailability]:
    return [check_component_availability(component) for component in components]


def summarize_blockers(availabilities: Sequence[ModelAvailability]) -> list[dict[str, Any]]:
    return [item.to_blocker() for item in availabilities if not item.available]


def format_blocker_markdown(availabilities: Sequence[ModelAvailability], *, title: str = "Model-backed diagnosis is blocked") -> str:
    blockers = [item for item in availabilities if not item.available]
    if not blockers:
        return ""
    lines = [f"⚠️ **{title}**", "", "The app will not emit a rule-based answer path while the sponsor model requirements are unmet.", "", "Missing components:"]
    for blocker in blockers:
        lines.append(f"- `{blocker.component}` → `{blocker.expected_model_id}`")
    lines.append("")
    lines.append("Checked local paths:")
    for blocker in blockers:
        checked = blocker.checked_paths[:3]
        suffix = " …" if len(blocker.checked_paths) > 3 else ""
        lines.append(f"- `{blocker.component}`: {', '.join(f'`{path}`' for path in checked)}{suffix}")
    return "\n".join(lines)


def _load_llama_cpp():
    try:
        import llama_cpp
    except Exception as exc:  # pragma: no cover - import error is environment-specific
        raise RuntimeError(
            "llama_cpp is required for GGUF-backed inference; install llama-cpp-python in the runtime environment.") from exc
    return llama_cpp


def _load_transformers():
    try:
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer
    except Exception as exc:  # pragma: no cover - import error is environment-specific
        raise RuntimeError(
            "transformers/torch are required for Hugging Face model-backed inference; install them in the runtime environment.") from exc
    return torch, AutoModelForCausalLM, AutoTokenizer


def _extract_generation_stats(response: Any, *, prompt_tokens: int | None = None, completion_tokens: int | None = None, duration_s: float | None = None, load_s: float | None = None, adapter_name: str | None = None, model_path: Path | None = None) -> dict[str, Any]:
    usage = response.get("usage") if isinstance(response, dict) else None
    if isinstance(usage, dict):
        prompt_tokens = prompt_tokens if prompt_tokens is not None else usage.get("prompt_tokens")
        completion_tokens = completion_tokens if completion_tokens is not None else usage.get("completion_tokens")
        total_tokens = usage.get("total_tokens")
    else:
        total_tokens = None
    stats = {
        "prompt_tokens": prompt_tokens,
        "completion_tokens": completion_tokens,
        "total_tokens": total_tokens,
        "duration_s": duration_s,
        "load_s": load_s,
        "adapter_name": adapter_name,
        "model_path": str(model_path) if model_path else None,
    }
    return {key: value for key, value in stats.items() if value is not None}


def generate_text(component: str, prompt: str, *, max_tokens: int = 384, temperature: float = 0.2, top_p: float = 0.9, seed: int = 13) -> tuple[str, dict[str, Any]]:
    availability = check_component_availability(component)
    if not availability.available or availability.resolved_path is None:
        raise ModelUnavailableError(availability)

    model_path = availability.resolved_path
    if model_path.suffix.lower() == ".gguf":
        llama_cpp = _load_llama_cpp()
        started = time.perf_counter()
        llm = llama_cpp.Llama(
            model_path=str(model_path),
            n_ctx=4096,
            seed=seed,
            verbose=False,
        )
        loaded = time.perf_counter()
        kwargs = {
            "max_tokens": max_tokens,
            "temperature": temperature,
            "top_p": top_p,
            "seed": seed,
        }
        try:
            response = llm.create_chat_completion(
                messages=[{"role": "user", "content": prompt}],
                **kwargs,
            )
            choice = response["choices"][0]
            message = choice.get("message") or {}
            text = message.get("content") or ""
        except Exception:
            response = llm(f"{prompt}\n", echo=False, **kwargs)
            choice = response["choices"][0]
            text = choice.get("text") or ""
        if not text.strip():
            raise RuntimeError(f"{component} returned empty text from GGUF model {model_path}")
        stats = _extract_generation_stats(
            response,
            duration_s=round(time.perf_counter() - started, 3),
            load_s=round(loaded - started, 3),
            adapter_name="llama_cpp",
            model_path=model_path,
        )
        return text.strip(), {
            "component": component,
            "model_name": availability.expected_model_id,
            "model_id": availability.expected_model_id,
            "expected_model_id": availability.expected_model_id,
            "adapter_name": "llama_cpp",
            "backend": availability.backend or "llama_cpp",
            "resolved_model_path": str(model_path),
            "generation_stats": stats,
        }

    torch, AutoModelForCausalLM, AutoTokenizer = _load_transformers()
    started = time.perf_counter()
    tokenizer = AutoTokenizer.from_pretrained(str(model_path), trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        str(model_path),
        trust_remote_code=True,
        torch_dtype="auto",
        device_map="auto",
    )
    loaded = time.perf_counter()
    if hasattr(tokenizer, "apply_chat_template"):
        prompt_text = tokenizer.apply_chat_template(
            [{"role": "user", "content": prompt}],
            tokenize=False,
            add_generation_prompt=True,
        )
    else:
        prompt_text = prompt
    inputs = tokenizer(prompt_text, return_tensors="pt")
    try:
        inputs = {key: value.to(model.device) for key, value in inputs.items()}
    except Exception:
        pass
    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            do_sample=temperature > 0,
            pad_token_id=tokenizer.eos_token_id,
        )
    prompt_len = int(inputs["input_ids"].shape[-1])
    completion_ids = output_ids[0][prompt_len:]
    text = tokenizer.decode(completion_ids, skip_special_tokens=True).strip()
    if not text:
        raise RuntimeError(f"{component} returned empty text from transformers model {model_path}")
    stats = _extract_generation_stats(
        {},
        prompt_tokens=prompt_len,
        completion_tokens=int(completion_ids.shape[-1]),
        duration_s=round(time.perf_counter() - started, 3),
        load_s=round(loaded - started, 3),
        adapter_name="transformers",
        model_path=model_path,
    )
    return text, {
        "component": component,
        "model_name": availability.expected_model_id,
        "model_id": availability.expected_model_id,
        "expected_model_id": availability.expected_model_id,
        "adapter_name": "transformers",
        "backend": availability.backend or "transformers",
        "resolved_model_path": str(model_path),
        "generation_stats": stats,
    }


def blocker_payload(availability: Sequence[ModelAvailability], *, status: str = "blocked") -> dict[str, Any]:
    blockers = summarize_blockers(availability)
    return {
        "status": status,
        "blocked_by": blockers,
    }


def stringify_blocked_response(availability: Sequence[ModelAvailability], *, title: str = "Model-backed diagnosis is blocked") -> str:
    body = format_blocker_markdown(availability, title=title)
    if body:
        body += "\n\nNo deterministic fallback will be used until the required model assets are mounted."
    return body