unsloth/studio/backend/picker/service.py

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

from __future__ import annotations

import json
import logging
import os
import re
from pathlib import Path
from typing import Optional

from hub.services.models.folder_browser import (
    _build_browse_allowlist,
    _is_path_inside_allowlist,
)
from hub.utils.gguf import extract_quant_label, iter_hf_cache_snapshots
from utils.models.gguf_metadata import read_gguf_chat_template
from utils.models.model_config import (
    _extract_quant_label,
    _is_big_endian_gguf_path,
    _is_mmproj,
    _is_mtp_drafter,
)
from utils.hf_cache_settings import active_hf_hub_cache
from utils.paths.path_utils import (
    is_local_path,
    normalize_path,
    resolve_cached_repo_id_case,
)

from .schemas import MAX_CHAT_TEMPLATE_BYTES, ValidateChatTemplateResponse

logger = logging.getLogger(__name__)

_VALID_REPO_ID = re.compile(r"^[A-Za-z0-9._-]+/[A-Za-z0-9._-]+$")


def _is_valid_repo_id(repo_id: str) -> bool:
    return bool(_VALID_REPO_ID.fullmatch(repo_id))


_TOKENIZER_CONFIG_PATHS = ("tokenizer_config.json", "LLM/tokenizer_config.json")
_JINJA_TEMPLATE_PATHS = ("chat_template.jinja", "LLM/chat_template.jinja")
_PROCESSOR_TEMPLATE_PATHS = ("chat_template.json", "LLM/chat_template.json")

# Cap sidecar reads so a malformed or hostile metadata file cannot exhaust memory
# before its template is size-checked. The JSON envelope may exceed a bare template
# (it carries other tokenizer metadata); the extracted template is still bounded by
# MAX_CHAT_TEMPLATE_BYTES downstream.
MAX_TEMPLATE_METADATA_BYTES = 4 * 1024 * 1024


def _read_bounded_text(path: Path, limit: int) -> Optional[str]:
    """Read at most `limit` bytes of UTF-8 text; None if larger or unreadable."""
    try:
        with path.open("rb") as f:
            data = f.read(limit + 1)
    except OSError:
        return None
    if len(data) > limit:
        return None
    try:
        return data.decode("utf-8")
    except UnicodeError:
        return None


def _leaf_inside_allowlist(path: Path, allow_roots: Optional[list[Path]]) -> bool:
    # Block symlinked children from escaping the validated directory (realpath-checked).
    # None = trusted caller (HF cache / remote download).
    return allow_roots is None or _is_path_inside_allowlist(path, allow_roots)


def validate_chat_template(template: str) -> ValidateChatTemplateResponse:
    text = (template or "").strip()
    if not text:
        return ValidateChatTemplateResponse(valid = True, error = None)
    # Import Jinja lazily: optional at runtime (e.g. GGUF-only installs), so a
    # missing dependency must not crash API startup.
    try:
        from jinja2 import TemplateError
        from jinja2.ext import Extension
        from jinja2.sandbox import ImmutableSandboxedEnvironment
    except ImportError:
        return ValidateChatTemplateResponse(valid = True, error = None)

    class _GenerationTag(Extension):
        # Accept Transformers' {% generation %} assistant-mask tag so a pasted HF
        # chat template validates (we only parse it).
        tags = {"generation"}

        def parse(self, parser):
            next(parser.stream)
            return parser.parse_statements(["name:endgeneration"], drop_needle = True)

    try:
        env = ImmutableSandboxedEnvironment(
            trim_blocks = True,
            lstrip_blocks = True,
            extensions = ["jinja2.ext.loopcontrols", _GenerationTag],
        )
        env.parse(text)
        return ValidateChatTemplateResponse(valid = True, error = None)
    except TemplateError as exc:
        message = getattr(exc, "message", None) or str(exc)
        lineno = getattr(exc, "lineno", None)
        if lineno:
            message = f"Line {lineno}: {message}"
        return ValidateChatTemplateResponse(valid = False, error = message)
    except Exception as exc:
        return ValidateChatTemplateResponse(valid = False, error = str(exc))


def _chat_template_from_tokenizer_config(config: dict) -> Optional[str]:
    if not isinstance(config, dict):
        return None
    raw = config.get("chat_template")
    if isinstance(raw, str) and raw.strip():
        return raw
    if isinstance(raw, list):
        fallback: Optional[str] = None
        for entry in raw:
            if not isinstance(entry, dict):
                continue
            template = entry.get("template")
            if not isinstance(template, str):
                continue
            if entry.get("name") == "default":
                return template
            if fallback is None:
                fallback = template
        return fallback
    return None


def _chat_template_from_jinja_file(
    dir_path: Path, allow_roots: Optional[list[Path]] = None
) -> Optional[str]:
    for rel in _JINJA_TEMPLATE_PATHS:
        template_file = dir_path / rel
        if not template_file.exists() or not _leaf_inside_allowlist(template_file, allow_roots):
            continue
        try:
            if template_file.stat().st_size > MAX_CHAT_TEMPLATE_BYTES:
                continue
            template = template_file.read_text(encoding = "utf-8")
        except Exception:
            continue
        if template.strip():
            return template
    return None


def _chat_template_from_processor_payload(payload: object) -> Optional[str]:
    # processor chat_template.json may be the template string itself or a
    # {name: template} map, not only a tokenizer_config-shaped object.
    if isinstance(payload, str):
        return payload if payload.strip() else None
    template = _chat_template_from_tokenizer_config(payload)  # type: ignore[arg-type]
    if template:
        return template
    if isinstance(payload, dict):
        # Named-template map: prefer "default", else the first non-empty entry
        # (mirrors the tokenizer-config list fallback).
        default = payload.get("default")
        if isinstance(default, str) and default.strip():
            return default
        for value in payload.values():
            if isinstance(value, str) and value.strip():
                return value
    return None


def _chat_template_from_processor_json(
    dir_path: Path, allow_roots: Optional[list[Path]] = None
) -> Optional[str]:
    for rel in _PROCESSOR_TEMPLATE_PATHS:
        config_file = dir_path / rel
        if not config_file.exists() or not _leaf_inside_allowlist(config_file, allow_roots):
            continue
        raw = _read_bounded_text(config_file, MAX_TEMPLATE_METADATA_BYTES)
        if raw is None:
            continue
        try:
            payload = json.loads(raw)
        except Exception:
            continue
        template = _chat_template_from_processor_payload(payload)
        if template:
            return template
    return None


def _chat_template_from_tokenizer_dir(
    dir_path: Path, allow_roots: Optional[list[Path]] = None
) -> Optional[str]:
    jinja = _chat_template_from_jinja_file(dir_path, allow_roots)
    if jinja:
        return jinja
    for rel in _TOKENIZER_CONFIG_PATHS:
        config_file = dir_path / rel
        if not config_file.exists() or not _leaf_inside_allowlist(config_file, allow_roots):
            continue
        raw = _read_bounded_text(config_file, MAX_TEMPLATE_METADATA_BYTES)
        if raw is None:
            continue
        try:
            config = json.loads(raw)
        except Exception:
            continue
        template = _chat_template_from_tokenizer_config(config)
        if template:
            return template
    return _chat_template_from_processor_json(dir_path, allow_roots)


_GGUF_SCAN_MAX_DEPTH = 2


def _iter_ggufs(dir_path: Path) -> list[Path]:
    if dir_path == dir_path.parent:
        return []
    root = str(dir_path)
    found: list[Path] = []
    for current, dirs, files in os.walk(root, followlinks = False):
        rel = os.path.relpath(current, root)
        depth = 0 if rel == os.curdir else rel.count(os.sep) + 1
        if depth >= _GGUF_SCAN_MAX_DEPTH:
            dirs[:] = []
        for name in files:
            if not name.lower().endswith(".gguf") or _is_mmproj(name):
                continue
            path = Path(current) / name
            try:
                rel = path.relative_to(dir_path).as_posix()
            except ValueError:
                rel = name
            quant = _extract_quant_label(rel)
            if _is_mtp_drafter(rel) or _is_big_endian_gguf_path(rel, quant):
                continue
            found.append(path)
    return found


def _variant_matches(relative_path: str, needle: str) -> bool:
    quant = _extract_quant_label(relative_path).lower()
    if quant == needle:
        return True
    if extract_quant_label(relative_path).lower() == needle:
        return True
    prefix = f"{needle}-"
    if not quant.startswith(prefix):
        return False
    suffix = quant[len(prefix) :]
    if not suffix.endswith("bpw"):
        return False
    value = suffix[:-3]
    return bool(value) and value.replace(".", "", 1).isdigit()


_GGUF_SPLIT_INDEX_RE = re.compile(r"-(\d{3,})-of-\d{3,}$", re.IGNORECASE)


def _is_nonfirst_gguf_split(path: Path) -> bool:
    match = _GGUF_SPLIT_INDEX_RE.search(path.stem)
    return match is not None and int(match.group(1)) != 1


def _find_gguf_in_dir(dir_path: Path, gguf_variant: Optional[str]) -> Optional[Path]:
    try:
        ggufs = sorted(_iter_ggufs(dir_path))
    except OSError:
        return None
    if not ggufs:
        return None
    needle = (gguf_variant or "").strip().lower()
    if needle:
        for path in ggufs:
            try:
                relative = path.relative_to(dir_path).as_posix()
            except ValueError:
                relative = path.name
            if _variant_matches(relative, needle):
                return path
        return None
    candidates = [path for path in ggufs if not _is_nonfirst_gguf_split(path)] or ggufs
    try:
        return max(candidates, key = lambda path: path.stat().st_size)
    except OSError:
        return candidates[0]


def _chat_template_from_dir(
    dir_path: Path,
    gguf_variant: Optional[str] = None,
    allow_roots: Optional[list[Path]] = None,
) -> Optional[str]:
    def from_gguf() -> Optional[str]:
        gguf = _find_gguf_in_dir(dir_path, gguf_variant)
        if gguf is None or not _leaf_inside_allowlist(gguf, allow_roots):
            return None
        return read_gguf_chat_template(str(gguf))

    # Sidecar tokenizer files (chat_template.jinja / tokenizer_config.json) are the
    # author's maintained template and supersede the GGUF's possibly-stale embedded
    # copy. The variant only picks the GGUF fallback, so tokenizer-first precedence
    # holds whether or not a variant is given.
    return _chat_template_from_tokenizer_dir(dir_path, allow_roots) or from_gguf()


def read_default_chat_template(
    model_name: str,
    hf_token: Optional[str] = None,
    gguf_variant: Optional[str] = None,
) -> Optional[str]:
    if not isinstance(model_name, str) or not model_name.strip():
        return None
    name = model_name.strip()

    if is_local_path(name):
        try:
            target = Path(normalize_path(name)).expanduser()
            allow_roots = _build_browse_allowlist()
            if not _is_path_inside_allowlist(target, allow_roots):
                logger.debug("Refused chat template read outside allowed folders: %s", name)
                return None
            if name.lower().endswith(".gguf"):
                # Prefer a maintained sidecar next to the file over the GGUF's
                # embedded copy (tokenizer-first precedence, as elsewhere).
                sidecar = _chat_template_from_tokenizer_dir(target.parent, allow_roots)
                if sidecar:
                    return sidecar
                return read_gguf_chat_template(str(target))
            return _chat_template_from_dir(target, gguf_variant, allow_roots)
        except Exception as exc:
            logger.debug("Could not read local chat template for %s: %s", name, exc)
            return None

    if not _is_valid_repo_id(name):
        return None

    resolved = resolve_cached_repo_id_case(name)

    try:
        # Resolve within each cached revision, newest first. A revision's sidecar
        # supersedes its own embedded GGUF copy, but must not override a newer
        # revision, so precedence stays per-snapshot rather than global.
        for snapshot in iter_hf_cache_snapshots(resolved):
            template = _chat_template_from_dir(snapshot, gguf_variant)
            if template:
                return template
    except Exception as exc:
        logger.debug("Could not read cached chat template for %s: %s", resolved, exc)

    try:
        from huggingface_hub import HfApi, hf_hub_download

        _api = HfApi()

        def _remote_exceeds_cap(rel: str) -> bool:
            # Best-effort: skip the download when the remote's advertised size
            # exceeds the cap, so a maliciously large sidecar is never fetched.
            try:
                infos = _api.get_paths_info(resolved, [rel], repo_type = "model", token = hf_token)
            except Exception:
                return False
            for info in infos:
                size = getattr(info, "size", None)
                if (
                    getattr(info, "path", None) == rel
                    and isinstance(size, int)
                    and size > MAX_TEMPLATE_METADATA_BYTES
                ):
                    return True
            return False

        def _download_text(rel: str) -> Optional[str]:
            if _remote_exceeds_cap(rel):
                return None
            try:
                path = hf_hub_download(
                    resolved,
                    rel,
                    token = hf_token,
                    cache_dir = active_hf_hub_cache(),
                )
                return _read_bounded_text(Path(path), MAX_TEMPLATE_METADATA_BYTES)
            except Exception:
                return None

        for rel in _JINJA_TEMPLATE_PATHS:
            template = _download_text(rel)
            if not template or not template.strip():
                continue
            # A raw Jinja sidecar is the whole template, so it must fit the route's
            # response cap (the local path skips oversized .jinja too). Download stays
            # bounded at MAX_TEMPLATE_METADATA_BYTES so a large JSON embedding a small
            # template still extracts below, but an over-cap Jinja is dropped so the
            # search falls through to the tokenizer/processor template.
            if len(template.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES:
                continue
            return template

        for rel in _TOKENIZER_CONFIG_PATHS:
            raw = _download_text(rel)
            if not raw:
                continue
            try:
                config = json.loads(raw)
            except Exception:
                continue
            template = _chat_template_from_tokenizer_config(config)
            if template:
                return template

        for rel in _PROCESSOR_TEMPLATE_PATHS:
            raw = _download_text(rel)
            if not raw:
                continue
            try:
                payload = json.loads(raw)
            except Exception:
                continue
            template = _chat_template_from_processor_payload(payload)
            if template:
                return template

        return None
    except Exception as exc:
        logger.debug("Could not fetch chat template for %s: %s", resolved, exc)
        return None