unsloth/unsloth_cli/_inference.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

"""Model loading and streaming shared by `inference` and `chat`."""

import asyncio
import os
import re
import sys
from pathlib import Path
from typing import List, Optional

import typer

_THINK_OPEN = "<think>"
_THINK_BLOCK = re.compile(rf"{re.escape(_THINK_OPEN)}.*?</think>", re.DOTALL)

# Cloudflare (in front of remote Studio proxies like RunPod) 403s the default
# "Python-urllib/X.Y" User-Agent as a bot; send a real one on every request.
_USER_AGENT = "unsloth-cli"

# Built lazily; urllib stays function-local to match this module.
_no_redirect_opener = None


def urlopen_no_redirect(request, timeout):
    """urlopen that errors on any redirect: following a 3xx would send a bearer
    token (or accept an identity proof) to a base we never vetted, letting a port
    squatter relay a real Studio's response."""
    global _no_redirect_opener
    if _no_redirect_opener is None:
        import urllib.error
        import urllib.request

        class _NoRedirect(urllib.request.HTTPRedirectHandler):
            def redirect_request(self, req, fp, code, msg, headers, newurl):
                raise urllib.error.HTTPError(
                    req.full_url, code, f"refusing redirect to {newurl}", headers, fp
                )

        _no_redirect_opener = urllib.request.build_opener(_NoRedirect)
    return _no_redirect_opener.open(request, timeout = timeout)


def ensure_studio_backend_path() -> None:
    backend_dir = str(Path(__file__).resolve().parents[1] / "studio" / "backend")
    if backend_dir not in sys.path:
        sys.path.insert(0, backend_dir)


def configure_quiet_logging() -> None:
    import logging

    import structlog

    # The CLI never configures structlog, so without this every backend INFO
    # line prints. LOG_LEVEL is exported so the worker subprocess inherits it.
    level_name = os.environ.setdefault("LOG_LEVEL", "WARNING").upper()
    level = getattr(logging, level_name, logging.WARNING)
    structlog.configure(wrapper_class = structlog.make_filtering_bound_logger(level))
    os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")


def visible_text(text: str, show_thinking: bool) -> str:
    if show_thinking:
        return text
    text = _THINK_BLOCK.sub("", text)
    # Hold back an unclosed trailing <think> so reasoning never leaks mid-stream.
    open_idx = text.find(_THINK_OPEN)
    if open_idx != -1:
        text = text[:open_idx]
    max_prefix = min(len(text), len(_THINK_OPEN) - 1)
    for size in range(max_prefix, 0, -1):
        if _THINK_OPEN.startswith(text[-size:]):
            return text[:-size]
    return text


def stream_to_stdout(stream, show_thinking: bool) -> str:
    # Backends yield the full text-so-far on each step (llama.cpp ends with a
    # metadata dict, skipped); print the growing tail, return the raw text.
    raw = ""
    shown = ""
    for chunk in stream:
        if not isinstance(chunk, str):
            continue
        raw = chunk
        rendered = visible_text(chunk, show_thinking)
        delta = rendered[len(shown) :]
        if delta:
            sys.stdout.write(delta)
            sys.stdout.flush()
        shown = rendered
    sys.stdout.write("\n")
    sys.stdout.flush()
    return raw


def stream_markdown(stream, show_thinking: bool, *, console) -> str:
    from rich.live import Live
    from rich.markdown import Markdown
    from rich.text import Text

    raw = ""
    with Live(console = console, refresh_per_second = 12, vertical_overflow = "visible") as live:
        for chunk in stream:
            if not isinstance(chunk, str):
                continue
            raw = chunk
            visible = visible_text(chunk, show_thinking)
            live.update(Markdown(visible) if visible.strip() else Text(""))
    return raw


def collect_stream(stream, show_thinking: bool) -> str:
    raw = ""
    for chunk in stream:
        if isinstance(chunk, str):
            raw = chunk
    return visible_text(raw, show_thinking)


def render_columns(
    left_label: str,
    left_text: str,
    right_label: str,
    right_text: str,
    *,
    console = None,
) -> None:
    from rich import box
    from rich.console import Console
    from rich.table import Table

    table = Table(box = box.MINIMAL, expand = True, padding = (0, 1), pad_edge = False)
    table.add_column(left_label, header_style = "bold yellow", ratio = 1, overflow = "fold")
    table.add_column(right_label, header_style = "bold magenta", ratio = 1, overflow = "fold")
    table.add_row(left_text or "", right_text or "")
    (console or Console()).print(table)


class ChatBackend:
    """Uniform stream()/close() over the llama-server and Unsloth backends."""

    def __init__(self, kind: str, backend) -> None:
        self._kind = kind  # "gguf" | "unsloth"
        self._backend = backend

    def stream(
        self,
        messages: list,
        *,
        system_prompt: str,
        temperature: float,
        top_p: float,
        top_k: int,
        max_new_tokens: int,
        repetition_penalty: float,
        enable_thinking: bool,
        use_adapter: Optional[bool] = None,
    ):
        if self._kind == "gguf":
            # llama-server takes the system prompt as the first message.
            msgs = list(messages)
            if system_prompt:
                msgs = [{"role": "system", "content": system_prompt}, *msgs]
            return self._backend.generate_chat_completion(
                messages = msgs,
                temperature = temperature,
                top_p = top_p,
                top_k = top_k,
                max_tokens = max_new_tokens,
                repetition_penalty = repetition_penalty,
                enable_thinking = enable_thinking,
            )
        gen_kwargs = dict(
            messages = messages,
            system_prompt = system_prompt,
            temperature = temperature,
            top_p = top_p,
            top_k = top_k,
            max_new_tokens = max_new_tokens,
            repetition_penalty = repetition_penalty,
            enable_thinking = enable_thinking,
        )
        if use_adapter is not None:
            return self._backend.generate_with_adapter_control(
                use_adapter = use_adapter, **gen_kwargs
            )
        return self._backend.generate_chat_response(**gen_kwargs)

    def close(self) -> None:
        # Shut the worker down directly: the graceful unload_model waits for
        # an ack that compare mode can swallow, hanging exit for minutes.
        try:
            if self._kind == "gguf":
                self._backend.unload_model()
            else:
                self._backend._shutdown_subprocess(timeout = 2.0)
        except Exception:
            pass


def resolve_model_config(model: str, *, hf_token: Optional[str]):
    ensure_studio_backend_path()
    from utils.models import ModelConfig

    model_config = ModelConfig.from_identifier(model_id = model, hf_token = hf_token)
    if not model_config:
        typer.echo("Could not resolve model config", err = True)
        raise typer.Exit(code = 1)
    return model_config


def _validate_llama_extra_args_or_exit(llama_extra_args: Optional[List[str]]) -> list[str]:
    from core.inference.llama_server_args import validate_extra_args
    try:
        return validate_extra_args(llama_extra_args)
    except ValueError as exc:
        typer.echo(f"Error: {exc}", err = True)
        raise typer.Exit(code = 1)


def _load_gguf_backend(
    model_config,
    *,
    hf_token,
    max_seq_length,
    tensor_parallel: bool = False,
    llama_extra_args: Optional[List[str]] = None,
):
    ensure_studio_backend_path()
    from core.inference.llama_cpp import LlamaCppBackend
    from core.inference.tensor_fallback import load_with_tensor_fallback

    llama_backend = LlamaCppBackend()
    extra_args = _validate_llama_extra_args_or_exit(llama_extra_args)
    common = dict(
        hf_variant = model_config.gguf_variant,
        model_identifier = model_config.identifier,
        is_vision = model_config.is_vision,
        n_ctx = max_seq_length,
    )

    async def _attempt_gguf_load(
        requested_tensor_parallel: bool, attempt_extra_args: Optional[List[str]]
    ) -> bool:
        attempt_common = dict(
            common,
            tensor_parallel = requested_tensor_parallel,
            extra_args = attempt_extra_args,
        )
        if model_config.gguf_hf_repo:
            return llama_backend.load_model(
                hf_repo = model_config.gguf_hf_repo,
                hf_token = hf_token,
                **attempt_common,
            )
        return llama_backend.load_model(
            gguf_path = model_config.gguf_file,
            mmproj_path = model_config.gguf_mmproj_file,
            mtp_draft_path = model_config.gguf_mtp_file,
            **attempt_common,
        )

    loaded = asyncio.run(
        load_with_tensor_fallback(
            _attempt_gguf_load,
            requested_tensor = tensor_parallel,
            extra_args = extra_args,
            label = model_config.identifier,
        )
    )
    if not loaded:
        typer.echo("Model load failed", err = True)
        raise typer.Exit(code = 1)
    return ChatBackend("gguf", llama_backend)


def load_chat_backend(
    model: str,
    *,
    hf_token: Optional[str],
    max_seq_length: int,
    load_in_4bit: bool,
    tensor_parallel: bool = False,
    llama_extra_args: Optional[List[str]] = None,
    model_config = None,
    fresh_backend: bool = False,
):
    """Load `model` in-process: GGUF via llama-server, else the orchestrator.

    fresh_backend uses a private orchestrator so a second model (compare's
    base column) can run alongside the main one.
    """
    if model_config is None:
        model_config = resolve_model_config(model, hf_token = hf_token)

    typer.echo(f"Loading {model}", err = True)

    if model_config.is_gguf:
        return _load_gguf_backend(
            model_config,
            hf_token = hf_token,
            max_seq_length = max_seq_length,
            tensor_parallel = tensor_parallel,
            llama_extra_args = llama_extra_args,
        )

    if fresh_backend:
        ensure_studio_backend_path()
        from core.inference import InferenceOrchestrator
        backend = InferenceOrchestrator()
    else:
        ensure_studio_backend_path()
        from core.inference import get_inference_backend
        backend = get_inference_backend()
    if not backend.load_model(
        config = model_config,
        max_seq_length = max_seq_length,
        load_in_4bit = load_in_4bit,
        hf_token = hf_token,
    ):
        typer.echo("Model load failed", err = True)
        raise typer.Exit(code = 1)
    return ChatBackend("unsloth", backend)


def _loopback_candidate_bases(base: str) -> list:
    """For a bare ``localhost`` base, the concrete IP bases to try, IPv4
    127.0.0.1 first (where ``unsloth studio`` binds by default). Pinning to one
    address up front means discovery, the identity check, and the credential we
    then send all target the same endpoint instead of racing IPv4/IPv6
    resolution -- which would otherwise let the health probe land on one address
    and the identity check on another. A literal IP or remote name is unchanged.
    """
    from urllib.parse import urlparse

    parsed = urlparse(base)
    if (parsed.hostname or "").lower() != "localhost":
        return [base]
    import socket

    port = parsed.port or (443 if parsed.scheme == "https" else 80)
    try:
        ips = {
            ai[4][0] for ai in socket.getaddrinfo(parsed.hostname, port, type = socket.SOCK_STREAM)
        }
    except Exception:
        return [base]
    ordered = sorted(ips, key = lambda ip: (ip != "127.0.0.1", ip))
    bases = [
        f"{parsed.scheme}://" + (f"[{ip}]:{port}" if ":" in ip else f"{ip}:{port}")
        for ip in ordered
    ]
    return bases or [base]


def find_studio_server(timeout: float = 3.0) -> Optional[str]:
    import urllib.request

    base = os.environ.get("UNSLOTH_STUDIO_URL", "http://127.0.0.1:8888").rstrip("/")
    # Try the concrete loopback addresses in order and return the first that
    # answers, so the rest of the flow talks to that exact address.
    for candidate in _loopback_candidate_bases(base):
        request = urllib.request.Request(
            f"{candidate}/api/health", headers = {"User-Agent": _USER_AGENT}
        )
        try:
            with urllib.request.urlopen(request, timeout = timeout):
                return candidate
        except Exception:
            continue
    return None


def is_loopback_url(base: str) -> bool:
    """True only when *base* resolves to loopback. find_studio_server() trusts a
    base after only a health probe, so credentials are auto-sent only to loopback
    (a local Studio or an SSH tunnel on 127.0.0.1), the targets the auto flows mean."""
    from urllib.parse import urlparse

    host = (urlparse(base).hostname or "").lower()
    if host in ("localhost", "127.0.0.1", "::1"):
        return True
    try:
        import ipaddress
        return ipaddress.ip_address(host).is_loopback
    except ValueError:
        return False


def verify_studio_identity(base: str, timeout: float = 3.0) -> bool:
    """Confirm `base` is really this machine's Studio before sending a secret.

    Send a random nonce to /api/auth/identity and check the returned HMAC against
    the one computed from the local same-user secret; an endpoint without that
    secret (port squatter, remote/fake) can't match. Fails closed on any error."""
    import base64
    import hmac as _hmac
    import json
    import secrets as _secrets
    import socket
    import urllib.request
    from urllib.parse import urlparse

    try:
        import studio.backend.core  # noqa: F401  puts studio/backend on sys.path
        from studio.backend.auth import storage
    except Exception:
        return False

    parsed = urlparse(base)
    host = parsed.hostname or ""
    port = parsed.port or (443 if parsed.scheme == "https" else 80)
    # Resolve to one concrete address and talk to *that* address, then bind the
    # proof to (address, port). A name like localhost can resolve to a squatter on
    # ::1 while the real Studio is on 127.0.0.1; connecting to the resolved IP and
    # binding to it means a proof relayed from a different address/port won't match.
    try:
        ip = socket.getaddrinfo(host, port, type = socket.SOCK_STREAM)[0][4][0]
    except Exception:
        return False
    netloc = f"[{ip}]:{port}" if ":" in ip else f"{ip}:{port}"
    nonce = _secrets.token_bytes(32)
    query = base64.urlsafe_b64encode(nonce).decode()
    request = urllib.request.Request(
        f"{parsed.scheme}://{netloc}/api/auth/identity?nonce={query}",
        headers = {"User-Agent": _USER_AGENT, "Host": parsed.netloc},
    )
    try:
        # No redirects: a 302 could relay a real Studio's proof (see urlopen_no_redirect).
        # Cap the read: the server is still unverified, so don't trust its length.
        with urlopen_no_redirect(request, timeout = timeout) as response:
            proof = json.loads(response.read(65536).decode() or "{}").get("proof")
    except Exception:
        return False
    if not isinstance(proof, str):
        return False
    try:
        expected = storage.compute_identity_proof(nonce, ip, port)
    except Exception:
        return False
    return _hmac.compare_digest(proof, expected)


def _studio_token() -> Optional[str]:
    """Self-issue a JWT: the CLI runs as the same OS user as the server, so it
    signs with the same stored secret the server validates against."""
    try:
        import studio.backend.core  # noqa: F401  puts studio/backend on sys.path

        from studio.backend.auth import storage
        from studio.backend.auth.authentication import create_access_token

        row = storage.get_connection().execute("SELECT username FROM auth_user LIMIT 1").fetchone()
        return create_access_token(row[0], desktop = True) if row else None
    except Exception:
        return None


class HttpChatBackend:
    """Chat against a running Studio server over its OpenAI-compatible API.

    close() leaves the model loaded on purpose — the next session (or the
    UI) starts instantly.
    """

    def __init__(self, base_url: str, token: str) -> None:
        self._base = base_url
        self._token = token

    def _request(
        self,
        method: str,
        path: str,
        payload = None,
        timeout = None,
    ):
        import json
        import urllib.request

        request = urllib.request.Request(
            self._base + path,
            data = None if payload is None else json.dumps(payload).encode(),
            headers = {
                "Authorization": f"Bearer {self._token}",
                "Content-Type": "application/json",
                "User-Agent": _USER_AGENT,
            },
            method = method,
        )
        # No redirects: this carries a bearer token (see urlopen_no_redirect).
        return urlopen_no_redirect(request, timeout = timeout)

    def ensure_loaded(
        self,
        model: str,
        *,
        hf_token,
        max_seq_length,
        load_in_4bit,
        tensor_parallel: bool = False,
        llama_extra_args: Optional[List[str]] = None,
    ) -> None:
        typer.echo(f"Loading {model} on the Studio server", err = True)
        payload = {
            "model_path": model,
            "hf_token": hf_token,
            "max_seq_length": max_seq_length,
            "load_in_4bit": load_in_4bit,
            "tensor_parallel": tensor_parallel,
        }
        if llama_extra_args:
            payload["llama_extra_args"] = llama_extra_args
        try:
            self._request(
                "POST",
                "/api/inference/load",
                payload,
            ).close()
        except Exception as exc:
            typer.echo(f"Model load failed: {exc}", err = True)
            raise typer.Exit(code = 1)

    def stream(
        self,
        messages: list,
        *,
        system_prompt: str,
        temperature: float,
        top_p: float,
        top_k: int,
        max_new_tokens: int,
        repetition_penalty: float,
        enable_thinking: bool,
        use_adapter: Optional[bool] = None,
    ):
        import json

        msgs = list(messages)
        if system_prompt:
            msgs = [{"role": "system", "content": system_prompt}, *msgs]
        resp = self._request(
            "POST",
            "/v1/chat/completions",
            {
                "model": "default",
                "messages": msgs,
                "stream": True,
                "temperature": temperature,
                "top_p": top_p,
                "top_k": top_k,
                "max_tokens": max_new_tokens,
                "repetition_penalty": repetition_penalty,
                "enable_thinking": enable_thinking,
            },
        )

        def cumulative():
            # Accumulate SSE deltas into the full-text-so-far convention the
            # stream helpers expect.
            text = ""
            with resp:
                for raw_line in resp:
                    line = raw_line.decode("utf-8", "replace").strip()
                    if not line.startswith("data:"):
                        continue
                    data = line[len("data:") :].strip()
                    if data == "[DONE]":
                        break
                    try:
                        parsed = json.loads(data)
                    except ValueError:
                        continue
                    if "error" in parsed:
                        raise RuntimeError(
                            f"Server error: {parsed['error'].get('message', 'Unknown server error')}"
                        )
                    try:
                        delta = parsed["choices"][0]["delta"].get("content")
                    except (KeyError, IndexError):
                        continue
                    if not delta:
                        continue
                    text += delta
                    # An emoji can arrive split across two deltas as lone
                    # surrogate halves: hold back a trailing half, merge pairs.
                    visible = text
                    if "\ud800" <= visible[-1] <= "\udbff":
                        visible = visible[:-1]
                    yield visible.encode("utf-16", "surrogatepass").decode("utf-16", "replace")

        return cumulative()

    def close(self) -> None:
        pass


def connect_studio_server(
    model: str,
    *,
    hf_token,
    max_seq_length,
    load_in_4bit,
    tensor_parallel: bool = False,
    llama_extra_args: Optional[List[str]] = None,
):
    """Backend on a running Studio server, or None (caller loads locally)."""
    base_url = find_studio_server()
    if not base_url:
        return None

    # Explicit server (UNSLOTH_STUDIO_URL) we can't safely attach to -> fail loudly;
    # opportunistic local discovery just falls back to a local load.
    explicit = bool(os.environ.get("UNSLOTH_STUDIO_URL"))

    def _refuse(reason: str):
        if not explicit:
            return None
        typer.echo(
            f"Can't attach to the Studio server at {base_url}: {reason} Run Studio "
            "on this machine, or unset UNSLOTH_STUDIO_URL to load the model locally.",
            err = True,
        )
        raise typer.Exit(code = 1)

    # Only hand the self-issued JWT (signed with the local secret) to loopback: a
    # remote URL is unverified and a real remote Studio would reject it anyway.
    if not is_loopback_url(base_url):
        return _refuse(
            "it isn't a local Studio, so a self-issued token can't "
            "authenticate to it and must not be sent to it."
        )
    # Confirm the loopback responder is really our Studio (not a port squatter).
    if not verify_studio_identity(base_url):
        return _refuse(
            "its identity couldn't be verified (it may be running as a "
            "different OS user, or another process took the port)."
        )
    token = _studio_token()
    if not token:
        return _refuse("couldn't self-issue a Studio token (is Studio set up here?).")
    backend = HttpChatBackend(base_url, token)
    backend.ensure_loaded(
        model,
        hf_token = hf_token,
        max_seq_length = max_seq_length,
        load_in_4bit = load_in_4bit,
        tensor_parallel = tensor_parallel,
        llama_extra_args = llama_extra_args,
    )
    return backend