vllm_mlx.constrained.json_schema_processor

JSONSchemaLogitsProcessor — a mlx_lm-compatible logits processor that masks the vocabulary so the model can only emit tokens forming a valid JSON value (optionally matching a JSON schema).

View the complete module source at #L1-L924.

API details

Each callable below includes its exact signature, type annotations, inputs, defaults, return contract, documented exceptions, implementation source, and parsed docstring sections when the source provides them.

vllm_mlx.constrained.json_schema_processor

JSONSchemaLogitsProcessor — a mlx_lm-compatible logits processor that masks the vocabulary so the model can only emit tokens forming a valid JSON value (optionally matching a JSON schema).

The processor implements the signature expected by mlx_lm.generate.generate_step and vllm_mlx's batched engine alike:

processor(tokens: mx.array, logits: mx.array) -> mx.array

tokens contains the full sequence generated for this request so far (prompt + previously emitted tokens), and logits is the last-step logits row.

vllm_mlx.constrained.json_schema_processor.logger module-attribute

logger = logging.getLogger(__name__)

vllm_mlx.constrained.json_schema_processor._parser_cache module-attribute

_parser_cache: dict[str, tuple[dict, Any]] = {}

vllm_mlx.constrained.json_schema_processor._MAX_NONPROGRESS_WHITESPACE_CHARS module-attribute

_MAX_NONPROGRESS_WHITESPACE_CHARS = 256

vllm_mlx.constrained.json_schema_processor._JSON_WHITESPACE module-attribute

_JSON_WHITESPACE = frozenset(' \t\r\n')

vllm_mlx.constrained.json_schema_processor._GENERIC_JSON_SCHEMA module-attribute

_GENERIC_JSON_SCHEMA: dict = {'anyOf': [{'type': 'object'}, {'type': 'array'}]}

vllm_mlx.constrained.json_schema_processor.LMFormatEnforcerNotAvailableError

Bases: RuntimeError

Raised when lm-format-enforcer is required but not installed.

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor

JSONSchemaLogitsProcessor(schema: dict | None, tokenizer: Any)

Logits processor that constrains generation to valid JSON.

Parameters

schema: The JSON Schema the output must match. When None, any valid JSON object/array is accepted (json_object mode). tokenizer: The tokenizer used for generation. Its vocabulary is iterated once (via :mod:vllm_mlx.constrained.cache) and cached for subsequent requests.

Source code in vllm_mlx/constrained/json_schema_processor.py

def __init__(
    self,
    schema: dict | None,
    tokenizer: Any,
) -> None:
    if not is_available():
        raise LMFormatEnforcerNotAvailableError(
            "lm-format-enforcer is not installed. "
            'Install it with `pip install "lm-format-enforcer>=0.10.9"`.'
        )

    from lmformatenforcer import TokenEnforcer

    self._tokenizer = tokenizer
    self._schema = schema
    self._tok_data = get_tokenizer_data(tokenizer)
    if self._tok_data is None:
        raise LMFormatEnforcerNotAvailableError(
            "Could not build TokenEnforcerTokenizerData for this tokenizer."
        )

    # Reuse a memoised JsonSchemaParser for this schema; the parser is
    # immutable after construction so it can be shared across requests.
    # Each request still gets its own TokenEnforcer (which carries the
    # per-sequence ``prefix_states`` and must not be shared).
    self._disabled = False
    try:
        parser_schema, self._parser = _get_or_build_parser(schema)
        self._enforcer = TokenEnforcer(self._tok_data, self._parser)
    except Exception as exc:
        logger.warning(
            "JSONSchemaLogitsProcessor: enforcer init failed (%s); "
            "falling back to unconstrained generation",
            exc,
        )
        self._disabled = True
        self._parser = None  # type: ignore[assignment]
        self._enforcer = None  # type: ignore[assignment]

    # Bootstrap the enforcer's ``prefix_states`` with the empty tuple so
    # that subsequent ``get_allowed_tokens([t1, t2, ...])`` calls can find
    # their ``prev_step_tuple`` and apply characters incrementally rather
    # than treating the whole sequence as a prompt and resetting to the
    # root parser.
    if not self._disabled:
        try:
            self._enforcer.get_allowed_tokens([])
        except Exception as exc:
            logger.warning(
                "TokenEnforcer bootstrap failed (%s); "
                "falling back to unconstrained generation",
                exc,
            )
            self._disabled = True

    self._prompt_len: int | None = None
    self._vocab_size: int = _get_vocab_size(tokenizer)

    # EOS/stop tokens cache.
    eos_id = getattr(self._tok_data, "eos_token_id", None)
    if isinstance(eos_id, (list, tuple, set)):
        self._eos_set: set[int] = {int(e) for e in eos_id}
    elif eos_id is not None:
        self._eos_set = {int(eos_id)}
    else:
        self._eos_set = set()

    # Pre-compute valid property name prefixes for key-start filtering.
    all_names = _collect_property_names(schema)
    self._valid_key_first_chars: set[str] = {n[0] for n in all_names if n}
    self._valid_key_names: set[str] = all_names

    # Lazy decode cache — populated on demand.
    self._token_decode_cache: dict[int, str | None] = {}

    # Suffix decode cache keyed by length.  Full tokenizer.decode()
    # is always used (incremental per-token decode is incorrect for
    # BPE/SentencePiece tokenizers where whitespace is a token prefix).
    self._cached_suffix_text: str = ""
    self._cached_suffix_len: int = 0

    # Incremental JSON context state — avoids re-scanning the full
    # decoded text on every step.
    self._json_ctx_in_string: bool = False
    self._json_ctx_last_quote_pos: int = -1
    self._json_ctx_scanned_len: int = 0

    # Bracket/brace depth counters for fast _suffix_is_complete_json
    # pre-check.  Updated by _get_json_context incrementally.  JSON
    # is only potentially complete when both are zero.
    self._brace_depth: int = 0
    self._bracket_depth: int = 0

    # Container nesting stack for distinguishing object vs array context.
    # Entries are ``"o"`` (object/brace) or ``"a"`` (array/bracket).
    # Used by ``_get_json_context`` to return ``"key_start"`` only when
    # inside an object — NOT when inside an array after ``,``.
    self._container_stack: list[str] = []

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._tokenizer instance-attribute

_tokenizer = tokenizer

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._schema instance-attribute

_schema = schema

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._tok_data instance-attribute

_tok_data = get_tokenizer_data(tokenizer)

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._disabled instance-attribute

_disabled = False

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._enforcer instance-attribute

_enforcer = TokenEnforcer(self._tok_data, self._parser)

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._parser instance-attribute

_parser = None

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._prompt_len instance-attribute

_prompt_len: int | None = None

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._vocab_size instance-attribute

_vocab_size: int = _get_vocab_size(tokenizer)

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._eos_set instance-attribute

_eos_set: set[int] = {int(e) for e in eos_id}

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._valid_key_first_chars instance-attribute

_valid_key_first_chars: set[str] = {n[0] for n in all_names if n}

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._valid_key_names instance-attribute

_valid_key_names: set[str] = all_names

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._token_decode_cache instance-attribute

_token_decode_cache: dict[int, str | None] = {}

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._cached_suffix_text instance-attribute

_cached_suffix_text: str = ''

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._cached_suffix_len instance-attribute

_cached_suffix_len: int = 0

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._json_ctx_in_string instance-attribute

_json_ctx_in_string: bool = False

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._json_ctx_last_quote_pos instance-attribute

_json_ctx_last_quote_pos: int = -1

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._json_ctx_scanned_len instance-attribute

_json_ctx_scanned_len: int = 0

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._brace_depth instance-attribute

_brace_depth: int = 0

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._bracket_depth instance-attribute

_bracket_depth: int = 0

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._container_stack instance-attribute

_container_stack: list[str] = []

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor.schema property

schema: dict | None

Return the normalized JSON Schema enforced for this request.

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor.vocab_size property

vocab_size: int

Return the tokenizer vocabulary size used to construct masks.

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._suffix

_suffix(tokens_list: list[int]) -> list[int]

Return the slice of tokens that corresponds to generated output.

Source code in vllm_mlx/constrained/json_schema_processor.py

def _suffix(self, tokens_list: list[int]) -> list[int]:
    """Return the slice of ``tokens`` that corresponds to generated output."""
    if self._prompt_len is None:
        # Use len(tokens_list) so the prompt is excluded entirely.
        # The previous ``- 1`` caused the last prompt token (e.g. Gemma-4's
        # <channel|> thinking-end special token) to be prepended to every
        # suffix fed to the enforcer, corrupting the JSON grammar state.
        self._prompt_len = len(tokens_list)
    return tokens_list[self._prompt_len :]

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._decode_token_cached

_decode_token_cached(tok_id: int) -> str | None

Return the decoded text for a single token (cached).

Source code in vllm_mlx/constrained/json_schema_processor.py

def _decode_token_cached(self, tok_id: int) -> str | None:
    """Return the decoded text for a single token (cached)."""
    cached = self._token_decode_cache.get(tok_id)
    if cached is not None:
        return cached
    if tok_id in self._token_decode_cache:
        return None  # previously cached as None
    try:
        decoded = self._tokenizer.decode([tok_id])
    except Exception:
        self._token_decode_cache[tok_id] = None
        return None
    result = decoded if isinstance(decoded, str) else None
    self._token_decode_cache[tok_id] = result
    return result

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._decode_suffix

_decode_suffix(suffix: list[int]) -> str | None

Decode suffix tokens to text.

Always uses full tokenizer.decode(suffix) which is correct for all tokenizer families (BPE, SentencePiece, etc.). Per-token concatenation is NOT safe because whitespace may be encoded as a token prefix (e.g. decode([1526]) = "world" but in context decode([22557, 1526]) = "Hello world").

Results are cached by suffix length to avoid redundant decodes within the same generation step (_get_json_context and _suffix_is_complete_json both call this method).

Source code in vllm_mlx/constrained/json_schema_processor.py

def _decode_suffix(self, suffix: list[int]) -> str | None:
    """Decode suffix tokens to text.

    Always uses full ``tokenizer.decode(suffix)`` which is correct for
    all tokenizer families (BPE, SentencePiece, etc.).  Per-token
    concatenation is NOT safe because whitespace may be encoded as a
    token prefix (e.g. ``decode([1526]) = "world"`` but in context
    ``decode([22557, 1526]) = "Hello world"``).

    Results are cached by suffix length to avoid redundant decodes
    within the same generation step (``_get_json_context`` and
    ``_suffix_is_complete_json`` both call this method).
    """
    if not suffix:
        self._cached_suffix_text = ""
        self._cached_suffix_len = 0
        return ""

    suffix_len = len(suffix)

    # Fast path: already decoded this exact suffix length.
    if suffix_len == self._cached_suffix_len:
        return self._cached_suffix_text

    # Full decode (correct for all tokenizer families).
    try:
        decoded = self._tokenizer.decode(list(suffix))
    except Exception:
        return None
    result = decoded if isinstance(decoded, str) else ""

    # Validate prefix stability for incremental JSON context scanning.
    # decode(tokens[:n]) must be a prefix of decode(tokens[:n+1]) for
    # the incremental scanner in _get_json_context to be correct.
    if (
        suffix_len > self._cached_suffix_len
        and self._cached_suffix_len > 0
        and not result.startswith(self._cached_suffix_text)
    ):
        # Prefix changed — reset incremental context state.
        self._json_ctx_scanned_len = 0
        self._json_ctx_in_string = False
        self._json_ctx_last_quote_pos = -1
        self._brace_depth = 0
        self._bracket_depth = 0
        self._container_stack = []

    self._cached_suffix_text = result
    self._cached_suffix_len = suffix_len
    return result

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._suffix_is_complete_json

_suffix_is_complete_json(suffix: list[int]) -> bool

Return True if the decoded suffix parses as a complete JSON value.

Uses cached bracket/brace depth from _get_json_context as a fast pre-check: JSON cannot be complete when brackets are unbalanced or we are inside a string. This avoids the expensive json.loads call on ~99% of steps.

Source code in vllm_mlx/constrained/json_schema_processor.py

def _suffix_is_complete_json(self, suffix: list[int]) -> bool:
    """Return True if the decoded ``suffix`` parses as a complete JSON value.

    Uses cached bracket/brace depth from ``_get_json_context`` as a
    fast pre-check: JSON cannot be complete when brackets are
    unbalanced or we are inside a string.  This avoids the expensive
    ``json.loads`` call on ~99% of steps.
    """
    if not suffix:
        return False
    # Fast pre-check using cached structural state.
    if self._brace_depth != 0 or self._bracket_depth != 0:
        return False
    if self._json_ctx_in_string:
        return False
    text = self._decode_suffix(suffix)
    if not text:
        return False
    text = text.strip()
    if not text:
        return False
    try:
        json.loads(text)
    except (ValueError, json.JSONDecodeError):
        return False
    return True

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._get_json_context

_get_json_context(suffix: list[int]) -> str

Determine the JSON structural context of the current suffix.

Processes only newly appended characters instead of re-scanning the full decoded text on every call (O(1) amortised per step instead of O(n)).

Returns one of: - "key_start": expecting a new key (after { or ,) - "in_key": inside an open key string - "other": any other position

Source code in vllm_mlx/constrained/json_schema_processor.py

def _get_json_context(self, suffix: list[int]) -> str:
    """Determine the JSON structural context of the current suffix.

    Processes only newly appended characters instead of re-scanning
    the full decoded text on every call (O(1) amortised per step
    instead of O(n)).

    Returns one of:
    - ``"key_start"``: expecting a new key (after ``{`` or ``,``)
    - ``"in_key"``: inside an open key string
    - ``"other"``: any other position
    """
    text = self._decode_suffix(suffix)
    if text is None or not text:
        return "other"

    text_len = len(text)

    if text_len > self._json_ctx_scanned_len and self._json_ctx_scanned_len > 0:
        # Incremental scan: process only new characters.
        in_string = self._json_ctx_in_string
        last_quote_pos = self._json_ctx_last_quote_pos
        brace_depth = self._brace_depth
        bracket_depth = self._bracket_depth
        container_stack = self._container_stack
        i = self._json_ctx_scanned_len
        while i < text_len:
            ch = text[i]
            if in_string:
                if ch == "\\" and i + 1 < text_len:
                    i += 2
                    continue
                if ch == '"':
                    in_string = False
            else:
                if ch == '"':
                    in_string = True
                    last_quote_pos = i
                elif ch == "{":
                    brace_depth += 1
                    container_stack.append("o")
                elif ch == "}":
                    brace_depth -= 1
                    if container_stack and container_stack[-1] == "o":
                        container_stack.pop()
                elif ch == "[":
                    bracket_depth += 1
                    container_stack.append("a")
                elif ch == "]":
                    bracket_depth -= 1
                    if container_stack and container_stack[-1] == "a":
                        container_stack.pop()
            i += 1
        self._json_ctx_in_string = in_string
        self._json_ctx_last_quote_pos = last_quote_pos
        self._json_ctx_scanned_len = text_len
        self._brace_depth = brace_depth
        self._bracket_depth = bracket_depth
    else:
        # Full scan (first call or text shrank/reset).
        in_string = False
        last_quote_pos = -1
        brace_depth = 0
        bracket_depth = 0
        container_stack: list[str] = []
        i = 0
        while i < text_len:
            ch = text[i]
            if in_string:
                if ch == "\\" and i + 1 < text_len:
                    i += 2
                    continue
                if ch == '"':
                    in_string = False
            else:
                if ch == '"':
                    in_string = True
                    last_quote_pos = i
                elif ch == "{":
                    brace_depth += 1
                    container_stack.append("o")
                elif ch == "}":
                    brace_depth -= 1
                    if container_stack and container_stack[-1] == "o":
                        container_stack.pop()
                elif ch == "[":
                    bracket_depth += 1
                    container_stack.append("a")
                elif ch == "]":
                    bracket_depth -= 1
                    if container_stack and container_stack[-1] == "a":
                        container_stack.pop()
            i += 1
        self._json_ctx_in_string = in_string
        self._json_ctx_last_quote_pos = last_quote_pos
        self._json_ctx_scanned_len = text_len
        self._brace_depth = brace_depth
        self._bracket_depth = bracket_depth
        self._container_stack = container_stack

    if self._json_ctx_in_string:
        before = text[: self._json_ctx_last_quote_pos].rstrip()
        if not before or before[-1] in ("{", ","):
            # Only treat as key if we're inside an object, not an array.
            if self._container_stack and self._container_stack[-1] == "a":
                return "other"
            return "in_key"
        return "other"

    stripped = text.rstrip()
    if not stripped:
        return "other"
    if stripped[-1] in ("{", ","):
        # Only return key_start when inside an object.  After a comma
        # inside an array (e.g. ``[{...},``), the next element is a
        # value, not a key — returning "key_start" here would cause
        # _filter_key_start_tokens to block valid array element tokens
        # like ``{``.
        if self._container_stack and self._container_stack[-1] == "a":
            return "other"
        return "key_start"
    return "other"

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._filter_at_key_context

_filter_at_key_context(context: str, suffix: list[int], allowed: list[int]) -> list[int]

Apply schema-aware filtering when in key-related context.

At key_start: only allow tokens that begin a valid key, whitespace, }, or just ". At in_key: only allow tokens compatible with continuing a valid property name (no leading whitespace; content must be a valid prefix).

Source code in vllm_mlx/constrained/json_schema_processor.py

def _filter_at_key_context(
    self, context: str, suffix: list[int], allowed: list[int]
) -> list[int]:
    """Apply schema-aware filtering when in key-related context.

    At ``key_start``: only allow tokens that begin a valid key, whitespace,
    ``}``, or just ``"``.
    At ``in_key``: only allow tokens compatible with continuing a valid
    property name (no leading whitespace; content must be a valid prefix).
    """
    if not self._valid_key_names:
        return allowed  # no schema info → skip filtering

    if context == "key_start":
        return self._filter_key_start_tokens(suffix, allowed)
    elif context == "in_key":
        return self._filter_in_key_tokens(suffix, allowed)
    return allowed

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._filter_key_start_tokens

_filter_key_start_tokens(suffix: list[int], allowed: list[int]) -> list[int]

Filter tokens at key-start position.

Only permit tokens that: - Are whitespace-only (before the key ") - Decode to } (close object) - Start a valid key: " followed by a valid first char

Source code in vllm_mlx/constrained/json_schema_processor.py

def _filter_key_start_tokens(
    self, suffix: list[int], allowed: list[int]
) -> list[int]:
    """Filter tokens at key-start position.

    Only permit tokens that:
    - Are whitespace-only (before the key ``"``)
    - Decode to ``}`` (close object)
    - Start a valid key: ``"`` followed by a valid first char
    """
    result = []
    for tok_id in allowed:
        if tok_id in self._eos_set:
            continue  # EOS at key-start is handled separately
        tok_text = self._decode_token_cached(tok_id)
        if tok_text is None:
            result.append(tok_id)
            continue
        stripped = tok_text.lstrip()
        if not stripped:
            # Pure whitespace — allowed before key
            result.append(tok_id)
            continue
        if stripped[0] == "}":
            # Closing brace — end of object
            result.append(tok_id)
            continue
        if stripped[0] == '"':
            # Opening a key — validate content
            rest = stripped[1:]
            if not rest:
                # Just ``"`` — will be validated on next step
                result.append(tok_id)
                continue
            # Check if rest starts with a valid key character
            if rest[0] in self._valid_key_first_chars:
                # Further check: does the key content (up to closing ``"``)
                # match a prefix of a known property name?
                close_idx = rest.find('"')
                if close_idx < 0:
                    # Key not yet closed — check prefix
                    if self._is_valid_key_prefix(rest):
                        result.append(tok_id)
                else:
                    # Key fully contained in this token
                    key_name = rest[:close_idx]
                    if key_name in self._valid_key_names:
                        result.append(tok_id)
                continue
            # First char not in valid set → skip
            continue
        # Other chars (digits, letters without quote) → skip at key-start
        continue
    return result if result else allowed  # safety: never return empty

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._filter_in_key_tokens

_filter_in_key_tokens(suffix: list[int], allowed: list[int]) -> list[int]

Filter tokens when we're inside an open key string.

Only allow tokens whose content continues a valid property name. Reject whitespace-only/leading-whitespace tokens.

Source code in vllm_mlx/constrained/json_schema_processor.py

def _filter_in_key_tokens(self, suffix: list[int], allowed: list[int]) -> list[int]:
    """Filter tokens when we're inside an open key string.

    Only allow tokens whose content continues a valid property name.
    Reject whitespace-only/leading-whitespace tokens.
    """
    # Figure out what key content we've accumulated so far.
    text = self._decode_suffix(suffix)
    if text is None:
        return allowed

    # Find the last unmatched ``"`` — everything after it is key content
    # accumulated so far.
    last_open = text.rfind('"')
    if last_open < 0:
        return allowed
    key_so_far = text[last_open + 1 :]

    result = []
    for tok_id in allowed:
        tok_text = self._decode_token_cached(tok_id)
        if tok_text is None:
            result.append(tok_id)
            continue
        # Token must not start with whitespace (no ws inside keys)
        if tok_text and tok_text[0] in (" ", "\t", "\n", "\r"):
            continue
        # Check if key_so_far + tok_text is a valid key prefix
        candidate = key_so_far + tok_text
        # If the closing ``"`` is in tok_text, extract the full key
        close_idx = tok_text.find('"')
        if close_idx >= 0:
            full_key = key_so_far + tok_text[:close_idx]
            if full_key in self._valid_key_names:
                result.append(tok_id)
        else:
            # Key still open — check if it's a valid prefix
            if self._is_valid_key_prefix(candidate):
                result.append(tok_id)
    return result if result else allowed

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._is_valid_key_prefix

_is_valid_key_prefix(prefix: str) -> bool

Return True if prefix is a prefix of at least one valid key name.

Source code in vllm_mlx/constrained/json_schema_processor.py

def _is_valid_key_prefix(self, prefix: str) -> bool:
    """Return True if *prefix* is a prefix of at least one valid key name."""
    return any(name.startswith(prefix) for name in self._valid_key_names)

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._filter_nonprogress_whitespace_tokens

_filter_nonprogress_whitespace_tokens(suffix: list[int], allowed: list[int]) -> list[int]

Stop constrained JSON from spending a long run on pure whitespace.

JSON permits arbitrary whitespace around structural tokens. That is valid, but with non-streaming requests a model can keep selecting whitespace-only tokens for minutes without producing useful JSON content. Once the decoded suffix has a long trailing whitespace run outside a string, remove pure-whitespace tokens from the next-step allowed set so generation must make structural/content progress.

Source code in vllm_mlx/constrained/json_schema_processor.py

def _filter_nonprogress_whitespace_tokens(
    self, suffix: list[int], allowed: list[int]
) -> list[int]:
    """Stop constrained JSON from spending a long run on pure whitespace.

    JSON permits arbitrary whitespace around structural tokens. That is
    valid, but with non-streaming requests a model can keep selecting
    whitespace-only tokens for minutes without producing useful JSON
    content. Once the decoded suffix has a long trailing whitespace run
    outside a string, remove pure-whitespace tokens from the next-step
    allowed set so generation must make structural/content progress.
    """
    text = self._decode_suffix(suffix)
    if text is None or not text:
        return allowed

    trailing = len(text) - len(text.rstrip(" \t\r\n"))
    if trailing < _MAX_NONPROGRESS_WHITESPACE_CHARS:
        return allowed

    filtered: list[int] = []
    for tok_id in allowed:
        tok_text = self._decode_token_cached(tok_id)
        if tok_text is None:
            filtered.append(tok_id)
            continue
        if tok_text == "" or all(ch in _JSON_WHITESPACE for ch in tok_text):
            continue
        filtered.append(tok_id)
    return filtered if filtered else allowed

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._build_allow_mask

_build_allow_mask(allowed: list[int], vocab_size: int) -> array

Build a 1-D mask of length vocab_size where allowed positions are 0 and disallowed positions are -inf.

Uses numpy for mask construction (C-level speed) instead of a Python loop over vocab_size elements.

Source code in vllm_mlx/constrained/json_schema_processor.py

def _build_allow_mask(self, allowed: list[int], vocab_size: int) -> mx.array:
    """
    Build a 1-D mask of length ``vocab_size`` where allowed positions are
    ``0`` and disallowed positions are ``-inf``.

    Uses numpy for mask construction (C-level speed) instead of a
    Python loop over ``vocab_size`` elements.
    """
    if not allowed:
        return mx.full((vocab_size,), -float("inf"))
    allowed_clamped = [i for i in allowed if 0 <= i < vocab_size]
    if not allowed_clamped:
        return mx.full((vocab_size,), -float("inf"))
    buf = np.full(vocab_size, -np.inf, dtype=np.float32)
    buf[allowed_clamped] = 0.0
    return mx.array(buf)

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor.__call__

__call__(tokens: array, logits: array) -> array

Apply the allowed-tokens mask to logits.

Source code in vllm_mlx/constrained/json_schema_processor.py

def __call__(self, tokens: mx.array, logits: mx.array) -> mx.array:
    """Apply the allowed-tokens mask to ``logits``."""
    if self._disabled:
        # Force EOS immediately — the enforcer got stuck, so continued
        # generation produces garbage.  Without this cap the model would
        # generate up to max_tokens (often 262 K) of useless output,
        # blocking the slot for minutes/hours.
        return _eos_logits_or_original(
            self._eos_set,
            logits,
            self._build_allow_mask,
        )

    try:
        tokens_list = tokens.tolist() if hasattr(tokens, "tolist") else list(tokens)
        if isinstance(tokens_list, int):
            tokens_list = [tokens_list]
        elif tokens_list and isinstance(tokens_list[0], list):
            tokens_list = tokens_list[0]

        suffix = self._suffix(tokens_list)
        eos_logits = _complete_json_eos_logits(
            self._eos_set,
            suffix,
            logits,
            self._suffix_is_complete_json,
            self._build_allow_mask,
        )
        if eos_logits is not None:
            return eos_logits

        # Use prompt_len directly instead of O(n) list comparison.
        pass_to_enforcer = suffix if self._prompt_len else tokens_list
        allowed_result = self._enforcer.get_allowed_tokens(pass_to_enforcer)
        allowed = getattr(allowed_result, "allowed_tokens", allowed_result)
        if allowed is None:
            return logits

        allowed_list = list(allowed)

        # --- Schema-aware key filter (before EOS guard so that the
        # incremental JSON context state and bracket depth counters
        # are up-to-date for the _suffix_is_complete_json pre-check).
        context = self._get_json_context(suffix)
        if not self._json_ctx_in_string:
            allowed_list = self._filter_nonprogress_whitespace_tokens(
                suffix, allowed_list
            )
        if context in ("key_start", "in_key"):
            allowed_list = self._filter_at_key_context(
                context, suffix, allowed_list
            )

        # --- EOS guard: only permit EOS when output is valid JSON ---
        if (
            self._eos_set
            and any(t in self._eos_set for t in allowed_list)
            and not self._suffix_is_complete_json(suffix)
        ):
            allowed_list = [t for t in allowed_list if t not in self._eos_set]

        # --- Recovery: if enforcer returns empty set AND output is not
        # complete JSON, the schema is likely unsupported — disable the
        # processor and let the model generate freely (system prompt +
        # post-validation still apply).  Only force EOS if the output
        # already parses as valid JSON (generation is done).
        if not allowed_list:
            if self._suffix_is_complete_json(suffix) and self._eos_set:
                allowed_list = sorted(self._eos_set)
            else:
                try:
                    decoded = self._tokenizer.decode(suffix)
                except Exception:
                    decoded = "<decode-error>"
                logger.warning(
                    "JSONLP: enforcer stuck (empty allowed-set at "
                    "suffix_len=%d, suffix_tokens=%s, decoded=%r); "
                    "disabling constrained decoding for this request",
                    len(suffix),
                    suffix[:10],
                    decoded[:80],
                )
                self._disabled = True
                return logits

        actual_vocab = logits.shape[-1]
        mask = self._build_allow_mask(allowed_list, actual_vocab)
        if logits.ndim == 2 and logits.shape[0] == 1:
            mask = mask[None, :]
        return logits + mask
    except Exception as exc:  # pragma: no cover - defensive
        logger.error(
            "JSONSchemaLogitsProcessor crashed; disabling for this request: %s",
            exc,
        )
        self._disabled = True
        return logits

vllm_mlx.constrained.json_schema_processor._canonical_schema_key

_canonical_schema_key(schema: dict | None) -> str

Source code in vllm_mlx/constrained/json_schema_processor.py

def _canonical_schema_key(schema: dict | None) -> str:
    if schema is None:
        return "__none__"
    blob = json.dumps(schema, sort_keys=True, separators=(",", ":")).encode()
    return hashlib.sha256(blob).hexdigest()

vllm_mlx.constrained.json_schema_processor._get_or_build_parser

_get_or_build_parser(schema: dict | None) -> tuple[dict, Any]

Return (parser_schema, JsonSchemaParser) for schema, memoised.

Source code in vllm_mlx/constrained/json_schema_processor.py

def _get_or_build_parser(schema: dict | None) -> tuple[dict, Any]:
    """Return (parser_schema, JsonSchemaParser) for ``schema``, memoised."""
    from lmformatenforcer import JsonSchemaParser

    key = _canonical_schema_key(schema)
    cached = _parser_cache.get(key)
    if cached is not None:
        return cached

    if schema is None:
        parser_schema = _GENERIC_JSON_SCHEMA
    else:
        parser_schema = _simplify_schema(schema)
        parser_schema = _force_no_additional_properties(parser_schema)
    parser = JsonSchemaParser(parser_schema)
    _parser_cache[key] = (parser_schema, parser)
    return parser_schema, parser

vllm_mlx.constrained.json_schema_processor.is_available

is_available() -> bool

Return True iff lm-format-enforcer is importable.

Source code in vllm_mlx/constrained/json_schema_processor.py

def is_available() -> bool:
    """Return ``True`` iff ``lm-format-enforcer`` is importable."""
    try:
        import lmformatenforcer  # noqa: F401
    except ImportError:
        return False
    return True

vllm_mlx.constrained.json_schema_processor._simplify_schema

_simplify_schema(schema: dict) -> dict

Pre-process a JSON Schema for lm-format-enforcer compatibility.

lm-format-enforcer does not support $ref, not, type as an array, or recursive definitions. This function:

1. Resolves $ref by inlining referenced definitions (with cycle detection so recursive definitions are truncated to {}).
2. Removes not sub-schemas (makes the schema more permissive).
3. Strips metadata / serialisation-hint keywords that the enforcer does not understand: default, examples, title, description, $schema, $id.
4. Converts type: [t1, t2, ...] to anyOf: [{type: t1}, ...].
5. Cleans up empty anyOf / oneOf branches.
6. Flattens nested anyOf/oneOf (e.g. anyOf: [{anyOf: [A, B]}, C] → anyOf: [A, B, C]).

Source code in vllm_mlx/constrained/json_schema_processor.py

def _simplify_schema(schema: dict) -> dict:
    """Pre-process a JSON Schema for ``lm-format-enforcer`` compatibility.

    ``lm-format-enforcer`` does not support ``$ref``, ``not``, ``type`` as an
    array, or recursive definitions.  This function:

    1. Resolves ``$ref`` by inlining referenced definitions (with cycle
       detection so recursive definitions are truncated to ``{}``).
    2. Removes ``not`` sub-schemas (makes the schema more permissive).
    3. Strips metadata / serialisation-hint keywords that the enforcer does
       not understand: ``default``, ``examples``, ``title``, ``description``,
       ``$schema``, ``$id``.
    4. Converts ``type: [t1, t2, ...]`` to ``anyOf: [{type: t1}, ...]``.
    5. Cleans up empty ``anyOf`` / ``oneOf`` branches.
    6. Flattens nested ``anyOf``/``oneOf`` (e.g.
       ``anyOf: [{anyOf: [A, B]}, C]`` → ``anyOf: [A, B, C]``).
    """
    schema = copy.deepcopy(schema)
    definitions: dict = {}
    definitions.update(schema.pop("definitions", {}))
    definitions.update(schema.pop("$defs", {}))

    resolving: set[str] = set()  # cycle guard

    def _resolve(node: Any, depth: int = 0) -> Any:
        if depth > 12 or not isinstance(node, dict):
            return node

        # --- resolve $ref --------------------------------------------------
        if "$ref" in node:
            ref: str = node["$ref"]
            parts = ref.split("/")
            if (
                len(parts) == 3
                and parts[0] == "#"
                and parts[1] in ("definitions", "$defs")
            ):
                name = parts[2]
                if name in definitions and ref not in resolving:
                    resolving.add(ref)
                    resolved = copy.deepcopy(definitions[name])
                    # Merge extra keys (e.g. ``default``) from the $ref node.
                    for k, v in node.items():
                        if k != "$ref" and k not in resolved:
                            resolved[k] = v
                    result = _resolve(resolved, depth + 1)
                    resolving.discard(ref)
                    return result
            # Circular or unresolvable — return empty (= any).
            return {}

        # --- remove unsupported keywords -----------------------------------
        node.pop("not", None)
        node.pop("$schema", None)
        node.pop("$id", None)
        # Metadata / serialisation hints that lm-format-enforcer doesn't
        # understand — keeping them causes the parser to mis-navigate.
        node.pop("default", None)
        node.pop("examples", None)
        node.pop("title", None)
        node.pop("description", None)

        # --- type array → anyOf --------------------------------------------
        if isinstance(node.get("type"), list):
            types = node.pop("type")
            items_schema = node.pop("items", None)
            branches: list[dict] = []
            for t in types:
                branch: dict[str, Any] = {"type": t}
                if t == "array" and items_schema is not None:
                    branch["items"] = _resolve(copy.deepcopy(items_schema), depth + 1)
                branches.append(branch)
            existing = node.pop("anyOf", [])
            node["anyOf"] = existing + branches

        # --- recurse into sub-schemas --------------------------------------
        if "properties" in node and isinstance(node["properties"], dict):
            for k in list(node["properties"]):
                node["properties"][k] = _resolve(node["properties"][k], depth + 1)

        for key in ("items", "additionalProperties"):
            if key in node and isinstance(node[key], dict):
                node[key] = _resolve(node[key], depth + 1)

        for key in ("allOf", "anyOf", "oneOf"):
            if key in node and isinstance(node[key], list):
                # Resolve each branch; drop empty dicts (= "any", redundant
                # inside anyOf since they make the whole constraint trivially
                # true — but keeping one "any" branch confuses the enforcer).
                resolved_items = [_resolve(item, depth + 1) for item in node[key]]
                node[key] = [it for it in resolved_items if it != {}]
                if not node[key]:
                    del node[key]

        # --- flatten nested anyOf/oneOf ------------------------------------
        # ``anyOf: [{anyOf: [A, B]}, C]`` → ``anyOf: [A, B, C]`` when the
        # wrapper dict has no extra keys.  This removes one level of nesting
        # that confuses lm-format-enforcer's UnionParser.
        for key in ("anyOf", "oneOf"):
            if key in node and isinstance(node[key], list):
                flattened: list[Any] = []
                for item in node[key]:
                    if isinstance(item, dict) and key in item and len(item) == 1:
                        flattened.extend(item[key])
                    else:
                        flattened.append(item)
                node[key] = flattened

        return node

    return _resolve(schema)

vllm_mlx.constrained.json_schema_processor._force_no_additional_properties

_force_no_additional_properties(schema: dict) -> dict

Return a deep copy of schema with additionalProperties: false injected into every object-type sub-schema that declares properties.

lm-format-enforcer has a bug where multi-character tokens spanning JSON structural boundaries (e.g., a single token that decodes to "") can produce empty or whitespace-only keys, causing KeyError crashes in jsonschemaparser.py. Setting additionalProperties: false tells the enforcer's trie traversal that only the declared property names are valid keys, which significantly narrows the allowed tokens and prevents most of these boundary-spanning issues.

Source code in vllm_mlx/constrained/json_schema_processor.py

def _force_no_additional_properties(schema: dict) -> dict:
    """Return a deep copy of *schema* with ``additionalProperties: false``
    injected into every object-type sub-schema that declares ``properties``.

    ``lm-format-enforcer`` has a bug where multi-character tokens spanning
    JSON structural boundaries (e.g., a single token that decodes to ``""``)
    can produce empty or whitespace-only keys, causing ``KeyError`` crashes in
    ``jsonschemaparser.py``.  Setting ``additionalProperties: false`` tells the
    enforcer's trie traversal that only the declared property names are valid
    keys, which significantly narrows the allowed tokens and prevents most of
    these boundary-spanning issues.
    """
    schema = copy.deepcopy(schema)
    _inject_no_additional_props(schema)
    return schema

vllm_mlx.constrained.json_schema_processor._inject_no_additional_props

_inject_no_additional_props(node: Any) -> None

Recursively inject additionalProperties: false into node.

Source code in vllm_mlx/constrained/json_schema_processor.py

def _inject_no_additional_props(node: Any) -> None:
    """Recursively inject ``additionalProperties: false`` into *node*."""
    if not isinstance(node, dict):
        return
    if "properties" in node and "additionalProperties" not in node:
        node["additionalProperties"] = False
    for value in node.values():
        if isinstance(value, dict):
            _inject_no_additional_props(value)
        elif isinstance(value, list):
            for item in value:
                _inject_no_additional_props(item)

vllm_mlx.constrained.json_schema_processor._collect_property_names

_collect_property_names(schema: dict | None) -> set[str]

Collect all property names declared anywhere in schema.

Source code in vllm_mlx/constrained/json_schema_processor.py

def _collect_property_names(schema: dict | None) -> set[str]:
    """Collect all property names declared anywhere in *schema*."""
    names: set[str] = set()
    if schema is None:
        return names
    _walk_properties(schema, names)
    return names

vllm_mlx.constrained.json_schema_processor._walk_properties

_walk_properties(node: Any, names: set[str]) -> None

Source code in vllm_mlx/constrained/json_schema_processor.py

def _walk_properties(node: Any, names: set[str]) -> None:
    if not isinstance(node, dict):
        return
    props = node.get("properties")
    if isinstance(props, dict):
        names.update(props.keys())
        for v in props.values():
            _walk_properties(v, names)
    for key in ("items", "additionalProperties", "not"):
        if key in node and isinstance(node[key], dict):
            _walk_properties(node[key], names)
    for key in ("allOf", "anyOf", "oneOf"):
        if key in node and isinstance(node[key], list):
            for item in node[key]:
                _walk_properties(item, names)

vllm_mlx.constrained.json_schema_processor._complete_json_eos_logits

_complete_json_eos_logits(eos_set: set[int], suffix: list[int], logits: array, is_complete_json, build_allow_mask) -> array | None

Source code in vllm_mlx/constrained/json_schema_processor.py

def _complete_json_eos_logits(
    eos_set: set[int],
    suffix: list[int],
    logits: mx.array,
    is_complete_json,
    build_allow_mask,
) -> mx.array | None:
    if not eos_set or not is_complete_json(suffix):
        return None
    return _eos_logits(eos_set, logits, build_allow_mask)

vllm_mlx.constrained.json_schema_processor._eos_logits

_eos_logits(eos_set: set[int], logits: array, build_allow_mask) -> array | None

Source code in vllm_mlx/constrained/json_schema_processor.py

def _eos_logits(
    eos_set: set[int],
    logits: mx.array,
    build_allow_mask,
) -> mx.array | None:
    if not eos_set:
        return None
    actual_vocab = logits.shape[-1]
    mask = build_allow_mask(sorted(eos_set), actual_vocab)
    if logits.ndim == 2 and logits.shape[0] == 1:
        mask = mask[None, :]
    return logits + mask

vllm_mlx.constrained.json_schema_processor._eos_logits_or_original

_eos_logits_or_original(eos_set: set[int], logits: array, build_allow_mask) -> array

Source code in vllm_mlx/constrained/json_schema_processor.py

def _eos_logits_or_original(
    eos_set: set[int],
    logits: mx.array,
    build_allow_mask,
) -> mx.array:
    masked = _eos_logits(eos_set, logits, build_allow_mask)
    return logits if masked is None else masked

Complete contract reference

Expand any definition for its exact inputs, annotations, defaults, return contract, directly raised exceptions, source-grounded behavior, and immutable line link. This section includes private and nested definitions that ordinary API generators omit.

vllm_mlx.constrained.json_schema_processor.LMFormatEnforcerNotAvailableError · class

vllm_mlx.constrained.json_schema_processor.LMFormatEnforcerNotAvailableError()

Raised when lm-format-enforcer is required but not installed.

Parameters

This callable has no explicit inputs.

Returns

• Constructs: vllm_mlx.constrained.json_schema_processor.LMFormatEnforcerNotAvailableError

Exceptions and behavior

Class LMFormatEnforcerNotAvailableError derives from RuntimeError and declares 0 direct member(s). No direct raise statement appears in this definition.

View source #L33-L34.

vllm_mlx.constrained.json_schema_processor._canonical_schema_key · function

vllm_mlx.constrained.json_schema_processor._canonical_schema_key(schema: dict | None) -> str

Function _canonical_schema_key calls json.dumps(schema, sort_keys=True, separators=(',', ':')).encode, json.dumps, hashlib.sha256(blob).hexdigest, hashlib.sha256; has 2 explicit return paths.

Parameters

Name	Type	Required	Default	Description
schema	dict \| None	yes	none	Required positional or keyword input.

Returns

• Type: str
• Direct return expressions: '__none__'; hashlib.sha256(blob).hexdigest()

Exceptions and behavior

Function _canonical_schema_key calls json.dumps(schema, sort_keys=True, separators=(',', ':')).encode, json.dumps, hashlib.sha256(blob).hexdigest, hashlib.sha256; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L50-L54.

vllm_mlx.constrained.json_schema_processor._get_or_build_parser · function

vllm_mlx.constrained.json_schema_processor._get_or_build_parser(schema: dict | None) -> tuple[dict, Any]

Return (parser_schema, JsonSchemaParser) for schema, memoised.

Parameters

Name	Type	Required	Default	Description
schema	dict \| None	yes	none	Required positional or keyword input.

Returns

• Type: tuple[dict, Any]
• Direct return expressions: cached; (parser_schema, parser)

Exceptions and behavior

Function _get_or_build_parser calls _canonical_schema_key, _parser_cache.get, _simplify_schema, _force_no_additional_properties; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L57-L73.

vllm_mlx.constrained.json_schema_processor.is_available · function

vllm_mlx.constrained.json_schema_processor.is_available() -> bool

Return True iff lm-format-enforcer is importable.

Parameters

This callable has no explicit inputs.

Returns

• Type: bool
• Direct return expressions: False; True

Exceptions and behavior

Function is_available has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L76-L82.

vllm_mlx.constrained.json_schema_processor._simplify_schema · function

vllm_mlx.constrained.json_schema_processor._simplify_schema(schema: dict) -> dict

Pre-process a JSON Schema for lm-format-enforcer compatibility.

Parameters

Name	Type	Required	Default	Description
schema	dict	yes	none	Required positional or keyword input.

Returns

• Type: dict
• Direct return expressions: _resolve(schema)

Exceptions and behavior

Function _simplify_schema calls copy.deepcopy, definitions.update, schema.pop, set; returns _resolve(schema). No direct raise statement appears in this definition.

View source #L97-L207.

vllm_mlx.constrained.json_schema_processor._simplify_schema._resolve · nested function

vllm_mlx.constrained.json_schema_processor._simplify_schema._resolve(node: Any, depth: int = 0) -> Any

Nested Function _simplify_schema._resolve calls isinstance, ref.split, len, resolving.add; has 3 explicit return paths.

Parameters

Name	Type	Required	Default	Description
node	Any	yes	none	Required positional or keyword input.
depth	int	no	0	Optional positional or keyword input; defaults to 0.

Returns

• Type: Any
• Direct return expressions: node; result; {}

Exceptions and behavior

Nested Function _simplify_schema._resolve calls isinstance, ref.split, len, resolving.add; has 3 explicit return paths. No direct raise statement appears in this definition.

View source #L121-L205.

vllm_mlx.constrained.json_schema_processor._force_no_additional_properties · function

vllm_mlx.constrained.json_schema_processor._force_no_additional_properties(schema: dict) -> dict

Return a deep copy of schema with additionalProperties: false injected into every object-type sub-schema that declares properties.

Parameters

Name	Type	Required	Default	Description
schema	dict	yes	none	Required positional or keyword input.

Returns

• Type: dict
• Direct return expressions: schema

Exceptions and behavior

Function _force_no_additional_properties calls copy.deepcopy, _inject_no_additional_props; returns schema. No direct raise statement appears in this definition.

View source #L210-L224.

vllm_mlx.constrained.json_schema_processor._inject_no_additional_props · function

vllm_mlx.constrained.json_schema_processor._inject_no_additional_props(node: Any) -> None

Recursively inject additionalProperties: false into node.

Parameters

Name	Type	Required	Default	Description
node	Any	yes	none	Required positional or keyword input.

Returns

• Type: None
• Direct return expressions: None

Exceptions and behavior

Function _inject_no_additional_props calls isinstance, node.values, _inject_no_additional_props; returns None. No direct raise statement appears in this definition.

View source #L227-L238.

vllm_mlx.constrained.json_schema_processor._collect_property_names · function

vllm_mlx.constrained.json_schema_processor._collect_property_names(schema: dict | None) -> set[str]

Collect all property names declared anywhere in schema.

Parameters

Name	Type	Required	Default	Description
schema	dict \| None	yes	none	Required positional or keyword input.

Returns

• Type: set[str]
• Direct return expressions: names

Exceptions and behavior

Function _collect_property_names calls set, _walk_properties; returns names. No direct raise statement appears in this definition.

View source #L241-L247.

vllm_mlx.constrained.json_schema_processor._walk_properties · function

vllm_mlx.constrained.json_schema_processor._walk_properties(node: Any, names: set[str]) -> None

Function _walk_properties calls isinstance, node.get, names.update, props.keys; returns None.

Parameters

Name	Type	Required	Default	Description
node	Any	yes	none	Required positional or keyword input.
names	set[str]	yes	none	Required positional or keyword input.

Returns

• Type: None
• Direct return expressions: None

Exceptions and behavior

Function _walk_properties calls isinstance, node.get, names.update, props.keys; returns None. No direct raise statement appears in this definition.

View source #L250-L264.

vllm_mlx.constrained.json_schema_processor._complete_json_eos_logits · function

vllm_mlx.constrained.json_schema_processor._complete_json_eos_logits(eos_set: set[int], suffix: list[int], logits: mx.array, is_complete_json, build_allow_mask) -> mx.array | None

Function _complete_json_eos_logits calls is_complete_json, _eos_logits; has 2 explicit return paths.

Parameters

Name	Type	Required	Default	Description
eos_set	set[int]	yes	none	Required positional or keyword input.
suffix	list[int]	yes	none	Required positional or keyword input.
logits	mx.array	yes	none	Required positional or keyword input.
is_complete_json	not annotated	yes	none	Required positional or keyword input.
build_allow_mask	not annotated	yes	none	Required positional or keyword input.

Returns

• Type: mx.array | None
• Direct return expressions: None; _eos_logits(eos_set, logits, build_allow_mask)

Exceptions and behavior

Function _complete_json_eos_logits calls is_complete_json, _eos_logits; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L267-L276.

vllm_mlx.constrained.json_schema_processor._eos_logits · function

vllm_mlx.constrained.json_schema_processor._eos_logits(eos_set: set[int], logits: mx.array, build_allow_mask) -> mx.array | None

Function _eos_logits calls build_allow_mask, sorted; has 2 explicit return paths.

Parameters

Name	Type	Required	Default	Description
eos_set	set[int]	yes	none	Required positional or keyword input.
logits	mx.array	yes	none	Required positional or keyword input.
build_allow_mask	not annotated	yes	none	Required positional or keyword input.

Returns

• Type: mx.array | None
• Direct return expressions: None; logits + mask

Exceptions and behavior

Function _eos_logits calls build_allow_mask, sorted; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L279-L290.

vllm_mlx.constrained.json_schema_processor._eos_logits_or_original · function

vllm_mlx.constrained.json_schema_processor._eos_logits_or_original(eos_set: set[int], logits: mx.array, build_allow_mask) -> mx.array

Function _eos_logits_or_original calls _eos_logits; returns logits if masked is None else masked.

Parameters

Name	Type	Required	Default	Description
eos_set	set[int]	yes	none	Required positional or keyword input.
logits	mx.array	yes	none	Required positional or keyword input.
build_allow_mask	not annotated	yes	none	Required positional or keyword input.

Returns

• Type: mx.array
• Direct return expressions: logits if masked is None else masked

Exceptions and behavior

Function _eos_logits_or_original calls _eos_logits; returns logits if masked is None else masked. No direct raise statement appears in this definition.

View source #L293-L299.

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor · class

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor(schema: dict | None, tokenizer: Any)

Logits processor that constrains generation to valid JSON.

Parameters

Name	Type	Required	Default	Description
schema	dict \| None	yes	none	Required positional or keyword input.
tokenizer	Any	yes	none	Required positional or keyword input.

Returns

• Constructs: vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor

Exceptions and behavior

Class JSONSchemaLogitsProcessor declares 15 direct member(s). No direct raise statement appears in this definition.

View source #L302-L924.

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor.__init__ · method

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor.__init__(schema: dict | None, tokenizer: Any) -> None

Method JSONSchemaLogitsProcessor.__init__ updates self._tokenizer, self._schema, self._tok_data, self._disabled; calls is_available, LMFormatEnforcerNotAvailableError, get_tokenizer_data, _get_or_build_parser; can raise LMFormatEnforcerNotAvailableError.

Parameters

Name	Type	Required	Default	Description
schema	dict \| None	yes	none	Required positional or keyword input.
tokenizer	Any	yes	none	Required positional or keyword input.

Returns

• Type: None

Exceptions and behavior

Method JSONSchemaLogitsProcessor.__init__ updates self._tokenizer, self._schema, self._tok_data, self._disabled; calls is_available, LMFormatEnforcerNotAvailableError, get_tokenizer_data, _get_or_build_parser; can raise LMFormatEnforcerNotAvailableError. Directly raised exceptions: LMFormatEnforcerNotAvailableError.

View source #L317-L414.

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._suffix · method

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._suffix(tokens_list: list[int]) -> list[int]

Return the slice of tokens that corresponds to generated output.

Parameters

Name	Type	Required	Default	Description
tokens_list	list[int]	yes	none	Required positional or keyword input.

Returns

• Type: list[int]
• Direct return expressions: tokens_list[self._prompt_len:]

Exceptions and behavior

Method JSONSchemaLogitsProcessor._suffix updates self._prompt_len; calls len; returns tokens_list[self._prompt_len:]. No direct raise statement appears in this definition.

View source #L418-L426.

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._decode_token_cached · method

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._decode_token_cached(tok_id: int) -> str | None

Return the decoded text for a single token (cached).

Parameters

Name	Type	Required	Default	Description
tok_id	int	yes	none	Required positional or keyword input.

Returns

• Type: str | None
• Direct return expressions: cached; None; result

Exceptions and behavior

Method JSONSchemaLogitsProcessor._decode_token_cached calls self._token_decode_cache.get, self._tokenizer.decode, isinstance; has 3 explicit return paths. No direct raise statement appears in this definition.

View source #L428-L442.

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._decode_suffix · method

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._decode_suffix(suffix: list[int]) -> str | None

Decode suffix tokens to text.

Parameters

Name	Type	Required	Default	Description
suffix	list[int]	yes	none	Required positional or keyword input.

Returns

• Type: str | None
• Direct return expressions: ''; self._cached_suffix_text; None; result

Exceptions and behavior

Method JSONSchemaLogitsProcessor._decode_suffix updates self._cached_suffix_text, self._cached_suffix_len, self._json_ctx_scanned_len, self._json_ctx_in_string; calls len, self._tokenizer.decode, list, isinstance; has 4 explicit return paths. No direct raise statement appears in this definition.

View source #L444-L493.

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._suffix_is_complete_json · method

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._suffix_is_complete_json(suffix: list[int]) -> bool

Return True if the decoded suffix parses as a complete JSON value.

Parameters

Name	Type	Required	Default	Description
suffix	list[int]	yes	none	Required positional or keyword input.

Returns

• Type: bool
• Direct return expressions: False; True

Exceptions and behavior

Method JSONSchemaLogitsProcessor._suffix_is_complete_json calls self._decode_suffix, text.strip, json.loads; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L495-L520.

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._get_json_context · method

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._get_json_context(suffix: list[int]) -> str

Determine the JSON structural context of the current suffix.

Parameters

Name	Type	Required	Default	Description
suffix	list[int]	yes	none	Required positional or keyword input.

Returns

• Type: str
• Direct return expressions: 'other'; 'in_key'; 'key_start'

Exceptions and behavior

Method JSONSchemaLogitsProcessor._get_json_context updates self._json_ctx_in_string, self._json_ctx_last_quote_pos, self._json_ctx_scanned_len, self._brace_depth; calls self._decode_suffix, len, container_stack.append, container_stack.pop; has 3 explicit return paths. No direct raise statement appears in this definition.

View source #L522-L643.

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._filter_at_key_context · method

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._filter_at_key_context(context: str, suffix: list[int], allowed: list[int]) -> list[int]

Apply schema-aware filtering when in key-related context.

Parameters

Name	Type	Required	Default	Description
context	str	yes	none	Required positional or keyword input.
suffix	list[int]	yes	none	Required positional or keyword input.
allowed	list[int]	yes	none	Required positional or keyword input.

Returns

• Type: list[int]
• Direct return expressions: allowed; self._filter_key_start_tokens(suffix, allowed); self._filter_in_key_tokens(suffix, allowed)

Exceptions and behavior

Method JSONSchemaLogitsProcessor._filter_at_key_context calls self._filter_key_start_tokens, self._filter_in_key_tokens; has 3 explicit return paths. No direct raise statement appears in this definition.

View source #L645-L662.

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._filter_key_start_tokens · method

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._filter_key_start_tokens(suffix: list[int], allowed: list[int]) -> list[int]

Filter tokens at key-start position.

Parameters

Name	Type	Required	Default	Description
suffix	list[int]	yes	none	Required positional or keyword input.
allowed	list[int]	yes	none	Required positional or keyword input.

Returns

• Type: list[int]
• Direct return expressions: result if result else allowed

Exceptions and behavior

Method JSONSchemaLogitsProcessor._filter_key_start_tokens calls self._decode_token_cached, result.append, tok_text.lstrip, rest.find; returns result if result else allowed. No direct raise statement appears in this definition.

View source #L664-L717.

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._filter_in_key_tokens · method

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._filter_in_key_tokens(suffix: list[int], allowed: list[int]) -> list[int]

Filter tokens when we're inside an open key string.

Parameters

Name	Type	Required	Default	Description
suffix	list[int]	yes	none	Required positional or keyword input.
allowed	list[int]	yes	none	Required positional or keyword input.

Returns

• Type: list[int]
• Direct return expressions: allowed; result if result else allowed

Exceptions and behavior

Method JSONSchemaLogitsProcessor._filter_in_key_tokens calls self._decode_suffix, text.rfind, self._decode_token_cached, result.append; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L719-L758.

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._is_valid_key_prefix · method

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._is_valid_key_prefix(prefix: str) -> bool

Return True if prefix is a prefix of at least one valid key name.

Parameters

Name	Type	Required	Default	Description
prefix	str	yes	none	Required positional or keyword input.

Returns

• Type: bool
• Direct return expressions: any((name.startswith(prefix) for name in self._valid_key_names))

Exceptions and behavior

Method JSONSchemaLogitsProcessor._is_valid_key_prefix calls any, name.startswith; returns any((name.startswith(prefix) for name in self._valid_key_names)). No direct raise statement appears in this definition.

View source #L760-L762.

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._filter_nonprogress_whitespace_tokens · method

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._filter_nonprogress_whitespace_tokens(suffix: list[int], allowed: list[int]) -> list[int]

Stop constrained JSON from spending a long run on pure whitespace.

Parameters

Name	Type	Required	Default	Description
suffix	list[int]	yes	none	Required positional or keyword input.
allowed	list[int]	yes	none	Required positional or keyword input.

Returns

• Type: list[int]
• Direct return expressions: allowed; filtered if filtered else allowed

Exceptions and behavior

Method JSONSchemaLogitsProcessor._filter_nonprogress_whitespace_tokens calls self._decode_suffix, len, text.rstrip, self._decode_token_cached; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L764-L793.

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._build_allow_mask · method

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor._build_allow_mask(allowed: list[int], vocab_size: int) -> mx.array

Build a 1-D mask of length vocab_size where allowed positions are 0 and disallowed positions are -inf.

Parameters

Name	Type	Required	Default	Description
allowed	list[int]	yes	none	Required positional or keyword input.
vocab_size	int	yes	none	Required positional or keyword input.

Returns

• Type: mx.array
• Direct return expressions: mx.full((vocab_size,), -float('inf')); mx.array(buf)

Exceptions and behavior

Method JSONSchemaLogitsProcessor._build_allow_mask calls mx.full, float, np.full, mx.array; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L795-L810.

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor.__call__ · method

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor.__call__(tokens: mx.array, logits: mx.array) -> mx.array

Apply the allowed-tokens mask to logits.

Parameters

Name	Type	Required	Default	Description
tokens	mx.array	yes	none	Required positional or keyword input.
logits	mx.array	yes	none	Required positional or keyword input.

Returns

• Type: mx.array
• Direct return expressions: _eos_logits_or_original(self._eos_set, logits, self._build_allow_mask); eos_logits; logits; logits + mask

Exceptions and behavior

Method JSONSchemaLogitsProcessor.__call__ updates self._disabled; calls _eos_logits_or_original, hasattr, tokens.tolist, list; has 4 explicit return paths. No direct raise statement appears in this definition.

View source #L814-L910.

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor.schema · method

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor.schema() -> dict | None

Return the normalized JSON Schema enforced for this request.

Parameters

This callable has no explicit inputs.

Returns

• Type: dict | None
• Direct return expressions: self._schema

Exceptions and behavior

Method JSONSchemaLogitsProcessor.schema returns self._schema. No direct raise statement appears in this definition.

View source #L915-L918.

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor.vocab_size · method

vllm_mlx.constrained.json_schema_processor.JSONSchemaLogitsProcessor.vocab_size() -> int

Return the tokenizer vocabulary size used to construct masks.

Parameters

This callable has no explicit inputs.

Returns

• Type: int
• Direct return expressions: self._vocab_size

Exceptions and behavior

Method JSONSchemaLogitsProcessor.vocab_size returns self._vocab_size. No direct raise statement appears in this definition.

View source #L921-L924.

Complete symbol map

This map also includes private definitions and nested helpers. The signature column exposes every explicit input even when an internal helper has no dedicated parameter prose.

Symbol	Kind	Signature and inputs	What it does	Source
LMFormatEnforcerNotAvailableError	class	LMFormatEnforcerNotAvailableError()	Raised when lm-format-enforcer is required but not installed.	#L33-L34
_canonical_schema_key	function	_canonical_schema_key(schema: dict \| None) -> str	Function _canonical_schema_key calls json.dumps(schema, sort_keys=True, separators=(',', ':')).encode, json.dumps, hashlib.sha256(blob).hexdigest, hashlib.sha256; has 2 explicit return paths.	#L50-L54
_get_or_build_parser	function	_get_or_build_parser(schema: dict \| None) -> tuple[dict, Any]	Return (parser_schema, JsonSchemaParser) for schema, memoised.	#L57-L73
is_available	function	is_available() -> bool	Return True iff lm-format-enforcer is importable.	#L76-L82
_simplify_schema	function	_simplify_schema(schema: dict) -> dict	Pre-process a JSON Schema for lm-format-enforcer compatibility.	#L97-L207
_simplify_schema._resolve	nested function	_simplify_schema._resolve(node: Any, depth: int = 0) -> Any	Nested Function _simplify_schema._resolve calls isinstance, ref.split, len, resolving.add; has 3 explicit return paths.	#L121-L205
_force_no_additional_properties	function	_force_no_additional_properties(schema: dict) -> dict	Return a deep copy of schema with additionalProperties: false injected into every object-type sub-schema that declares properties.	#L210-L224
_inject_no_additional_props	function	_inject_no_additional_props(node: Any) -> None	Recursively inject additionalProperties: false into node.	#L227-L238
_collect_property_names	function	_collect_property_names(schema: dict \| None) -> set[str]	Collect all property names declared anywhere in schema.	#L241-L247
_walk_properties	function	_walk_properties(node: Any, names: set[str]) -> None	Function _walk_properties calls isinstance, node.get, names.update, props.keys; returns None.	#L250-L264
_complete_json_eos_logits	function	_complete_json_eos_logits(eos_set: set[int], suffix: list[int], logits: mx.array, is_complete_json, build_allow_mask) -> mx.array \| None	Function _complete_json_eos_logits calls is_complete_json, _eos_logits; has 2 explicit return paths.	#L267-L276
_eos_logits	function	_eos_logits(eos_set: set[int], logits: mx.array, build_allow_mask) -> mx.array \| None	Function _eos_logits calls build_allow_mask, sorted; has 2 explicit return paths.	#L279-L290
_eos_logits_or_original	function	_eos_logits_or_original(eos_set: set[int], logits: mx.array, build_allow_mask) -> mx.array	Function _eos_logits_or_original calls _eos_logits; returns logits if masked is None else masked.	#L293-L299
JSONSchemaLogitsProcessor	class	JSONSchemaLogitsProcessor(schema: dict \| None, tokenizer: Any)	Logits processor that constrains generation to valid JSON.	#L302-L924
JSONSchemaLogitsProcessor.__init__	method	JSONSchemaLogitsProcessor.__init__(schema: dict \| None, tokenizer: Any) -> None	Method JSONSchemaLogitsProcessor.__init__ updates self._tokenizer, self._schema, self._tok_data, self._disabled; calls is_available, LMFormatEnforcerNotAvailableError, get_tokenizer_data, _get_or_build_parser; can raise LMFormatEnforcerNotAvailableError.	#L317-L414
JSONSchemaLogitsProcessor._suffix	method	JSONSchemaLogitsProcessor._suffix(tokens_list: list[int]) -> list[int]	Return the slice of tokens that corresponds to generated output.	#L418-L426
JSONSchemaLogitsProcessor._decode_token_cached	method	JSONSchemaLogitsProcessor._decode_token_cached(tok_id: int) -> str \| None	Return the decoded text for a single token (cached).	#L428-L442
JSONSchemaLogitsProcessor._decode_suffix	method	JSONSchemaLogitsProcessor._decode_suffix(suffix: list[int]) -> str \| None	Decode suffix tokens to text.	#L444-L493
JSONSchemaLogitsProcessor._suffix_is_complete_json	method	JSONSchemaLogitsProcessor._suffix_is_complete_json(suffix: list[int]) -> bool	Return True if the decoded suffix parses as a complete JSON value.	#L495-L520
JSONSchemaLogitsProcessor._get_json_context	method	JSONSchemaLogitsProcessor._get_json_context(suffix: list[int]) -> str	Determine the JSON structural context of the current suffix.	#L522-L643
JSONSchemaLogitsProcessor._filter_at_key_context	method	JSONSchemaLogitsProcessor._filter_at_key_context(context: str, suffix: list[int], allowed: list[int]) -> list[int]	Apply schema-aware filtering when in key-related context.	#L645-L662
JSONSchemaLogitsProcessor._filter_key_start_tokens	method	JSONSchemaLogitsProcessor._filter_key_start_tokens(suffix: list[int], allowed: list[int]) -> list[int]	Filter tokens at key-start position.	#L664-L717
JSONSchemaLogitsProcessor._filter_in_key_tokens	method	JSONSchemaLogitsProcessor._filter_in_key_tokens(suffix: list[int], allowed: list[int]) -> list[int]	Filter tokens when we're inside an open key string.	#L719-L758
JSONSchemaLogitsProcessor._is_valid_key_prefix	method	JSONSchemaLogitsProcessor._is_valid_key_prefix(prefix: str) -> bool	Return True if prefix is a prefix of at least one valid key name.	#L760-L762
JSONSchemaLogitsProcessor._filter_nonprogress_whitespace_tokens	method	JSONSchemaLogitsProcessor._filter_nonprogress_whitespace_tokens(suffix: list[int], allowed: list[int]) -> list[int]	Stop constrained JSON from spending a long run on pure whitespace.	#L764-L793
JSONSchemaLogitsProcessor._build_allow_mask	method	JSONSchemaLogitsProcessor._build_allow_mask(allowed: list[int], vocab_size: int) -> mx.array	Build a 1-D mask of length vocab_size where allowed positions are 0 and disallowed positions are -inf.	#L795-L810
JSONSchemaLogitsProcessor.__call__	method	JSONSchemaLogitsProcessor.__call__(tokens: mx.array, logits: mx.array) -> mx.array	Apply the allowed-tokens mask to logits.	#L814-L910
JSONSchemaLogitsProcessor.schema	method	JSONSchemaLogitsProcessor.schema() -> dict \| None	Return the normalized JSON Schema enforced for this request.	#L915-L918
JSONSchemaLogitsProcessor.vocab_size	method	JSONSchemaLogitsProcessor.vocab_size() -> int	Return the tokenizer vocabulary size used to construct masks.	#L921-L924