nisse/nisse-core/src/nisse_core/probe.py

"""Check a model's declared context window against what its server does.

LiteLLM's ``max_input_tokens`` for a local model is only a declaration; the
real limit is whatever the model server allocates. This sends one prompt
half again as long as the declared window and reads back what the server did
with it:

- **It refused.** Over-long prompts fail loudly, which is the safe behaviour.
- **It cut the prompt.** The reported prompt size is where it cut, which is
  the server's real window. Ollama does this, silently.
- **It took the whole prompt.** The server allows more than is declared, so
  the declaration is conservative.

Every answer is checked for which model gave it. A proxy that falls back to
another model on an error returns a normal-looking answer under that model's
name, and once made this probe report a refusing server as roomy.

Run it with ``scripts/probe_context.py``. It spends one small and one large
request on the model, so it refuses windows above a ceiling unless told
otherwise: commercial models' limits are already known to LiteLLM.
"""

from dataclasses import dataclass

import httpx

from nisse_core.catalog import ModelFacts

# Of a probe's size, how much must come back for it to count as uncut. A
# chat template adds a little, and a tokenizer counts a little differently.
_UNCUT = 0.9
# How close the server's cut must be to the declaration to count as a match.
_MATCH = 0.02


@dataclass(frozen=True)
class ProbeResult:
    verdict: str  # "matches", "declared_too_high", "declared_conservative", ...
    message: str
    server_window: int | None = None


def interpret(
    declared: int, sent: int, reported: int | None, error: str | None
) -> ProbeResult:
    """What one over-long probe says about the server's window.

    ``sent`` is roughly how many tokens the probe held; ``reported`` is the
    prompt size the server reported, if it answered; ``error`` its error
    text, if it refused.
    """
    if error is not None:
        return ProbeResult(
            "refuses",
            f"the server refused a ~{sent}-token prompt ({error[:200]}). Over-long "
            "prompts fail loudly here, which is safe; the declared window of "
            f"{declared} was not contradicted.",
        )
    assert reported is not None
    if reported >= sent * _UNCUT:
        return ProbeResult(
            "declared_conservative",
            f"the server took the whole ~{sent}-token prompt ({reported} reported), "
            f"more than the declared {declared}. The declaration is safe; raise it "
            "if you want the room.",
            server_window=None,
        )
    if abs(reported - declared) <= declared * _MATCH:
        return ProbeResult(
            "matches",
            f"the server cut the prompt at {reported} tokens, matching the declared "
            f"{declared}.",
            server_window=reported,
        )
    if reported < declared:
        return ProbeResult(
            "declared_too_high",
            f"the server cut the prompt at {reported} tokens, but LiteLLM declares "
            f"{declared}. Prompts between the two would be silently cut: set "
            f"max_input_tokens to {reported}.",
            server_window=reported,
        )
    return ProbeResult(
        "declared_conservative",
        f"the server cut the prompt at {reported} tokens, above the declared "
        f"{declared}. The declaration is safe; it could be raised to {reported}.",
        server_window=reported,
    )


def _prompt(words: int) -> str:
    return " ".join(f"item{i}" for i in range(words)) + "\n\nReply with just: ok"


class Fallback(Exception):
    def __init__(self, answered_by: str) -> None:
        self.answered_by = answered_by


def _send(
    http: httpx.Client, model: str, words: int, backing: str
) -> tuple[int | None, str | None]:
    response = http.post(
        "/chat/completions",
        json={
            "model": model,
            "messages": [{"role": "user", "content": _prompt(words)}],
            "max_tokens": 5,
        },
    )
    if response.status_code != 200:
        return None, f"{response.status_code}: {response.text}"
    data = response.json()
    answered = str(data.get("model") or "")
    facts = ModelFacts(alias=model, backing=backing, provider="", supports_tools=None)
    if not facts.answered(answered):
        raise Fallback(answered)
    usage = data.get("usage") or {}
    return int(usage.get("prompt_tokens") or 0), None


def probe(
    http: httpx.Client, model: str, declared: int, backing: str = ""
) -> ProbeResult:
    """Measure how many tokens a word of the probe costs on this model, then
    send a prompt half again as long as the declared window."""
    calibration_words = 200
    try:
        tokens, error = _send(http, model, calibration_words, backing)
        if error is not None or not tokens:
            return ProbeResult("failed", f"a short calibration request failed: {error}")
        per_word = tokens / calibration_words
        sent = int(declared * 1.5)
        reported, error = _send(http, model, int(sent / per_word), backing)
    except Fallback as fallback:
        return ProbeResult(
            "fallback",
            f"{fallback.answered_by} answered instead of {model}: the proxy fell "
            "back to another model, so nothing here describes this one. Remove "
            "the fallback, or probe the server directly.",
        )
    return interpret(declared, sent, reported, error)