app_kit/model_runtime.py

from __future__ import annotations

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

DEFAULT_MODEL_REPO_ID = "Abiray/MiniCPM5-1B-GGUF"
DEFAULT_MODEL_FILENAME = "minicpm5-1b-Q4_K_M.gguf"
DEFAULT_MODEL_ID = "Abiray/MiniCPM5-1B-GGUF:Q4_K_M"
DEFAULT_MODEL_CONTEXT = 4096


@dataclass(frozen=True)
class LoadedModel:
    model_id: str
    model_path: Path
    source: str
    backend: str = "llama-cpp-python"


def _repo_root() -> Path:
    return Path(__file__).resolve().parents[1]


def _candidate_roots() -> list[Path]:
    roots: list[Path] = []
    env_cache = os.environ.get("MODEL_CACHE_DIR")
    if env_cache:
        roots.append(Path(env_cache).expanduser())
    roots.append(_repo_root() / "models")
    roots.append(Path("/opt/data/workspace/model-cache"))
    roots.append(Path("/opt/data/model-cache"))
    roots.append(Path.home() / ".cache" / "huggingface" / "hub")
    return roots


def _resolve_from_roots(filename: str) -> tuple[Path | None, str | None]:
    patterns = [
        filename,
        filename.lower(),
        filename.upper(),
        "*MiniCPM5-1B*Q4_K_M*.gguf",
        "*minicpm5-1b*Q4_K_M*.gguf",
        "*MiniCPM5-1B*.gguf",
        "*minicpm5-1b*.gguf",
    ]
    for root in _candidate_roots():
        if not root.exists():
            continue
        for pattern in patterns:
            for candidate in root.rglob(pattern):
                if candidate.is_file():
                    return candidate, f"local-cache:{root}"
    return None, None


def resolve_model_path(*, model_id: str = DEFAULT_MODEL_ID, repo_id: str = DEFAULT_MODEL_REPO_ID, filename: str = DEFAULT_MODEL_FILENAME, env_var: str = "P1_MODEL_PATH") -> LoadedModel:
    explicit = os.environ.get(env_var, "").strip()
    if explicit:
        path = Path(explicit).expanduser()
        if path.exists():
            return LoadedModel(model_id=model_id, model_path=path, source=f"env:{env_var}")
        raise FileNotFoundError(f"{env_var} points to missing model path: {path}")

    cached, source = _resolve_from_roots(filename)
    if cached is not None:
        return LoadedModel(model_id=model_id, model_path=cached, source=source or "local-cache")

    allow_download = os.environ.get("P1_ALLOW_MODEL_DOWNLOAD", "1").strip().lower() not in {"0", "false", "no"}
    if not allow_download:
        raise FileNotFoundError(
            f"Missing model checkpoint for {model_id}. Set {env_var} or place {filename} in MODEL_CACHE_DIR."
        )

    try:
        from huggingface_hub import hf_hub_download
    except Exception as exc:  # pragma: no cover - exercised in environments without the dependency
        raise RuntimeError(
            f"Could not import huggingface_hub to download {model_id}; install huggingface_hub or mount the model locally."
        ) from exc

    cache_dir = _candidate_roots()[0]
    cache_dir.mkdir(parents=True, exist_ok=True)
    try:
        downloaded = hf_hub_download(
            repo_id=repo_id,
            filename=filename,
            local_dir=str(cache_dir),
            local_dir_use_symlinks=False,
        )
    except Exception as exc:
        raise RuntimeError(
            f"Failed to download {model_id} from {repo_id}/{filename}. Mount a local checkpoint or pre-download the model."
        ) from exc

    downloaded_path = Path(downloaded)
    if not downloaded_path.exists():
        raise RuntimeError(f"Download for {model_id} completed but file is missing: {downloaded_path}")
    return LoadedModel(model_id=model_id, model_path=downloaded_path, source=f"huggingface:{repo_id}")


@lru_cache(maxsize=2)
def load_llama(model_path: str, n_ctx: int = DEFAULT_MODEL_CONTEXT):
    try:
        from llama_cpp import Llama
    except Exception as exc:  # pragma: no cover - import is exercised in runtime smoke tests
        raise RuntimeError(
            "llama-cpp-python is required for P1 model inference; install it in the runtime environment."
        ) from exc
    return Llama(model_path=model_path, n_ctx=n_ctx, verbose=False)


def _extract_json_object(text: str) -> dict[str, Any]:
    start = text.find("{")
    end = text.rfind("}")
    if start == -1 or end == -1 or end <= start:
        raise RuntimeError("Model output did not contain a JSON object")
    raw = text[start : end + 1]
    try:
        payload = json.loads(raw)
    except Exception as exc:
        raise RuntimeError(f"Failed to parse JSON from model output: {exc}") from exc
    if not isinstance(payload, dict):
        raise RuntimeError("Model output JSON must be an object")
    return payload


def _require_text(payload: dict[str, Any], field: str) -> str:
    value = payload.get(field)
    if not isinstance(value, str):
        raise RuntimeError(f"Model output missing required '{field}' field")
    value = value.strip()
    if not value:
        raise RuntimeError(f"Model output field '{field}' was empty")
    return value


def generate_text_completion(
    *,
    llm,
    model: LoadedModel,
    system_prompt: str,
    user_prompt: str,
    temperature: float = 0.2,
    max_tokens: int = 256,
) -> tuple[str, dict[str, Any]]:
    started_at = time.perf_counter()
    prompt = f"{system_prompt.strip()}\n\n{user_prompt.strip()}\n\n### Response\n"
    if hasattr(llm, 'create_chat_completion'):
        response = llm.create_chat_completion(
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
            temperature=temperature,
            max_tokens=max_tokens,
        )
        message = str(response["choices"][0]["message"]["content"]).strip()
    else:
        response = llm.create_completion(prompt=prompt, temperature=temperature, max_tokens=max_tokens)
        message = str(response["choices"][0].get("text", "")).strip()
    usage = response.get("usage") or {}
    generation_stats = {
        "prompt_tokens": int(usage.get("prompt_tokens", 0) or 0),
        "completion_tokens": int(usage.get("completion_tokens", 0) or 0),
        "total_tokens": int(usage.get("total_tokens", 0) or 0),
        "elapsed_ms": round((time.perf_counter() - started_at) * 1000.0, 2),
        "backend": "llama-cpp-python",
        "model_path": str(model.model_path),
        "n_ctx": DEFAULT_MODEL_CONTEXT,
    }
    meta = {
        "model_id": model.model_id,
        "model_path": str(model.model_path),
        "model_source": model.source,
        "backend": model.backend,
        "generation_stats": generation_stats,
    }
    return message, meta


def generate_json_completion(
    *,
    llm,
    model: LoadedModel,
    system_prompt: str,
    user_prompt: str,
    temperature: float = 0.2,
    max_tokens: int = 512,
) -> tuple[dict[str, Any], dict[str, Any]]:
    started_at = time.perf_counter()
    response = llm.create_chat_completion(
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
        temperature=temperature,
        max_tokens=max_tokens,
    )
    message = response["choices"][0]["message"]["content"]
    payload = _extract_json_object(message)

    usage = response.get("usage") or {}
    generation_stats = {
        "prompt_tokens": int(usage.get("prompt_tokens", 0) or 0),
        "completion_tokens": int(usage.get("completion_tokens", 0) or 0),
        "total_tokens": int(usage.get("total_tokens", 0) or 0),
        "elapsed_ms": round((time.perf_counter() - started_at) * 1000.0, 2),
        "backend": "llama-cpp-python",
        "model_path": str(model.model_path),
        "n_ctx": DEFAULT_MODEL_CONTEXT,
    }
    meta = {
        "model_id": model.model_id,
        "model_path": str(model.model_path),
        "model_source": model.source,
        "backend": model.backend,
        "generation_stats": generation_stats,
    }
    return payload, meta


def validate_p1_payload(payload: dict[str, Any]) -> dict[str, Any]:
    triage = _require_text(payload, "triage")
    summary = _require_text(payload, "summary")
    qa = payload.get("qa")
    if not isinstance(qa, list) or not qa:
        raise RuntimeError("Model output must include a non-empty 'qa' list")
    normalized_qa: list[dict[str, str]] = []
    for idx, item in enumerate(qa, start=1):
        if not isinstance(item, dict):
            raise RuntimeError(f"qa[{idx}] must be an object")
        question = _require_text(item, "question")
        answer = _require_text(item, "answer")
        citation = _require_text(item, "citation")
        normalized_qa.append({"question": question, "answer": answer, "citation": citation})
    citations = payload.get("citations")
    if not isinstance(citations, list) or not citations:
        raise RuntimeError("Model output must include a non-empty 'citations' list")
    normalized_citations: list[dict[str, str]] = []
    for idx, item in enumerate(citations, start=1):
        if not isinstance(item, dict):
            raise RuntimeError(f"citations[{idx}] must be an object")
        question = _require_text(item, "question")
        snippet = _require_text(item, "snippet")
        normalized_citations.append({"question": question, "snippet": snippet})
    payload = dict(payload)
    payload["triage"] = triage
    payload["summary"] = summary
    payload["qa"] = normalized_qa
    payload["citations"] = normalized_citations
    return payload